Update app.py

app.py

@@ -1,64 +1,67 @@
 import os
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from sentence_transformers import SentenceTransformer
-from pinecone import Pinecone, ServerlessSpec
+import pinecone
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from sentence_transformers import SentenceTransformer, util
 
-# --- Load environment variables ---
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# Environment variables
 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")
 
-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!"
+assert PINECONE_API_KEY is not None, "\u274c PINECONE_API_KEY is missing!"
+assert PINECONE_INDEX_NAME is not None, "\u274c PINECONE_INDEX_NAME is missing!"
+assert HF_TOKEN is not None, "\u274c HF_TOKEN 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)
+# Load embedding model
+embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
 
-# --- Connect to Pinecone ---
-pc = Pinecone(api_key=PINECONE_API_KEY)
+# Init Pinecone
+pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
 index = pc.Index(PINECONE_INDEX_NAME)
 
-# --- Inference pipeline ---
-def chat(query):
-    # 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 "پاسخی برای سوال شما پیدا نشد. لطفا تماس بگیرید."
+# Load language model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/gpt2-fa")
+model = AutoModelForCausalLM.from_pretrained("HooshvareLab/gpt2-fa")
+text_generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_length=100,
+    do_sample=True,
+    top_p=0.95,
+    temperature=0.8,
+    return_full_text=False,
+)
 
-    # Retrieve matched content
-    context = matches[0]['metadata']['text']
+def generate_answer(question):
+    try:
+        # Step 1: Embed the question
+        question_embedding = embedding_model.encode(question).tolist()
 
-    # Prepare prompt
-    prompt = f"سوال: {query}\nپاسخ بر اساس اطلاعات زیر بده: {context}\nپاسخ:"
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+        # Step 2: Search similar questions in Pinecone
+        search_result = index.query(vector=question_embedding, top_k=1, include_metadata=True)
 
-    # Generate response
-    output_ids = model.generate(**inputs, max_new_tokens=100)
-    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        if search_result and search_result.matches:
+            best_match = search_result.matches[0].metadata.get("answer", "")
 
-    # Post-process to remove prompt
-    if "پاسخ:" in answer:
-        answer = answer.split("پاسخ:", 1)[-1].strip()
+            # Step 3: Rewrite with the language model
+            prompt = f"پرسش: {question}\nپاسخ: {best_match}\nپاسخ نهایی:"
+            output = text_generator(prompt, max_new_tokens=50)[0]["generated_text"]
+            return output.strip()
+        else:
+            return "پاسخی برای این پرسش در پایگاه داده یافت نشد. لطفاً با پشتیبانی تماس بگیرید."
 
-    return answer
+    except Exception as e:
+        return f"خطا: {str(e)}"
 
-# --- 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(fn=chat, inputs=question, outputs=output)
-    gr.ClearButton([question, output])
+# Gradio Interface
+demo = gr.Interface(
+    fn=generate_answer,
+    inputs=gr.Textbox(label="question", placeholder="سوال خود را وارد کنید..."),
+    outputs=gr.Textbox(label="output"),
+    title="چت‌بات هوشمند تیام",
+    description="سوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید."
+)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()