import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Select a model
model_name = "Salesforce/codegen-2B-mono"  # Ensure this model is available
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function to generate code based on a prompt
def generate_code(prompt):
    # Adjust parameters to improve output quality
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,       # Adjust as needed for code length
        temperature=0.3,          # Lower temperature for more deterministic output
        top_p=0.9,                # Top-p filtering to focus on more likely completions
        repetition_penalty=1.2,   # Penalizes repetitive phrases
        do_sample=True            # Enables sampling for a creative touch
    )
    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_code

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=5, label="Enter your prompt"),
    outputs=gr.Code(language="python", label="Generated Code"),
    title="Python Code Generator",
    description="Enter a description of the Python code you want to generate."
)

# Launch the interface
iface.launch()