from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# Load the CodeGen model and tokenizer. This model has 2B parameters and is specialized for code generation:contentReference[oaicite:16]{index=16}.
# Note: Downloading and loading this model may be slow, and it may require a GPU for reasonable performance.
tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-mono")
model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-2B-mono")

# Define the code generation function.
def generate_code(prompt):
    # Format the prompt as a comment, because CodeGen is trained to take a code-related prompt in comment form:contentReference[oaicite:17]{index=17}.
    formatted_prompt = f"# {prompt}\n"  
    # Tokenize the prompt and get input IDs for the model
    input_ids = tokenizer.encode(formatted_prompt, return_tensors="pt")
    # Use the model to generate code. We set a limit on max_length for the output.
    # We also use a low temperature (0.2) to make the output more deterministic and focused.
    output_ids = model.generate(input_ids, max_length=256, num_beams=1, do_sample=True, temperature=0.2)
    # Decode the generated tokens back into a string of code.
    generated_code = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    return generated_code

# Set up Gradio interface with a textbox for the task description and a code output component.
input_desc = gr.Textbox(lines=2, label="Task Description", placeholder="Describe the code you need...")
output_code = gr.Code(language="python", label="Generated Code")

demo = gr.Interface(
    fn=generate_code,
    inputs=input_desc,
    outputs=output_code,
    title="💻 Code Generation Assistant (CodeGen-2B)",
    description="**Description:** Provide a natural language description of a programming task, "
                "and the model will generate Python code to accomplish the task. "
                "Uses Salesforce's CodeGen-2B-mono model (2B parameters) for code generation."
)

demo.launch()