Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import torch
+import re
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # ==============================
@@ -9,7 +10,7 @@ st.set_page_config(page_title="💻 AI Coding Assistant", layout="wide")
 st.title("💻 AI Coding Assistant")
 
 # ==============================
-# LOAD MODEL (LIGHTWEIGHT)
+# LOAD MODEL (HF SAFE)
 # ==============================
 @st.cache_resource
 def load_model():
@@ -31,54 +32,72 @@ with st.spinner("🔄 Loading model..."):
 st.success("✅ Ready")
 
 # ==============================
-# SESSION STATE (CHAT HISTORY)
+# SESSION STATE
 # ==============================
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
 # ==============================
-# CHAT DISPLAY
+# CLEAN CODE EXTRACTION
 # ==============================
-for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        if msg["role"] == "assistant":
-            st.code(msg["content"])
-        else:
-            st.markdown(msg["content"])
+def extract_code(text):
+    match = re.search(r"```python(.*?)```", text, re.DOTALL)
+
+    if match:
+        return match.group(1).strip()
+
+    # fallback cleaning
+    text = re.sub(r"[^\x00-\x7F]+", "", text)
+    return text.strip()
 
 # ==============================
 # GENERATE RESPONSE
 # ==============================
 def generate_response(user_input):
 
-    # Build conversation prompt
-    conversation = ""
-    for msg in st.session_state.messages:
-        role = "User" if msg["role"] == "user" else "Assistant"
-        conversation += f"{role}: {msg['content']}\n"
+    prompt = f"""
+You are a professional coding assistant.
+
+Return ONLY valid code.
+
+Rules:
+- Only code
+- No explanation
+- Use simple syntax
+- Must be complete
 
-    conversation += f"User: {user_input}\nAssistant:"
+User Request:
+{user_input}
 
-    inputs = tokenizer(conversation, return_tensors="pt", truncation=True)
+Output:
+```python
+"""
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=200,
-            do_sample=True,
-            temperature=0.3,
-            top_p=0.9,
-            repetition_penalty=1.1,
+            max_new_tokens=120,
+            do_sample=False,          # 🔥 deterministic (important)
+            temperature=0.0,
+            repetition_penalty=1.2,
             pad_token_id=tokenizer.eos_token_id
         )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract only assistant reply
-    if "Assistant:" in result:
-        result = result.split("Assistant:")[-1]
+    return extract_code(result)
 
-    return result.strip()
+# ==============================
+# DISPLAY CHAT
+# ==============================
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        if msg["role"] == "assistant":
+            st.code(msg["content"], language="python")
+        else:
+            st.markdown(msg["content"])
 
 # ==============================
 # CHAT INPUT
@@ -86,18 +105,44 @@ def generate_response(user_input):
 user_input = st.chat_input("Ask your coding question...")
 
 if user_input:
-    # Add user message
+    # Show user message
     st.session_state.messages.append({"role": "user", "content": user_input})
 
     with st.chat_message("user"):
         st.markdown(user_input)
 
     # Generate response
-    with st.spinner("💡 Thinking..."):
+    with st.spinner("💡 Generating code..."):
         response = generate_response(user_input)
 
-    # Add assistant message
+    # Show assistant response
     st.session_state.messages.append({"role": "assistant", "content": response})
 
     with st.chat_message("assistant"):
-        st.code(response)
+        st.code(response, language="python")
+
+# ==============================
+# QUICK ACTION BUTTONS
+# ==============================
+st.sidebar.title("⚡ Actions")
+
+if st.sidebar.button("🧹 Clear Chat"):
+    st.session_state.messages = []
+
+if st.sidebar.button("📖 Explain Last Code"):
+    if st.session_state.messages:
+        last_code = st.session_state.messages[-1]["content"]
+        st.session_state.messages.append({
+            "role": "user",
+            "content": f"Explain this code:\n{last_code}"
+        })
+        st.experimental_rerun()
+
+if st.sidebar.button("🛠 Fix Last Code"):
+    if st.session_state.messages:
+        last_code = st.session_state.messages[-1]["content"]
+        st.session_state.messages.append({
+            "role": "user",
+            "content": f"Fix this code:\n{last_code}"
+        })
+        st.experimental_rerun()