import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# استفاده از مدل باز جایگزین
model_name = "mistralai/Mistral-7B-Instruct-v0.2"  # یا "google/gemma-7b"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

def generate_response(prompt, max_new_tokens=512, temperature=0.7):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        do_sample=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

with gr.Blocks() as demo:
    gr.Markdown("# چت بات هوشمند")
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("پاک کردن")

    def respond(message, chat_history):
        response = generate_response(message)
        chat_history.append((message, response))
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()