Create app.py

app.py

@@ -0,0 +1,36 @@
+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()