Update app.py

app.py

@@ -2,49 +2,45 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# ✅ Load SantaCoder model
+# ✅ Load model and tokenizer
 model_name = "bigcode/santacoder"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+model.eval()
 
-# ✅ Define the assistant function
+# ✅ Function to generate suggestions
 def code_assistant(code, instruction, language):
     if not code.strip() or not instruction.strip():
-        return "⚠️ Please enter both code and instruction."
+        return "⚠️ Please provide both code and instruction."
 
-    prompt = f"""### Language: {language}
-### Instruction: {instruction}
-### Code:
-{code}
-### Suggestion:
-"""
+    # Clear and short prompt
+    prompt = f"# Language: {language}\n# Task: {instruction}\n{code}\n# Suggestion:\n"
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-    input_ids = inputs["input_ids"].to(model.device)
+    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.to(model.device)
 
     with torch.no_grad():
         output_ids = model.generate(
-            input_ids=input_ids,
-            max_new_tokens=256,
-            pad_token_id=tokenizer.eos_token_id,
+            input_ids,
+            max_new_tokens=128,
             do_sample=True,
-            temperature=0.7
+            temperature=0.6,
+            pad_token_id=tokenizer.eos_token_id,
         )
 
     full_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    # Extract only the suggestion
-    if "### Suggestion:" in full_output:
-        suggestion = full_output.split("### Suggestion:")[-1].strip()
-    else:
-        suggestion = full_output[len(prompt):].strip()
+    # Remove prompt from output
+    suggestion = full_output.replace(prompt, "").strip()
 
-    return suggestion or "⚠️ No suggestion generated. Try again with better instruction."
+    # Avoid garbage suggestions
+    if "Constant" in suggestion[:30]:
+        return "⚠️ Couldn't generate a valid suggestion. Try rephrasing your instruction."
+
+    return suggestion or "⚠️ No meaningful suggestion generated."
 
 # ✅ Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## 🤖 Multilingual AI Code Assistant")
-    gr.Markdown("Paste your code, describe what you want (e.g., Explain, Fix bugs), and get instant help!")
 
     with gr.Row():
         lang = gr.Dropdown(
@@ -53,8 +49,8 @@ with gr.Blocks() as demo:
             value="Python"
         )
         task = gr.Textbox(label="📌 Task (Explain, Fix, Improve...)", placeholder="Explain this code")
-
-    code_input = gr.Code(label="✍️ Your Code", language="python")
+    
+    code_input = gr.Code(label="✍️ Paste your code here", language="python")
     output = gr.Textbox(label="💡 AI Suggestion", lines=10)
 
     submit = gr.Button("🚀 Generate Suggestion")