app.py

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

# بارگذاری کلیدها
PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
INDEX_NAME = os.environ.get("INDEX_NAME")

# مدل embedding
embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# مدل تولید متن فارسی MT5
tokenizer = AutoTokenizer.from_pretrained("m3hrdadfi/mt5-small-finetuned-persian-qa")
model = AutoModelForSeq2SeqLM.from_pretrained("m3hrdadfi/mt5-small-finetuned-persian-qa")

def generate_answer_mt5(question, context):
    input_text = f"پرسش: {question} \n متن: {context}"
    inputs = tokenizer.encode(input_text, return_tensors="pt", truncation=True)
    outputs = model.generate(inputs, max_length=128, num_beams=4, early_stopping=True)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# اتصال به Pinecone
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(INDEX_NAME)

# تابع چت
def rag_chatbot(message, history):
    greetings = ["سلام", "سلام وقت بخیر", "درود", "خسته نباشید"]
    if message.strip() in greetings:
        return "سلام! چطور می‌تونم کمکتون کنم؟ 😊"

    query_vector = embedder.encode(message).tolist()
    result = index.query(vector=query_vector, top_k=3, include_metadata=True)

    if not result['matches']:
        return "متأسفم، پاسخی در پایگاه داده پیدا نکردم."

    context = "\n".join([match['metadata'].get("پاسخ", "") for match in result['matches']])

    if not context.strip():
        return "متأسفم، اطلاعات کافی در این زمینه ندارم. لطفاً با ما تماس بگیرید."

    answer = generate_answer_mt5(message, context)
    if len(answer) < 3:
        return "متأسفم، پاسخ مشخصی برای این سوال ندارم."
    return answer

# رابط کاربری Gradio
chat_ui = gr.ChatInterface(
    fn=rag_chatbot,
    title="🤖 چت‌بات هوشمند تیام (نسخه RAG با MT5)",
    description="پاسخ‌گویی بر پایه داده‌های واقعی شرکت + تولید متن با مدل فارسی",
    theme="soft"
)

chat_ui.launch()