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

model_name = "Qwen/Qwen2.5-1.5B-Instruct"

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

SYSTEM_PROMPT = """
أنت مساعد ذكي تتحدث العربية بطلاقة.
رد على المستخدم بالعربية الفصحى أو باللهجة السعودية حسب أسلوبه.
"""

def chat(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def respond(message, history):
    full_prompt = SYSTEM_PROMPT + "\n"
    for user, bot in history:
        full_prompt += f"User: {user}\nBot: {bot}\n"
    full_prompt += f"User: {message}\nBot:"
    reply = chat(full_prompt)
    return reply

demo = gr.ChatInterface(respond, title="Arabic Chatbot")

if __name__ == "__main__":
    demo.launch()