Update app.py

app.py

@@ -1,58 +1,64 @@
+import os
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
-import pinecone
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
-import torch
-import os
+from pinecone import Pinecone, ServerlessSpec
 
-# Load secrets and environment variables
+# --- Load environment variables ---
+HF_TOKEN = os.environ.get("HF_TOKEN")
 PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
 PINECONE_INDEX_NAME = os.environ.get("PINECONE_INDEX_NAME")
-HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# Step 1: Load embedding model and Pinecone
-embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-pinecone.init(api_key=PINECONE_API_KEY)
-pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
+assert HF_TOKEN is not None, "❌ HF_TOKEN is missing!"
+assert PINECONE_API_KEY is not None, "❌ PINECONE_API_KEY is missing!"
+assert PINECONE_INDEX_NAME is not None, "❌ Pinecone index name is missing!"
+
+# --- Load models ---
+embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/gpt2-fa", token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained("HooshvareLab/gpt2-fa", token=HF_TOKEN)
+
+# --- Connect to Pinecone ---
+pc = Pinecone(api_key=PINECONE_API_KEY)
 index = pc.Index(PINECONE_INDEX_NAME)
 
-# Step 2: Load GPT-2 language model
-model_name = "HooshvareLab/gpt2-fa"
-tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
-model = GPT2LMHeadModel.from_pretrained(model_name, use_auth_token=HF_TOKEN)
-model.eval()
-
-# Function: Embed input and search in Pinecone
-def retrieve_context(query, top_k=1):
-    xq = embedding_model.encode(query).tolist()
-    res = index.query(vector=xq, top_k=top_k, include_metadata=True)
-    if res.matches:
-        return res.matches[0].metadata['text']
-    return ""
-
-# Function: Generate response using GPT-2
-def generate_response(query, context):
-    prompt = f"پرسش: {query}\nپاسخ با توجه به اطلاعات زیر: {context}\nپاسخ:"
-    input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512)
-    output_ids = model.generate(input_ids, max_length=256, num_beams=4, no_repeat_ngram_size=2, early_stopping=True)
-    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return output.split("پاسخ:")[-1].strip()
-
-# Gradio interface
+# --- Inference pipeline ---
 def chat(query):
-    context = retrieve_context(query)
-    response = generate_response(query, context)
-    return response
-
-# UI
-with gr.Blocks() as demo:
-    gr.Markdown("## چت‌بات هوشمند تیام\nسوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید.")
-    with gr.Row():
-        inp = gr.Textbox(label="question", placeholder="سوال خود را وارد کنید")
-        out = gr.Textbox(label="output")
+    # Embed user question
+    xq = embedder.encode(query).tolist()
+
+    # Search in Pinecone
+    res = index.query(vector=xq, top_k=1, include_metadata=True)
+    matches = res.get("matches", [])
+
+    if not matches:
+        return "پاسخی برای سوال شما پیدا نشد. لطفا تماس بگیرید."
+
+    # Retrieve matched content
+    context = matches[0]['metadata']['text']
+
+    # Prepare prompt
+    prompt = f"سوال: {query}\nپاسخ بر اساس اطلاعات زیر بده: {context}\nپاسخ:"
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+
+    # Generate response
+    output_ids = model.generate(**inputs, max_new_tokens=100)
+    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # Post-process to remove prompt
+    if "پاسخ:" in answer:
+        answer = answer.split("پاسخ:", 1)[-1].strip()
+
+    return answer
+
+# --- Gradio UI ---
+with gr.Blocks(title="چت‌بات هوشمند تیام") as demo:
+    gr.Markdown("""## چت‌بات هوشمند تیام\nسوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید""")
+    question = gr.Textbox(label="question", placeholder="سوال خود را وارد کنید")
+    output = gr.Textbox(label="output")
     submit = gr.Button("Submit")
-    submit.click(chat, inputs=inp, outputs=out)
+    submit.click(fn=chat, inputs=question, outputs=output)
+    gr.ClearButton([question, output])
 
-# Launch
 if __name__ == "__main__":
     demo.launch()