Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,18 +8,18 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# LOAD MODEL (DeepSeek - BEST)
+# LOAD MODEL
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
+    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"   # 🔥 BEST MODEL
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float32,
-        device_map="cpu"
+        device_map="auto"   # auto uses GPU if available
     )
 
     return tokenizer, model
@@ -27,14 +27,24 @@ def load_model():
 tokenizer, model = load_model()
 
 # ==============================
-# AUTO FIX CODE (IMPORTANT)
+# CLEAN OUTPUT
 # ==============================
-def fix_incomplete_code(code):
-    # Fix missing brackets (basic handling)
-    if code.count("(") > code.count(")"):
-        code += ")"
-    if code.count("{") > code.count("}"):
-        code += "}"
+def clean_code(code):
+    code = code.strip()
+
+    remove_words = [
+        "Explanation:",
+        "Here is the code:",
+        "Output:",
+        "Answer:"
+    ]
+
+    for word in remove_words:
+        code = code.replace(word, "")
+
+    # Remove markdown formatting
+    code = code.replace("```python", "").replace("```", "")
+
     return code.strip()
 
 # ==============================
@@ -43,38 +53,53 @@ def fix_incomplete_code(code):
 def generate_code(prompt, language):
 
     full_prompt = f"""
-### Instruction:
-Write a correct {language} function.
+You are an expert {language} developer.
 
-### Task:
+Generate clean, correct, and complete code.
+
+Rules:
+- Only return code
+- No explanation
+- Proper syntax
+- Complete working solution
+
+Task:
 {prompt}
 
-### Response:
+Code:
 """
 
-    inputs = tokenizer(full_prompt, return_tensors="pt")
+    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=200,
-            do_sample=False,
-            temperature=0.0,
+            max_new_tokens=350,
+            do_sample=True,
+            temperature=0.2,
+            top_p=0.9,
+            repetition_penalty=1.1,
             eos_token_id=tokenizer.eos_token_id,
             pad_token_id=tokenizer.eos_token_id
         )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Remove prompt
-    result = result.replace(full_prompt, "").strip()
+    if "Code:" in result:
+        result = result.split("Code:")[-1]
+
+    return clean_code(result)
 
-    return fix_incomplete_code(result)
+# ==============================
+# SESSION STATE (CHAT HISTORY)
+# ==============================
+if "history" not in st.session_state:
+    st.session_state.history = []
 
 # ==============================
 # UI
 # ==============================
-st.title("💻 AI Code Generator (HF Optimized)")
+st.title("💻 AI Code Generator (Advanced)")
 
 col1, col2 = st.columns(2)
 
@@ -87,12 +112,21 @@ with col2:
         ["Python", "JavaScript", "SQL", "Java", "C++", "HTML", "CSS"]
     )
 
+# ==============================
+# GENERATE BUTTON
+# ==============================
 if st.button("Generate Code"):
     if not user_prompt.strip():
-        st.warning("⚠️ Enter a task")
+        st.warning("⚠️ Please enter a task")
     else:
-        with st.spinner("⚡ Generating code..."):
+        with st.spinner("⚡ Generating high-quality code..."):
             code = generate_code(user_prompt, language)
 
-        st.success("✅ Generated Code")
-        st.code(code, language=language.lower())
+        st.session_state.history.append((user_prompt, code))
+
+# ==============================
+# DISPLAY HISTORY
+# ==============================
+for q, c in reversed(st.session_state.history):
+    st.markdown(f"### 🧑 Task:\n{q}")
+    st.code(c, language=language.lower())