Update app.py

app.py

@@ -1,59 +1,72 @@
-
 import os
 import gradio as gr
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
 from pinecone import Pinecone
 
-# بارگذاری کلیدها
+# بارگذاری مقادیر محرمانه از محیط
 PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
 INDEX_NAME = os.environ.get("INDEX_NAME")
 
 # مدل embedding
 embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
 
-# مدل تولید متن فارسی 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)
+# مدل پاسخ به سوال فارسی
+tokenizer = AutoTokenizer.from_pretrained("SeyedAli/Persian-QA-Bert-V1")
+model = AutoModelForQuestionAnswering.from_pretrained("SeyedAli/Persian-QA-Bert-V1")
 
 # اتصال به Pinecone
 pc = Pinecone(api_key=PINECONE_API_KEY)
 index = pc.Index(INDEX_NAME)
 
-# تابع چت
+# پاسخ به سلام و احوال‌پرسی
+greetings = ["سلام", "سلام وقت بخیر", "درود", "خسته نباشید"]
+def is_greeting(text):
+    return any(text.strip().startswith(greet) for greet in greetings)
+
+# تابع اصلی چت‌بات
 def rag_chatbot(message, history):
-    greetings = ["سلام", "سلام وقت بخیر", "درود", "خسته نباشید"]
-    if message.strip() in greetings:
+    if is_greeting(message):
         return "سلام! چطور می‌تونم کمکتون کنم؟ 😊"
 
+    # مرحله 1: بازیابی داده‌های مرتبط از Pinecone
     query_vector = embedder.encode(message).tolist()
     result = index.query(vector=query_vector, top_k=3, include_metadata=True)
 
-    if not result['matches']:
-        return "متأسفم، پاسخی در پایگاه داده پیدا نکردم."
+    if not result["matches"]:
+        return "متأسفم، پاسخ مشخصی برای این سؤال در پایگاه داده ندارم."
+
+    # مرحله 2: استخراج متون بازیابی‌شده
+    contexts = [match["metadata"].get("پاسخ", "") for match in result["matches"]]
+    best_answer = ""
+    best_score = -float("inf")
+
+    for context in contexts:
+        inputs = tokenizer.encode_plus(message, context, return_tensors="pt", truncation=True)
+        outputs = model(**inputs)
+        answer_start = int(outputs.start_logits.argmax())
+        answer_end = int(outputs.end_logits.argmax())
+        score = float(outputs.start_logits[0][answer_start] + outputs.end_logits[0][answer_end])
+
+        answer = tokenizer.convert_tokens_to_string(
+            tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end + 1])
+        )
 
-    context = "\n".join([match['metadata'].get("پاسخ", "") for match in result['matches']])
+        if score > best_score and answer.strip() != "":
+            best_score = score
+            best_answer = answer
 
-    if not context.strip():
-        return "متأسفم، اطلاعات کافی در این زمینه ندارم. لطفاً با ما تماس بگیرید."
+    if best_answer.strip() == "":
+        return "متأسفم، اطلاعات کافی برای پاسخ دقیق ندارم. لطفاً تماس بگیرید."
 
-    answer = generate_answer_mt5(message, context)
-    if len(answer) < 3:
-        return "متأسفم، پاسخ مشخصی برای این سوال ندارم."
-    return answer
+    return best_answer.strip()
 
 # رابط کاربری Gradio
 chat_ui = gr.ChatInterface(
     fn=rag_chatbot,
-    title="🤖 چت‌بات هوشمند تیام (نسخه RAG با MT5)",
-    description="پاسخ‌گویی بر پایه داده‌های واقعی شرکت + تولید متن با مدل فارسی",
-    theme="soft"
+    title="🤖 چت‌بات هوشمند تیام (نسخه RAG)",
+    description="پاسخ‌گویی ترکیبی بر پایه داده‌های واقعی شرکت + تولید متن با مدل فارسی BERT",
+    theme="soft",
 )
 
 chat_ui.launch()