pcasale/CodeGenerationAssistant

1from transformers import AutoTokenizer, AutoModelForCausalLM2import gradio as gr3import re4 5# Load CodeGen (Python-specialised)6tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-2B-mono")7model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-2B-mono")8 9def generate_code(user_request):10    """11    Produce clean Python code from a natural language instruction.12    """13 14    # A structured prompt works significantly better with CodeGen.15    prompt = (16        "# Task: Write Python code that accomplishes the following:\n"17        f"# {user_request}\n"18        "# Code:\n"19    )20 21    input_ids = tokenizer(prompt, return_tensors="pt").input_ids22 23    # Deterministic decoding avoids messy repetition.24    output_ids = model.generate(25        input_ids,26        max_length=256,27        num_beams=4,28        do_sample=False,29        eos_token_id=tokenizer.eos_token_id30    )31 32    full_output = tokenizer.decode(output_ids[0])33 34    # Remove the prompt section so that only the generated code remains.35    code_only = full_output.split("# Code:\n", 1)[-1]36 37    # Strip trailing text the model sometimes adds.38    code_only = code_only.strip()39 40    # Remove accidental markdown or stray tokens41    code_only = re.sub(r"<\|.*?\|>", "", code_only)42 43    return code_only44 45 46with gr.Blocks(title="Code Generation with CodeGen-2B") as demo:47 48    gr.Markdown(49        """### Code Generation Assistant  50Provide a description of the code you need, and the model will return Python code only.51        """52    )53 54    task = gr.Textbox(55        lines=2,56        label="Task Description",57        placeholder="For example: create a function that prints the first n Fibonacci numbers."58    )59 60    output = gr.Code(61        label="Generated Python Code",62        language="python"63    )64 65    btn = gr.Button("Generate Code")66    btn.click(generate_code, inputs=task, outputs=output)67 68demo.launch()69 70