Update app.py

app.py

@@ -1,61 +1,64 @@
 import os
 import json
-import torch
 import gradio as gr
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 from pinecone import Pinecone
-from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# بارگذاری توکن‌ها از محیط امن
+# دریافت توکن‌ها و متغیرهای محیطی
 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")
 
+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!"
+
 # بارگذاری مدل embedding
 embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", use_auth_token=HF_TOKEN)
 
+# بارگذاری مدل زبانی (مثل MT5 یا Gemma)
+model_name = "google/mt5-small"
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
+language_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, token=HF_TOKEN)
+
 # اتصال به Pinecone
 pc = Pinecone(api_key=PINECONE_API_KEY)
 index = pc.Index(PINECONE_INDEX_NAME)
 
-# بارگذاری مدل زبانی MT5
-tokenizer = T5Tokenizer.from_pretrained("google/mt5-small", token=HF_TOKEN)
-language_model = T5ForConditionalGeneration.from_pretrained("google/mt5-small", token=HF_TOKEN)
-
-# تابع جستجو در Pinecone
+# دریافت پاسخ اولیه از Pinecone
 def retrieve_answer(query, threshold=0.65, top_k=3):
     query_embedding = embedding_model.encode([query])[0]
     result = index.query(vector=query_embedding.tolist(), top_k=top_k, include_metadata=True)
 
     if result['matches'] and result['matches'][0]['score'] > threshold:
-        print(f"📊 Similarity: {result['matches'][0]['score']:.3f}")
-        metadata = result['matches'][0]['metadata']
-        return metadata.get('answer', 'پاسخی یافت نشد.')
+        return result['matches'][0]['metadata'].get('answer', "پاسخی یافت نشد.")
     else:
-        return "متأسفم، پاسخ دقیقی برای این سوال نداریم. لطفاً با ما تماس بگیرید."
-
-# تابع تولید پاسخ طبیعی با MT5
-def generate_natural_answer(question, raw_answer):
-    prompt = f"پرسش: {question}\nپاسخ دقیق: {raw_answer}\nپاسخ طبیعی:"
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(language_model.device)
+        return "متأسفم، پاسخی برای این سوال در پایگاه داده یافت نشد."
 
+# بازنویسی پاسخ به صورت طبیعی با مدل زبانی
+def rewrite_answer(question, raw_answer):
+    input_text = f"پرسش: {question}\nپاسخ: {raw_answer}"
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
-        outputs = language_model.generate(**inputs, max_new_tokens=128, do_sample=False)
-
+        outputs = language_model.generate(**inputs, max_new_tokens=100)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# اتصال همه‌چیز در رابط Gradio
-def chat_interface(user_question):
-    raw_answer = retrieve_answer(user_question)
-    return generate_natural_answer(user_question, raw_answer)
-
 # رابط Gradio
+def chat_interface(question):
+    raw_answer = retrieve_answer(question)
+    if "متأسفم" in raw_answer:
+        return raw_answer
+    final_answer = rewrite_answer(question, raw_answer)
+    return final_answer
+
 demo = gr.Interface(
     fn=chat_interface,
     inputs="text",
     outputs="text",
     title="چت‌بات تیام",
-    description="سؤالات خود را از آژانس دیجیتال مارکتینگ تیام بپرسید."
+    description="پرسش خود را درباره خدمات دیجیتال مارکتینگ تیام بپرسید."
 )
 
 demo.launch()