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

# Load model
model_name = "itsmeussa/AdabTranslate-Darija"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def translate_darija_to_msa(text):
    if not text.strip():
        return "⚠️ Please enter some text."
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Simple Gradio chat interface
with gr.Blocks(css=".user{background:#ff78ff;color:white} .bot{background:#eee}") as demo:
    gr.Markdown("## ✨ Darija → MSA Translator ✨")
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(placeholder="اكتب باللهجة المغربية...")
    clear = gr.Button("مسح")

    def respond(user_message, chat_history):
        bot_message = translate_darija_to_msa(user_message)
        chat_history.append((user_message, bot_message))
        return "", chat_history

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

demo.launch()