Update app.py

app.py

@@ -1,69 +1,59 @@
+
 import os
 import gradio as gr
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from pinecone import Pinecone
 
-# بارگذاری کلیدها از محیط
+# بارگذاری کلیدها
 PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
-INDEX_NAME = "tiyam-chat"
+INDEX_NAME = os.environ.get("INDEX_NAME")
 
-# اتصال به Pinecone
-pc = Pinecone(api_key=PINECONE_API_KEY)
-index = pc.Index(INDEX_NAME)
-
-# مدل بردارساز
+# مدل embedding
 embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
 
-# مدل زبانی فارسی
-generator = pipeline("text-generation", model="HooshvareLab/gpt2-fa")
+# مدل تولید متن فارسی MT5
+tokenizer = AutoTokenizer.from_pretrained("m3hrdadfi/mt5-small-finetuned-persian-qa")
+model = AutoModelForSeq2SeqLM.from_pretrained("m3hrdadfi/mt5-small-finetuned-persian-qa")
+
+def generate_answer_mt5(question, context):
+    input_text = f"پرسش: {question} \n متن: {context}"
+    inputs = tokenizer.encode(input_text, return_tensors="pt", truncation=True)
+    outputs = model.generate(inputs, max_length=128, num_beams=4, early_stopping=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# پاسخ‌های سلام
-greetings = ["سلام", "سلام وقت بخیر", "درود", "خسته نباشید", "سلام عزیز"]
+# اتصال به Pinecone
+pc = Pinecone(api_key=PINECONE_API_KEY)
+index = pc.Index(INDEX_NAME)
 
 # تابع چت
 def rag_chatbot(message, history):
+    greetings = ["سلام", "سلام وقت بخیر", "درود", "خسته نباشید"]
     if message.strip() in greetings:
         return "سلام! چطور می‌تونم کمکتون کنم؟ 😊"
 
-    try:
-        query_vector = embedder.encode(message).tolist()
-        result = index.query(vector=query_vector, top_k=3, include_metadata=True)
-
-        if not result['matches']:
-            return "متأسفم، پاسخ مشخصی پیدا نکردم."
-
-        context = "\n".join([
-            match['metadata'].get("answer", "") or match['metadata'].get("پاسخ", "")
-            for match in result['matches'] if match['score'] > 0.7
-        ])
-
-        if not context.strip():
-            return "متأسفم، اطلاعات کافی در این زمینه ندارم. لطفاً با ما تماس بگیرید."
-
-        prompt = f"""اطلاعات زیر را بخوان و فقط بر اساس همین اطلاعات پاسخ بده.
-اگر جواب در اطلاعات نبود، بگو اطلاعی ندارم.
-
-اطلاعات:
-{context}
+    query_vector = embedder.encode(message).tolist()
+    result = index.query(vector=query_vector, top_k=3, include_metadata=True)
 
-سؤال: {message}
-پاسخ:"""
+    if not result['matches']:
+        return "متأسفم، پاسخی در پایگاه داده پیدا نکردم."
 
-        response = generator(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)[0]['generated_text']
-        answer = response.split("پاسخ:")[-1].strip()
+    context = "\n".join([match['metadata'].get("پاسخ", "") for match in result['matches']])
 
-        return answer if len(answer) > 3 else "متأسفم، پاسخ مشخصی برای این سوال ندارم."
+    if not context.strip():
+        return "متأسفم، اطلاعات کافی در این زمینه ندارم. لطفاً با ما تماس بگیرید."
 
-    except Exception as e:
-        return f"🚫 خطا: {str(e)}"
+    answer = generate_answer_mt5(message, context)
+    if len(answer) < 3:
+        return "متأسفم، پاسخ مشخصی برای این سوال ندارم."
+    return answer
 
-# رابط چت Gradio
+# رابط کاربری Gradio
 chat_ui = gr.ChatInterface(
     fn=rag_chatbot,
-    title="🤖 چت‌بات هوشمند تیام",
+    title="🤖 چت‌بات هوشمند تیام (نسخه RAG با MT5)",
     description="پاسخ‌گویی بر پایه داده‌های واقعی شرکت + تولید متن با مدل فارسی",
-    theme="soft",
+    theme="soft"
 )
 
 chat_ui.launch()