Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
-MODEL_ID = "codellama/CodeLlama-7b-Instruct-hf"
+MODEL_ID = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 
 # Load once at startup (ZeroGPU handles GPU allocation automatically)
 bnb_config = BitsAndBytesConfig(
@@ -51,29 +51,36 @@ Please review the following {language} code:
 
 
 def review_code(code: str, language: str, max_new_tokens: int = 800):
-    if not code.strip():
-        return "⚠️ Please paste some code to review."
-
-    prompt = build_prompt(code, language)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.2,       # low temp = more deterministic, better for analysis
-            do_sample=True,
-            repetition_penalty=1.1,
-            pad_token_id=tokenizer.eos_token_id,
+    try:
+        if not code.strip():
+            return "⚠️ Please paste some code to review."
+    
+        prompt = build_prompt(code, language)
+        inputs = tokenizer(prompt, return_tensors="pt")
+    
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=512,
+                temperature=0.2,
+                do_sample=True,
+                top_p=0.95,
+                repetition_penalty=1.1,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+    
+        # Decode only the newly generated tokens (skip the prompt)
+        response = tokenizer.decode(
+            outputs[0][inputs["input_ids"].shape[1]:],
+            skip_special_tokens=True
         )
-
-    # Decode only the newly generated tokens (skip the prompt)
-    response = tokenizer.decode(
-        outputs[0][inputs["input_ids"].shape[1]:],
-        skip_special_tokens=True
-    )
-    return response.strip() or "⚠️ Model returned an empty response."
-
+        return response.strip() or "⚠️ Model returned an empty response."
+    except Exception as e:
+        return f"❌ Error:\n\n{str(e)}"
 
 # --- Gradio UI ---
 with gr.Blocks(title="AI Code Reviewer", theme=gr.themes.Soft()) as demo:
@@ -113,4 +120,4 @@ with gr.Blocks(title="AI Code Reviewer", theme=gr.themes.Soft()) as demo:
         label="Try an example"
     )
 
-demo.launch()
+demo.queue().launch()