Update app.py

app.py

@@ -1,67 +1,63 @@
 import os
-import gradio as gr
-import pinecone
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer, util
+from pinecone import Pinecone
+import gradio as gr
 
-# Environment variables
-PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
-PINECONE_INDEX_NAME = os.environ.get("PINECONE_INDEX_NAME")
-HF_TOKEN = os.environ.get("HF_TOKEN")
-
-assert PINECONE_API_KEY is not None, "\u274c PINECONE_API_KEY is missing!"
-assert PINECONE_INDEX_NAME is not None, "\u274c PINECONE_INDEX_NAME is missing!"
-assert HF_TOKEN is not None, "\u274c HF_TOKEN is missing!"
+# انتشار متغیرها از Hugging Face secrets
+HF_TOKEN = os.getenv("HF_TOKEN")
+PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
+PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME")
 
-# Load embedding model
+# مدل کوچک برای embedding (sentence-transformers)
 embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
 
-# Init Pinecone
-pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
-index = pc.Index(PINECONE_INDEX_NAME)
-
-# Load language model and tokenizer
+# مدل زبانی GPT2 فارسی
 tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/gpt2-fa")
 model = AutoModelForCausalLM.from_pretrained("HooshvareLab/gpt2-fa")
-text_generator = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_length=45,
-    do_sample=True,
-    top_p=0.95,
-    temperature=0.8,
-    return_full_text=False,
-)
 
-def generate_answer(question):
-    try:
-        # Step 1: Embed the question
-        question_embedding = embedding_model.encode(question).tolist()
-
-        # Step 2: Search similar questions in Pinecone
-        search_result = index.query(vector=question_embedding, top_k=1, include_metadata=True)
-
-        if search_result and search_result.matches:
-            best_match = search_result.matches[0].metadata.get("answer", "")
-
-            # Step 3: Rewrite with the language model
-            prompt = f"پرسش: {question}\nپاسخ: {best_match}\nپاسخ نهایی:"
-            output = text_generator(prompt, max_new_tokens=50)[0]["generated_text"]
-            return output.strip()
-        else:
-            return "پاسخی برای این پرسش در پایگاه داده یافت نشد. لطفاً با پشتیبانی تماس بگیرید."
+# اتصال به Pinecone
+pc = Pinecone(api_key=PINECONE_API_KEY)
+index = pc.Index(PINECONE_INDEX_NAME)
 
-    except Exception as e:
-        return f"خطا: {str(e)}"
+# توابع
 
-# Gradio Interface
-demo = gr.Interface(
-    fn=generate_answer,
-    inputs=gr.Textbox(label="question", placeholder="سوال خود را وارد کنید..."),
+def retrieve_from_pinecone(query):
+    query_embedding = embedding_model.encode(query).tolist()
+    search_result = index.query(vector=query_embedding, top_k=1, include_metadata=True)
+    try:
+        return search_result['matches'][0]['metadata']['answer']
+    except:
+        return "پاسخی برای این سوال پیدا نشد."
+
+def generate_response(query):
+    base_answer = retrieve_from_pinecone(query)
+    prompt = f"{query}\n{base_answer}"
+
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    output = model.generate(
+        inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_new_tokens=30,  # کمک به تسریع پاسخگویی
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # جدا کردن پاسخ تولیدی از prompt
+    return response.replace(prompt, "").strip()
+
+# رابط کاربری Gradio
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(label="question", placeholder="سوال خود را وارد کنید"),
     outputs=gr.Textbox(label="output"),
-    title="چت‌بات هوشمند تیام",
-    description="سوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید."
+    title="چتبات هوشمند تیام",
+    description="سوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید"
 )
 
-demo.launch()
+iface.launch()