Update app.py

app.py

@@ -1,12 +1,11 @@
 import os
-import json
-import gradio as gr
+import pinecone
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-from pinecone import Pinecone
+import gradio as gr
 
-# دریافت توکن‌ها و متغیرهای محیطی
+# --- Load environment variables ---
 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")
@@ -15,50 +14,40 @@ 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)
-
-# دریافت پاسخ اولیه از 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:
-        return result['matches'][0]['metadata'].get('answer', "پاسخی یافت نشد.")
-    else:
-        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=100)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# رابط 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="پرسش خود را درباره خدمات دیجیتال مارکتینگ تیام بپرسید."
+# --- Initialize Pinecone ---
+pinecone.init(api_key=PINECONE_API_KEY, environment="gcp-starter")
+index = pinecone.Index(PINECONE_INDEX_NAME)
+
+# --- Load models ---
+embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", token=HF_TOKEN)
+tokenizer = T5Tokenizer.from_pretrained("google/mt5-small", token=HF_TOKEN)
+model = T5ForConditionalGeneration.from_pretrained("google/mt5-small", token=HF_TOKEN)
+
+def generate_answer(question):
+    # Embed the question
+    question_embedding = embedding_model.encode(question).tolist()
+
+    # Query Pinecone for similar content
+    response = index.query(vector=question_embedding, top_k=3, include_metadata=True)
+    contexts = [match['metadata']['text'] for match in response['matches'] if 'text' in match['metadata']]
+
+    # Concatenate context
+    context = "\n".join(contexts)
+    input_text = f"پرسش: {question}\nاطلاعات: {context}\nپاسخ:"
+
+    # Tokenize and generate
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
+    outputs = model.generate(**inputs, max_new_tokens=200)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer
+
+# --- Gradio UI ---
+iface = gr.Interface(
+    fn=generate_answer,
+    inputs=gr.Textbox(label="question"),
+    outputs=gr.Textbox(label="output"),
+    title="💬 چت‌بات هوشمند تیام",
+    description="سؤالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید."
 )
 
-demo.launch()
+iface.launch()