Update src/streamlit_app.py

src/streamlit_app.py

@@ -20,7 +20,7 @@ footer {visibility: hidden;}
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# LOAD MODEL (FAST + CPU SAFE)
+# LOAD MODEL
 # ==============================
 @st.cache_resource
 def load_model():
@@ -30,8 +30,8 @@ def load_model():
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,   # CPU safe
-        device_map="cpu"             # force CPU for Spaces
+        torch_dtype=torch.float32,
+        device_map="cpu"
     )
 
     return tokenizer, model
@@ -39,42 +39,28 @@ def load_model():
 tokenizer, model = load_model()
 
 # ==============================
-# CLEAN OUTPUT FUNCTION
+# CLEAN OUTPUT (IMPORTANT FIX)
 # ==============================
-def clean_output(text, prompt):
-    text = text.replace(prompt, "").strip()
-
-    # Remove unwanted prefixes if model adds them
-    unwanted_phrases = ["```", "code:", "Code:"]
-    for phrase in unwanted_phrases:
-        text = text.replace(phrase, "")
+def extract_code(text):
+    # Try to extract code block if exists
+    if "```" in text:
+        parts = text.split("```")
+        if len(parts) >= 2:
+            return parts[1].strip()
 
     return text.strip()
 
 # ==============================
-# CODE GENERATION FUNCTION (FIXED)
+# GENERATE CODE (SIMPLIFIED PROMPT)
 # ==============================
 def generate_code(prompt, language):
 
     full_prompt = f"""
-You are a highly accurate {language} code generator.
-
-Example:
-Task: add two numbers
-Code:
-function add(a, b) {{
-    return a + b;
-}}
-
-Now solve:
+Write a {language} function for the following task:
 
-Task: {prompt}
+{prompt}
 
-Instructions:
-- Generate correct and complete {language} code
-- Do exactly what is asked
-- Do NOT change the logic
-- Return ONLY code
+Only return code.
 """
 
     inputs = tokenizer(full_prompt, return_tensors="pt")
@@ -82,14 +68,16 @@ Instructions:
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=200,     # prevent truncation
-            do_sample=False,        # deterministic (IMPORTANT)
+            max_new_tokens=120,
+            do_sample=False,
             temperature=0.0
         )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return clean_output(result, full_prompt)
+    result = result.replace(full_prompt, "").strip()
+
+    return extract_code(result)
 
 # ==============================
 # UI
@@ -108,16 +96,14 @@ with col2:
     )
 
 # ==============================
-# BUTTON ACTION
+# BUTTON
 # ==============================
 if st.button("Generate Code"):
     if not user_prompt.strip():
         st.warning("⚠️ Please enter a task")
     else:
-        with st.spinner("⚡ Generating high-quality code..."):
+        with st.spinner("⚡ Generating clean code..."):
             code = generate_code(user_prompt, language)
 
         st.success("✅ Generated Code")
-
-        # Display properly formatted code
         st.code(code, language=language.lower())