Update app.py

app.py

@@ -1,36 +1,69 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import re
 
-# 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.
+# Load CodeGen (Python-specialised)
 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."
-)
+def generate_code(user_request):
+    """
+    Produce clean Python code from a natural language instruction.
+    """
+
+    # A structured prompt works significantly better with CodeGen.
+    prompt = (
+        "# Task: Write Python code that accomplishes the following:\n"
+        f"# {user_request}\n"
+        "# Code:\n"
+    )
+
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+    # Deterministic decoding avoids messy repetition.
+    output_ids = model.generate(
+        input_ids,
+        max_length=256,
+        num_beams=4,
+        do_sample=False,
+        eos_token_id=tokenizer.eos_token_id
+    )
+
+    full_output = tokenizer.decode(output_ids[0])
+
+    # Remove the prompt section so that only the generated code remains.
+    code_only = full_output.split("# Code:\n", 1)[-1]
+
+    # Strip trailing text the model sometimes adds.
+    code_only = code_only.strip()
+
+    # Remove accidental markdown or stray tokens
+    code_only = re.sub(r"<\|.*?\|>", "", code_only)
+
+    return code_only
+
+
+with gr.Blocks(title="Code Generation with CodeGen-2B") as demo:
+
+    gr.Markdown(
+        """### Code Generation Assistant  
+Provide a description of the code you need, and the model will return Python code only.
+        """
+    )
+
+    task = gr.Textbox(
+        lines=2,
+        label="Task Description",
+        placeholder="For example: create a function that prints the first n Fibonacci numbers."
+    )
+
+    output = gr.Code(
+        label="Generated Python Code",
+        language="python"
+    )
+
+    btn = gr.Button("Generate Code")
+    btn.click(generate_code, inputs=task, outputs=output)
 
 demo.launch()
+