Update src/streamlit_app.py

src/streamlit_app.py

@@ -14,7 +14,7 @@ st.title("🤖 Simple AI Assistant")
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # 🔥 BEST FOR HF FREE
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # ✅ Best for HF free
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
@@ -41,32 +41,44 @@ if "messages" not in st.session_state:
 # CLEAN TEXT
 # ==============================
 def clean_text(text):
-    text = re.sub(r"[^\x00-\x7F]+", "", text)
-    return text.strip()
+    text = re.sub(r"[^\x00-\x7F]+", "", text).strip()
+
+    # Ensure response completes nicely
+    if not text.endswith((".", "!", "?")):
+        text += "..."
+
+    return text
 
 # ==============================
 # GENERATE RESPONSE
 # ==============================
 def generate_response(user_input):
 
-    prompt = f"<|user|>\n{user_input}\n<|assistant|>\n"
+    prompt = f"""
+<|user|>
+{user_input}
+
+Give a clear and complete answer.
+<|assistant|>
+"""
 
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=150,
+            max_new_tokens=300,   # 🔥 prevents cut-off
             do_sample=True,
             temperature=0.7,
             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)
 
-    # Extract assistant part
+    # Extract assistant response
     if "<|assistant|>" in result:
         result = result.split("<|assistant|>")[-1]
 
@@ -80,19 +92,22 @@ for msg in st.session_state.messages:
         st.markdown(msg["content"])
 
 # ==============================
-# INPUT
+# INPUT BOX
 # ==============================
 user_input = st.chat_input("Type your message...")
 
 if user_input:
+    # Add 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..."):
         response = generate_response(user_input)
 
+    # Add assistant response
     st.session_state.messages.append({"role": "assistant", "content": response})
 
     with st.chat_message("assistant"):