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

# تحميل الموديل
model_name = "Salesforce/codegen-350M-mono"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")

# دالة التوليد
def generate_code(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=200,            # زيادة عدد التوكنز ليرجع كود أطول
        temperature=0.7,               # ضبط العشوائية (أقل من 1 = كود أدق)
        do_sample=True,                # تفعيل التوليد العشوائي
        top_k=50,                      # استخدام top-k sampling
        top_p=0.95                     # أو استخدام nucleus sampling
    )
    code = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # إزالة البادئة من الـ prompt لو تكررت في الإخراج
    return code[len(prompt):] if code.startswith(prompt) else code

# واجهة Gradio
gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=5, placeholder="Describe what code you want...", label="Prompt"),
    outputs=gr.Textbox(label="Generated Code"),
    title="Code Generator - Mono Model",
    description="Generate Python code from a text description using CodeGen-350M-Mono model"
).launch()