Update app.py

app.py

@@ -1,57 +1,58 @@
-import os
-import pinecone
+import gradio as gr
 from sentence_transformers import SentenceTransformer
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+import pinecone
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import torch
-import gradio as gr
+import os
 
-# Load environment variables
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# Load secrets and 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 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!"
-
-# Load embedding model
-embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", use_auth_token=HF_TOKEN)
-
-# Load tokenizer and model
-tokenizer = T5Tokenizer.from_pretrained("google/mt5-small", token=HF_TOKEN)
-model = T5ForConditionalGeneration.from_pretrained("google/mt5-small", token=HF_TOKEN)
-
-# Initialize Pinecone client
+# Step 1: Load embedding model and Pinecone
+embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+pinecone.init(api_key=PINECONE_API_KEY)
 pc = pinecone.Pinecone(api_key=PINECONE_API_KEY)
 index = pc.Index(PINECONE_INDEX_NAME)
 
-def query_index(question):
-    # Embed question
-    question_embedding = embedder.encode(question).tolist()
-
-    # Query Pinecone
-    results = index.query(vector=question_embedding, top_k=1, include_metadata=True)
-
-    if results.matches:
-        retrieved_text = results.matches[0].metadata.get("text", "")
-    else:
-        retrieved_text = "متاسفم، پاسخ مناسبی پیدا نکردم."
-
-    # Generate answer
-    input_text = f"پرسش: {question} \n پاسخ بر اساس دانش: {retrieved_text}"
-    input_ids = tokenizer(input_text, return_tensors="pt", truncation=True).input_ids
-    output_ids = model.generate(input_ids, max_length=100)
-    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-    return answer
-
-# Gradio UI
-iface = gr.Interface(
-    fn=query_index,
-    inputs=gr.Textbox(label="question", placeholder="سوال خود را وارد کنید"),
-    outputs=gr.Textbox(label="output"),
-    title="چت‌بات هوشمند تیام",
-    description="سوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید."
-)
-
-iface.launch(server_name="0.0.0.0", server_port=7860)
+# Step 2: Load GPT-2 language model
+model_name = "HooshvareLab/gpt2-fa"
+tokenizer = GPT2Tokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+model = GPT2LMHeadModel.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+model.eval()
+
+# Function: Embed input and search in Pinecone
+def retrieve_context(query, top_k=1):
+    xq = embedding_model.encode(query).tolist()
+    res = index.query(vector=xq, top_k=top_k, include_metadata=True)
+    if res.matches:
+        return res.matches[0].metadata['text']
+    return ""
+
+# Function: Generate response using GPT-2
+def generate_response(query, context):
+    prompt = f"پرسش: {query}\nپاسخ با توجه به اطلاعات زیر: {context}\nپاسخ:"
+    input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512)
+    output_ids = model.generate(input_ids, max_length=256, num_beams=4, no_repeat_ngram_size=2, early_stopping=True)
+    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return output.split("پاسخ:")[-1].strip()
+
+# Gradio interface
+def chat(query):
+    context = retrieve_context(query)
+    response = generate_response(query, context)
+    return response
+
+# UI
+with gr.Blocks() as demo:
+    gr.Markdown("## چت‌بات هوشمند تیام\nسوالات خود درباره خدمات دیجیتال مارکتینگ تیام را بپرسید.")
+    with gr.Row():
+        inp = gr.Textbox(label="question", placeholder="سوال خود را وارد کنید")
+        out = gr.Textbox(label="output")
+    submit = gr.Button("Submit")
+    submit.click(chat, inputs=inp, outputs=out)
+
+# Launch
+if __name__ == "__main__":
+    demo.launch()