Update app.py

app.py

@@ -2,34 +2,32 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Point this to the repository you just created in Step 1
-model_id = "lakshraina2/leetcodeAI"
+# Replace with your actual merged model repo
+model_id = "lakshraina2/leetcode-coder-1.5B" 
 
-print("Downloading and loading model...")
+print("Loading model on CPU...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-# We load in float32 because the free tier uses standard CPU instances
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    torch_dtype=torch.float32, # CPU needs float32
+    device_map={"": "cpu"}     # Force CPU
+)
 
-def solve_problem(problem_text):
+def solve(problem_text):
     prompt = f"### Instruction:\nSolve this LeetCode problem:\n{problem_text}\n\n### Response:\n"
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt") # No .to("cuda")!
     
-    outputs = model.generate(
-        **inputs, 
-        max_new_tokens=512,
-        temperature=0.2, 
-        do_sample=True
-    )
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs, 
+            max_new_tokens=512,
+            temperature=0.2,
+            do_sample=True
+        )
     
     solution = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    try:
-        code_only = solution.split("### Response:\n")[1].strip()
-    except IndexError:
-        code_only = solution
-        
-    return code_only
+    return solution.split("### Response:\n")[-1].strip()
 
-# Gradio automatically builds an API around this function
-iface = gr.Interface(fn=solve_problem, inputs="text", outputs="text")
+# Gradio 4 interface
+iface = gr.Interface(fn=solve, inputs="text", outputs="text")
 iface.launch()