Update app.py

app.py

@@ -2,10 +2,9 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_id = "lakshraina2/leetcodeAI" 
-
-print("Loading model on CPU...")
+model_id = "lakshraina2/leetcode-coder-1.5B" 
 
+print("Loading model...")
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=False)
 model = AutoModelForCausalLM.from_pretrained(
     model_id, 
@@ -14,30 +13,32 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 def solve(problem_text):
-    # Basic check to ensure input isn't empty
-    if not problem_text:
-        return "No problem text detected."
-        
-    prompt = f"### Instruction:\nSolve this LeetCode problem:\n{problem_text}\n\n### Response:\n"
+    if not problem_text or len(problem_text) < 10:
+        return "// Error: Problem text too short or not scraped correctly."
+
+    # Standard Alpaca/Llama prompt format
+    prompt = f"Below is a LeetCode problem. Write a complete Python solution.\n\n### Problem:\n{problem_text}\n\n### Solution:\n"
+    
     inputs = tokenizer(prompt, return_tensors="pt")
     
+    # Generate with specific constraints to prevent empty output
     with torch.no_grad():
         outputs = model.generate(
-            **inputs, 
-            max_new_tokens=512,
-            temperature=0.2,
-            do_sample=True
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            max_new_tokens=1024, # Increased for complex problems
+            min_new_tokens=50,   # Force the model to talk
+            temperature=0.1,     # Lower temperature = more focused/less random
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
     
-    solution = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return solution.split("### Response:\n")[-1].strip()
-
-# THE FIX: We use gr.Interface but explicitly name the API endpoint 'predict'
-demo = gr.Interface(
-    fn=solve, 
-    inputs=gr.Textbox(), 
-    outputs=gr.Textbox(),
-    api_name="predict" # This matches the /predict in your content.js URL
-)
+    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the part after our '### Solution:' marker
+    if "### Solution:" in full_text:
+        return full_text.split("### Solution:")[-1].strip()
+    return full_text.strip()
 
+demo = gr.Interface(fn=solve, inputs="text", outputs="text", api_name="predict")
 demo.launch()