Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,11 +10,11 @@ st.set_page_config(page_title="🤖 AI Assistant", layout="wide")
 st.title("🤖 Simple AI Assistant")
 
 # ==============================
-# LOAD MODEL
+# LOAD MODEL (CHAT MODEL ✅)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # 🔥 BEST FOR HF FREE
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -38,11 +38,10 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 
 # ==============================
-# CLEAN OUTPUT
+# CLEAN TEXT
 # ==============================
 def clean_text(text):
     text = re.sub(r"[^\x00-\x7F]+", "", text)
-    text = text.replace("```", "")
     return text.strip()
 
 # ==============================
@@ -50,17 +49,16 @@ def clean_text(text):
 # ==============================
 def generate_response(user_input):
 
-    prompt = f"""User: {user_input}
-Assistant:"""
+    prompt = f"<|user|>\n{user_input}\n<|assistant|>\n"
 
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=120,
+            max_new_tokens=150,
             do_sample=True,
-            temperature=0.5,
+            temperature=0.7,
             top_p=0.9,
             repetition_penalty=1.1,
             pad_token_id=tokenizer.eos_token_id
@@ -68,9 +66,9 @@ Assistant:"""
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract only assistant reply
-    if "Assistant:" in result:
-        result = result.split("Assistant:")[-1]
+    # Extract assistant part
+    if "<|assistant|>" in result:
+        result = result.split("<|assistant|>")[-1]
 
     return clean_text(result)
 
@@ -82,18 +80,16 @@ for msg in st.session_state.messages:
         st.markdown(msg["content"])
 
 # ==============================
-# INPUT BOX
+# INPUT
 # ==============================
 user_input = st.chat_input("Type your message...")
 
 if user_input:
-    # User message
     st.session_state.messages.append({"role": "user", "content": user_input})
 
     with st.chat_message("user"):
         st.markdown(user_input)
 
-    # AI response
     with st.spinner("🤖 Thinking..."):
         response = generate_response(user_input)