Update app.py

app.py

@@ -1,20 +1,25 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Use a lightweight, free-tier friendly model
+model_id = "Salesforce/codegen-350M-mono"
 
-model_id = "Salesforce/codegen-350M-mono"  # SMALL model that works on HF free tier
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id)
 
 def generate_code(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_new_tokens=256, do_sample=True)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    outputs = model.generate(**inputs, max_new_tokens=300, do_sample=True)
+    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return code
 
-gr.Interface(
+# Gradio app
+demo = gr.Interface(
     fn=generate_code,
-    inputs=gr.Textbox(label="Enter Prompt", lines=4, placeholder="e.g., make a login form using HTML and CSS"),
+    inputs=gr.Textbox(label="Describe your code idea"),
     outputs=gr.Code(label="Generated Code"),
-    title="IMMORTAL Text to Code AI",
-    description="Enter a prompt to generate code using AI (based on CodeGen 350M)."
-).launch()
+    title="Text to Code Generator",
+    description="Enter any code-related prompt and watch it generate live."
+)
+
+demo.launch()