Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,11 +6,11 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # ==============================
 # PAGE CONFIG
 # ==============================
-st.set_page_config(page_title="💻 AI Coding Assistant", layout="wide")
-st.title("💻 AI Coding Assistant")
+st.set_page_config(page_title="🤖 AI Assistant", layout="wide")
+st.title("🤖 Simple AI Assistant")
 
 # ==============================
-# LOAD MODEL (HF SAFE)
+# LOAD MODEL
 # ==============================
 @st.cache_resource
 def load_model():
@@ -38,16 +38,11 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 
 # ==============================
-# CLEAN CODE EXTRACTION
+# CLEAN OUTPUT
 # ==============================
-def extract_code(text):
-    match = re.search(r"```python(.*?)```", text, re.DOTALL)
-
-    if match:
-        return match.group(1).strip()
-
-    # fallback cleaning
+def clean_text(text):
     text = re.sub(r"[^\x00-\x7F]+", "", text)
+    text = text.replace("```", "")
     return text.strip()
 
 # ==============================
@@ -55,23 +50,8 @@ def extract_code(text):
 # ==============================
 def generate_response(user_input):
 
-    prompt = f"""
-You are a professional coding assistant.
-
-Return ONLY valid code.
-
-Rules:
-- Only code
-- No explanation
-- Use simple syntax
-- Must be complete
-
-User Request:
-{user_input}
-
-Output:
-```python
-"""
+    prompt = f"""User: {user_input}
+Assistant:"""
 
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
@@ -79,70 +59,53 @@ Output:
         outputs = model.generate(
             **inputs,
             max_new_tokens=120,
-            do_sample=False,          # 🔥 deterministic (important)
-            temperature=0.0,
-            repetition_penalty=1.2,
+            do_sample=True,
+            temperature=0.5,
+            top_p=0.9,
+            repetition_penalty=1.1,
             pad_token_id=tokenizer.eos_token_id
         )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return extract_code(result)
+    # Extract only assistant reply
+    if "Assistant:" in result:
+        result = result.split("Assistant:")[-1]
+
+    return clean_text(result)
 
 # ==============================
 # 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"])
+        st.markdown(msg["content"])
 
 # ==============================
-# CHAT INPUT
+# INPUT BOX
 # ==============================
-user_input = st.chat_input("Ask your coding question...")
+user_input = st.chat_input("Type your message...")
 
 if user_input:
-    # Show user message
+    # 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("💡 Generating code..."):
+    # AI response
+    with st.spinner("🤖 Thinking..."):
         response = generate_response(user_input)
 
-    # Show assistant response
     st.session_state.messages.append({"role": "assistant", "content": response})
 
     with st.chat_message("assistant"):
-        st.code(response, language="python")
+        st.markdown(response)
 
 # ==============================
-# QUICK ACTION BUTTONS
+# SIDEBAR
 # ==============================
-st.sidebar.title("⚡ Actions")
+st.sidebar.title("⚙️ Options")
 
 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()
+    st.session_state.messages = []