Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,20 +1,20 @@
 import streamlit as st
 import torch
+import re
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # ==============================
 # PAGE CONFIG
 # ==============================
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
-
-st.title("💻 AI Code Generator (Stable Version)")
+st.title("💻 AI Code Generator (Advanced Clean Version)")
 
 # ==============================
-# LOAD MODEL (SAFE)
+# LOAD MODEL (HF SAFE)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"  # ✅ HF Free Safe
+    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -26,33 +26,42 @@ def load_model():
     model.eval()
     return tokenizer, model
 
-# Load with spinner (IMPORTANT)
-with st.spinner("🔄 Loading AI model... Please wait"):
+# Load model with spinner
+with st.spinner("🔄 Loading AI model..."):
     tokenizer, model = load_model()
 
 st.success("✅ Model Loaded Successfully")
 
 # ==============================
-# CLEAN OUTPUT
+# CLEAN FUNCTION
 # ==============================
-def clean_code(code):
-    code = code.strip()
+def extract_clean_code(text):
+
+    if "Code:" in text:
+        text = text.split("Code:")[-1]
 
-    # Remove unwanted text
-    remove_words = [
-        "Explanation:",
-        "Here is the code:",
-        "Output:",
-        "Answer:"
-    ]
+    # Remove non-ascii junk
+    text = re.sub(r"[^\x00-\x7F]+", "", text)
 
-    for word in remove_words:
-        code = code.replace(word, "")
+    # Remove unwanted sections
+    text = re.sub(r"(Explanation:.*)", "", text, flags=re.DOTALL)
 
     # Remove markdown
-    code = code.replace("```python", "").replace("```", "")
+    text = text.replace("```python", "").replace("```", "")
+
+    # Filter useful lines only
+    lines = text.split("\n")
+    clean_lines = []
+
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        if any(word in line.lower() for word in ["instruction", "task", "response"]):
+            continue
+        clean_lines.append(line)
 
-    return code.strip()
+    return "\n".join(clean_lines).strip()
 
 # ==============================
 # GENERATE CODE
@@ -60,19 +69,19 @@ def clean_code(code):
 def generate_code(prompt, language):
 
     full_prompt = f"""
+### Instruction:
 You are an expert {language} programmer.
 
-Write clean, correct, and complete code.
-
-Rules:
+### Rules:
 - Only return code
-- No explanations
-- Complete solution
+- No explanation
+- Use simple and correct syntax
+- Provide complete solution
 
-Task:
+### Task:
 {prompt}
 
-Code:
+### Code:
 """
 
     inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
@@ -81,20 +90,16 @@ Code:
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=200,
+                max_new_tokens=150,
                 do_sample=True,
-                temperature=0.3,
-                top_p=0.9,
-                repetition_penalty=1.1,
+                temperature=0.2,
+                top_p=0.85,
+                repetition_penalty=1.2,
                 pad_token_id=tokenizer.eos_token_id
             )
 
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        if "Code:" in result:
-            result = result.split("Code:")[-1]
-
-        return clean_code(result)
+        return extract_clean_code(result)
 
     except Exception as e:
         return f"# ERROR: {str(e)}"
@@ -126,7 +131,7 @@ if st.button("Generate Code"):
     if not user_prompt.strip():
         st.warning("⚠️ Please enter a task")
     else:
-        with st.spinner("⚡ Generating code..."):
+        with st.spinner("⚡ Generating clean code..."):
             code = generate_code(user_prompt, language)
 
         st.session_state.history.append((user_prompt, code))