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

# ----------------------
# 1. تحميل المودل
# ----------------------
model_name = "bigcode/starcoder2-7b"

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,   # FP16 لتقليل استهلاك VRAM
    device_map="auto",
    trust_remote_code=True
)
model.eval()

# ----------------------
# 2. دالة التوليد
# ----------------------
def generate_code(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=700,
            temperature=0.2,
            do_sample=False,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id,
            repetition_penalty=1.1
        )
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# ----------------------
# 3. واجهة Gradio بدون allow_flagging
# ----------------------
title = "StarCoder2 Flutter Code Generator"
description = """
Generate Dart / Flutter code using StarCoder2.
Type your prompt describing the widget or functionality you want.
"""

demo = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=8, placeholder="Write your Flutter prompt here..."),
    outputs=gr.Textbox(lines=20),
    title=title,
    description=description
)

# ----------------------
# 4. تشغيل الواجهة
# ----------------------
demo.launch()