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

MODEL_ID = "MaxKio/Mio-1.0-Pro"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16)

@spaces.GPU
def generate_text(prompt):
    if not prompt.strip():
        return "من فضلك اكتب أمراً أولاً."
    try:
        model.to("cuda")
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True
        )
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return result
    except Exception as e:
        return f"خطأ أثناء التوليد: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# <center>واجهة تشغيل نموذج Mio Pro الخفيفة</center>")

    with gr.Row():
        input_text = gr.Textbox(label="أدخل أمر الكود الموجه:", lines=5)
        output_text = gr.Textbox(label="استجابة النموذج:", lines=7, interactive=False)

    submit_btn = gr.Button("إرسال", variant="primary")
    submit_btn.click(fn=generate_text, inputs=input_text, outputs=output_text)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())