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
+# LOAD MODEL (FAST + STABLE)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "google/codegemma-2b"
+    model_name = "microsoft/phi-2"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,
-        device_map="cpu"
+        torch_dtype=torch.float32,   # CPU friendly
+        device_map=None
     )
 
     return tokenizer, model
@@ -39,50 +39,55 @@ def load_model():
 tokenizer, model = load_model()
 
 # ==============================
-# CLEAN OUTPUT (IMPORTANT FIX)
-# ==============================
-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()
-
-# ==============================
-# GENERATE CODE (SIMPLIFIED PROMPT)
+# CODE GENERATION FUNCTION
 # ==============================
 def generate_code(prompt, language):
 
-    full_prompt = f"""
-Write a {language} function for the following task:
+    full_prompt = f"""### Instruction:
+Write ONLY valid {language} code.
+
+Do not include explanations.
+Do not include special tokens.
 
+Task:
 {prompt}
 
-Only return code.
+### Response:
 """
 
     inputs = tokenizer(full_prompt, return_tensors="pt")
 
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=120,
-            do_sample=False,
-            temperature=0.0
-        )
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,
+        temperature=0.2,
+        top_p=0.9,
+        do_sample=False
+    )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    result = result.replace(full_prompt, "").strip()
+    # Extract response
+    if "### Response:" in result:
+        result = result.split("### Response:")[-1]
+
+    # Remove unwanted tokens
+    unwanted_tokens = [
+        "<|endoftext|>",
+        "<|file_separator|>",
+        "<|assistant|>",
+        "<|system|>"
+    ]
 
-    return extract_code(result)
+    for token in unwanted_tokens:
+        result = result.replace(token, "")
+
+    return result.strip()
 
 # ==============================
 # UI
 # ==============================
-st.title("💻 AI Code Generator (Fast & Accurate)")
+st.title("💻 AI Code Generator (Fast & Clean)")
 
 col1, col2 = st.columns(2)
 
@@ -96,14 +101,18 @@ with col2:
     )
 
 # ==============================
-# BUTTON
+# GENERATE BUTTON
 # ==============================
-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 clean code..."):
-            code = generate_code(user_prompt, language)
+        with st.spinner("Generating code..."):
+            try:
+                code = generate_code(user_prompt, language)
+
+                st.success("✅ Generated Code")
+                st.code(code, language=language.lower())
 
-        st.success("✅ Generated Code")
-        st.code(code, language=language.lower())
+            except Exception as e:
+                st.error(f"❌ Error: {str(e)}")