Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,16 +1,15 @@
 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 (Advanced Clean Version)")
+st.set_page_config(page_title="💻 AI Coding Assistant", layout="wide")
+st.title("💻 AI Coding Assistant")
 
 # ==============================
-# LOAD MODEL (HF SAFE)
+# LOAD MODEL (LIGHTWEIGHT)
 # ==============================
 @st.cache_resource
 def load_model():
@@ -26,122 +25,79 @@ def load_model():
     model.eval()
     return tokenizer, model
 
-# Load model with spinner
-with st.spinner("🔄 Loading AI model..."):
+with st.spinner("🔄 Loading model..."):
     tokenizer, model = load_model()
 
-st.success("✅ Model Loaded Successfully")
+st.success("✅ Ready")
 
 # ==============================
-# CLEAN FUNCTION
+# SESSION STATE (CHAT HISTORY)
 # ==============================
-def extract_clean_code(text):
-
-    if "Code:" in text:
-        text = text.split("Code:")[-1]
-
-    # Remove non-ascii junk
-    text = re.sub(r"[^\x00-\x7F]+", "", text)
-
-    # Remove unwanted sections
-    text = re.sub(r"(Explanation:.*)", "", text, flags=re.DOTALL)
-
-    # Remove markdown
-    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 "\n".join(clean_lines).strip()
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
 # ==============================
-# GENERATE CODE
+# CHAT DISPLAY
 # ==============================
-def generate_code(prompt, language):
+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"])
 
-    full_prompt = f"""
-### Instruction:
-You are an expert {language} programmer.
-
-### Rules:
-- Only return code
-- No explanation
-- Use simple and correct syntax
-- Provide complete solution
+# ==============================
+# GENERATE RESPONSE
+# ==============================
+def generate_response(user_input):
 
-### Task:
-{prompt}
+    # 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"
 
-### Code:
-"""
+    conversation += f"User: {user_input}\nAssistant:"
 
-    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
+    inputs = tokenizer(conversation, return_tensors="pt", truncation=True)
 
-    try:
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=150,
-                do_sample=True,
-                temperature=0.2,
-                top_p=0.85,
-                repetition_penalty=1.2,
-                pad_token_id=tokenizer.eos_token_id
-            )
+    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,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return extract_clean_code(result)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    except Exception as e:
-        return f"# ERROR: {str(e)}"
+    # Extract only assistant reply
+    if "Assistant:" in result:
+        result = result.split("Assistant:")[-1]
 
-# ==============================
-# SESSION STATE
-# ==============================
-if "history" not in st.session_state:
-    st.session_state.history = []
+    return result.strip()
 
 # ==============================
-# UI INPUT
+# CHAT INPUT
 # ==============================
-col1, col2 = st.columns(2)
+user_input = st.chat_input("Ask your coding question...")
 
-with col1:
-    user_prompt = st.text_area("Describe your task", height=200)
+if user_input:
+    # Add user message
+    st.session_state.messages.append({"role": "user", "content": user_input})
 
-with col2:
-    language = st.selectbox(
-        "Select Programming Language",
-        ["Python", "JavaScript", "SQL", "Java", "C++", "HTML", "CSS"]
-    )
+    with st.chat_message("user"):
+        st.markdown(user_input)
 
-# ==============================
-# GENERATE BUTTON
-# ==============================
-if st.button("Generate Code"):
-    if not user_prompt.strip():
-        st.warning("⚠️ Please enter a task")
-    else:
-        with st.spinner("⚡ Generating clean code..."):
-            code = generate_code(user_prompt, language)
-
-        st.session_state.history.append((user_prompt, code))
+    # Generate response
+    with st.spinner("💡 Thinking..."):
+        response = generate_response(user_input)
 
-# ==============================
-# OUTPUT DISPLAY
-# ==============================
-if st.session_state.history:
-    st.subheader("📌 Generated Results")
+    # Add assistant message
+    st.session_state.messages.append({"role": "assistant", "content": response})
 
-    for q, c in reversed(st.session_state.history):
-        st.markdown(f"**🧑 Task:** {q}")
-        st.code(c, language=language.lower())
+    with st.chat_message("assistant"):
+        st.code(response)