Update src/streamlit_app.py

src/streamlit_app.py

@@ -20,18 +20,18 @@ footer {visibility: hidden;}
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# LOAD MODEL (FAST)
+# LOAD MODEL (FAST + CPU SAFE)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "microsoft/phi-2"
+    model_name = "google/codegemma-2b"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,   # CPU friendly
-        device_map=None              # avoid GPU issues
+        torch_dtype=torch.float32,   # CPU safe
+        device_map="cpu"             # force CPU for Spaces
     )
 
     return tokenizer, model
@@ -39,36 +39,57 @@ def load_model():
 tokenizer, model = load_model()
 
 # ==============================
-# CODE GENERATION FUNCTION
+# CLEAN OUTPUT FUNCTION
+# ==============================
+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, "")
+
+    return text.strip()
+
+# ==============================
+# CODE GENERATION FUNCTION (FIXED)
 # ==============================
 def generate_code(prompt, language):
 
-    full_prompt = f"""### Instruction:
-Write a {language} program for the following task.
+    full_prompt = f"""
+You are a highly accurate {language} code generator.
 
-Task:
-{prompt}
+Example:
+Task: add two numbers
+Code:
+function add(a, b) {{
+    return a + b;
+}}
 
-### Response:
+Now solve:
+
+Task: {prompt}
+
+Instructions:
+- Generate correct and complete {language} code
+- Do exactly what is asked
+- Do NOT change the logic
+- Return ONLY code
 """
 
     inputs = tokenizer(full_prompt, return_tensors="pt")
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=150,
-        temperature=0.2,
-        top_p=0.9,
-        do_sample=False
-    )
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,     # prevent truncation
+            do_sample=False,        # deterministic (IMPORTANT)
+            temperature=0.0
+        )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract only response part
-    if "### Response:" in result:
-        result = result.split("### Response:")[-1]
-
-    return result.strip()
+    return clean_output(result, full_prompt)
 
 # ==============================
 # UI
@@ -87,18 +108,16 @@ with col2:
     )
 
 # ==============================
-# GENERATE BUTTON
+# BUTTON ACTION
 # ==============================
-if st.button("🚀 Generate Code"):
+if st.button("Generate Code"):
     if not user_prompt.strip():
-        st.warning("Please enter a task")
+        st.warning("⚠️ Please enter a task")
     else:
-        with st.spinner("Generating code..."):
-            try:
-                code = generate_code(user_prompt, language)
+        with st.spinner("⚡ Generating high-quality code..."):
+            code = generate_code(user_prompt, language)
 
-                st.success("✅ Generated Code")
-                st.code(code, language=language.lower())
+        st.success("✅ Generated Code")
 
-            except Exception as e:
-                st.error(f"❌ Error: {str(e)}")
+        # Display properly formatted code
+        st.code(code, language=language.lower())