app.py

# app.py

import json
import gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from pinecone import Pinecone
import os

# ===============================
# 🌐 اتصال به Pinecone
# ===============================
PINECONE_API_KEY = "pcsk_6p6AmJ_Qua4tQN69badNHEGZTj3tt5Bd7LiyiDGcXDj92LxSaBzK2ypYxTRx2rafTEJhjL"
PINECONE_INDEX_NAME = "tiyam-chat"

pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(PINECONE_INDEX_NAME)

# ===============================
# 🔤 بارگذاری مدل embedding
# ===============================
embedding_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')

# ===============================
# 🧠 بارگذاری مدل google/mt5-small برای بازنویسی
# ===============================
hf_token = os.getenv("HF_TOKEN")

tokenizer = AutoTokenizer.from_pretrained(
    "google/mt5-small", 
    use_auth_token=hf_token
)
mt5_model = AutoModelForSeq2SeqLM.from_pretrained(
    "google/mt5-small", 
    use_auth_token=hf_token
)

# ===============================
# 🔍 بازیابی پاسخ از Pinecone
# ===============================
def retrieve_answer(query, threshold=0.65, top_k=1):
    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 None

# ===============================
# ✨ بازنویسی پاسخ با MT5
# ===============================
def rewrite_with_mt5(answer, question):
    prompt = f"پاسخ به سوال: {question} بر اساس اطلاعات: {answer}"
    input_ids = tokenizer.encode(prompt, return_tensors="pt", truncation=True)
    output_ids = mt5_model.generate(input_ids, max_new_tokens=60)
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# ===============================
# 💬 منطق نهایی پاسخ‌دهی
# ===============================
def final_answer(user_question):
    answer = retrieve_answer(user_question)

    if answer:
        return rewrite_with_mt5(answer, user_question)
    else:
        # پاسخ عمومی برای سوالات چتی یا ناموجود
        general_prompts = {
            "سلام": "سلام! خوشحالیم که اینجایی 😊",
            "خوبی؟": "مرسی! من خوبم، شما چطورید؟",
            "مرسی": "خواهش می‌کنم، در خدمت شما هستیم."
        }
        for key in general_prompts:
            if key in user_question:
                return general_prompts[key]

        return "متأسفم، پاسخ دقیقی برای این سوال نداریم. لطفاً با ما تماس بگیرید."

# ===============================
# 🎛️ رابط Gradio
# ===============================
demo = gr.Interface(
    fn=final_answer,
    inputs=gr.Textbox(lines=2, label="سؤال شما"),
    outputs=gr.Textbox(label="پاسخ تیام"),
    title="💬 چت‌بات هوشمند تیام",
    description="سؤالات خود را از آژانس دیجیتال مارکتینگ تیام بپرسید. سیستم از ترکیب جستجوی دقیق و بازنویسی طبیعی استفاده می‌کند."
)

demo.launch()