Update app.py

app.py

@@ -1,37 +1,65 @@
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+import os
+import gradio as gr
 import torch
+from pinecone import Pinecone
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+
+# بارگذاری کلیدها
+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")
 
+# مدل QA فارسی
 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)
+
+# تابع RAG
 def rag_chatbot(message, history):
-    greetings = ["سلام", "درود", "سلام وقت بخیر", "خسته نباشید"]
+    greetings = ["سلام", "سلام وقت بخیر", "درود", "خسته نباشید"]
     if message.strip() in greetings:
         return "سلام! چطور می‌تونم کمکتون کنم؟ 😊"
 
+    # مرحله ۱: بازیابی نزدیک‌ترین اسناد
     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("پاسخ", "") for match in result['matches']])
+    context = "\n".join([match['metadata'].get("پاسخ", "") for match in result['matches'] if match['metadata'].get("پاسخ")])
 
     if not context.strip():
         return "متأسفم، اطلاعات کافی در این زمینه ندارم."
 
-    # پردازش مدل QA
+    # مرحله ۲: استفاده از مدل پاسخ‌گویی BERT
     inputs = tokenizer(message, context, return_tensors="pt", truncation=True, max_length=512)
     with torch.no_grad():
         outputs = model(**inputs)
 
-    answer_start = torch.argmax(outputs.start_logits)
-    answer_end = torch.argmax(outputs.end_logits)
-    answer = tokenizer.convert_tokens_to_string(
-        tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end+1])
-    )
+    start_idx = torch.argmax(outputs.start_logits)
+    end_idx = torch.argmax(outputs.end_logits)
+
+    answer_tokens = inputs["input_ids"][0][start_idx:end_idx + 1]
+    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
 
-    if not answer.strip() or answer.strip() == "[CLS]":
+    if not answer.strip() or answer.strip() in ["[CLS]", "[SEP]"]:
         return "متأسفم، نتونستم پاسخ دقیقی پیدا کنم."
 
     return answer.strip()
+
+# رابط Gradio
+chat_ui = gr.ChatInterface(
+    fn=rag_chatbot,
+    title="🤖 چت‌بات هوشمند تیام (نسخه RAG)",
+    description="پاسخ‌گویی ترکیبی بر پایه داده‌های واقعی شرکت + تولید پاسخ با مدل BERT فارسی",
+    theme="soft",
+)
+
+chat_ui.launch()