Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,36 +2,24 @@ import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ==============================
-# 🔐 HIDE STREAMLIT MENU
-# ==============================
-st.markdown("""
-<style>
-#MainMenu {visibility: hidden;}
-header {visibility: hidden;}
-footer {visibility: hidden;}
-.stDeployButton {display:none;}
-</style>
-""", unsafe_allow_html=True)
-
 # ==============================
 # PAGE CONFIG
 # ==============================
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# LOAD MODEL (FAST + STABLE)
+# LOAD MODEL (DeepSeek - BEST)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "microsoft/phi-2"
+    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,   # CPU friendly
-        device_map=None
+        torch_dtype=torch.float32,
+        device_map="cpu"
     )
 
     return tokenizer, model
@@ -39,17 +27,26 @@ def load_model():
 tokenizer, model = load_model()
 
 # ==============================
-# CODE GENERATION FUNCTION
+# AUTO FIX CODE (IMPORTANT)
 # ==============================
-def generate_code(prompt, language):
+def fix_incomplete_code(code):
+    # Fix missing brackets (basic handling)
+    if code.count("(") > code.count(")"):
+        code += ")"
+    if code.count("{") > code.count("}"):
+        code += "}"
+    return code.strip()
 
-    full_prompt = f"""### Instruction:
-Write ONLY valid {language} code.
+# ==============================
+# GENERATE CODE
+# ==============================
+def generate_code(prompt, language):
 
-Do not include explanations.
-Do not include special tokens.
+    full_prompt = f"""
+### Instruction:
+Write a correct {language} function.
 
-Task:
+### Task:
 {prompt}
 
 ### Response:
@@ -57,37 +54,27 @@ Task:
 
     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,
+            do_sample=False,
+            temperature=0.0,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract response
-    if "### Response:" in result:
-        result = result.split("### Response:")[-1]
-
-    # Remove unwanted tokens
-    unwanted_tokens = [
-        "<|endoftext|>",
-        "<|file_separator|>",
-        "<|assistant|>",
-        "<|system|>"
-    ]
-
-    for token in unwanted_tokens:
-        result = result.replace(token, "")
+    # Remove prompt
+    result = result.replace(full_prompt, "").strip()
 
-    return result.strip()
+    return fix_incomplete_code(result)
 
 # ==============================
 # UI
 # ==============================
-st.title("💻 AI Code Generator (Fast & Clean)")
+st.title("💻 AI Code Generator (HF Optimized)")
 
 col1, col2 = st.columns(2)
 
@@ -100,19 +87,12 @@ with col2:
         ["Python", "JavaScript", "SQL", "Java", "C++", "HTML", "CSS"]
     )
 
-# ==============================
-# 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("⚠️ Enter a task")
     else:
-        with st.spinner("Generating code..."):
-            try:
-                code = generate_code(user_prompt, language)
-
-                st.success("✅ Generated Code")
-                st.code(code, language=language.lower())
+        with st.spinner("⚡ Generating code..."):
+            code = generate_code(user_prompt, language)
 
-            except Exception as e:
-                st.error(f"❌ Error: {str(e)}")
+        st.success("✅ Generated Code")
+        st.code(code, language=language.lower())