Update chatbot.py

chatbot.py

@@ -4,13 +4,14 @@ import torch
 from huggingface_hub import login
 import os
 
-# Force authentication with your HF token (secret in Space settings)
-login(token=os.getenv("HF_TOKEN"))
+# Force authentication
+if "HF_TOKEN" in os.environ:
+    login(token=os.getenv("HF_TOKEN"))
 
 # ================= CACHE THE MODEL =================
 @st.cache_resource
 def load_model():
-    model_id = "ammoncoder123/IPTchatbotModel-1.7B"  # Your correct repo
+    model_id = "ammoncoder123/IPTchatbotModel-1.7B" 
 
     quantization_config = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -18,6 +19,10 @@ def load_model():
     )
 
     tokenizer = AutoTokenizer.from_pretrained(model_id)
+    # Ensure tokenizer has a chat template, or use a default one
+    if tokenizer.chat_template is None:
+        tokenizer.chat_template = "{% for message in messages %}{{'<|' + message['role'] + '|>' + '\n' + message['content'] + '</s>\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% endif %}"
+
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         quantization_config=quantization_config,
@@ -29,81 +34,67 @@ def load_model():
     return pipeline(
         "text-generation",
         model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=300,
-        temperature=0.7,
-        do_sample=True,
-        top_p=0.9
+        tokenizer=tokenizer
     )
 
 pipe = load_model()
 
 # ==================== CHAT INTERFACE ====================
 st.title("IPT Chatbot Assistance")
-st.info("Answers may vary — please verify important facts.")
+st.info("Ask about logbooks, placement in Tanzania, or report writing.")
 
-# Display chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+# Display chat history
 for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+    if message["role"] != "system":  # Don't show system prompt to user
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
 
 # User input
 if prompt := st.chat_input("Ask about Industrial Practical Training..."):
-    # Add user message to history and display
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    # Generate assistant response
     with st.chat_message("assistant"):
         with st.spinner("Thinking..."):
-            # Strong system prompt (customized for your IPT dataset)
-            system_prompt = """
-            You are a helpful assistant for engineering and ICT students in Tanzania who are preparing for or doing Industrial Practical Training (IPT), also known as Industrial Attachment.
-            IPT means Industrial Practical Training — a mandatory work placement where students gain real-world experience in companies related to their field of study.
-            Always answer questions about:
-            - What IPT is
-            - How to do IPT (logbook, daily/weekly reports, technical report, presentation)
-            - Placement suggestions for different engineering fields (ICT, Mechatronics, Electrical, Mechanical, Civil, Biomedical, etc.)
-            - Choosing IPT centers/companies
-            - Tips for success in IPT
-            - Any other directly related IPT topic
-            If the question is clearly unrelated to IPT (e.g., politics, sports, personal life), politely reply:
-            "Sorry, I can only help with questions about Industrial Practical Training (IPT). Please ask something related to IPT, logbook, placement, or reports."
-            For placement suggestions (e.g., for Mechatronics, Electrical, ICT), give practical, realistic company types or industries in Tanzania that match the field.
-            Be concise, accurate, and helpful.
-            """
-
-            # Build messages: system prompt first, then user prompt
-            chat_messages = [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": prompt}
-            ]
+            # The context for the AI
+            system_message = {
+                "role": "system", 
+                "content": "You are a helpful assistant for engineering and ICT students in Tanzania. "
+                           "Answer questions about IPT, logbooks, and placements (e.g., TANESCO, TTCL, Halotel, TARURA). "
+                           "If the question is unrelated to IPT, politely decline. Be concise."
+            }
+
+            # Prepare the template-ready messages
+            # We only send the system prompt and the current prompt to keep it simple, 
+            # or you can send st.session_state.messages for full history.
+            input_messages = [system_message] + st.session_state.messages[-3:] # Last 3 messages for context
+
+            # Apply the chat template
+            formatted_prompt = pipe.tokenizer.apply_chat_template(
+                input_messages, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
 
             # Generate response
             outputs = pipe(
-                chat_messages,
+                formatted_prompt,
                 max_new_tokens=300,
                 temperature=0.7,
                 do_sample=True,
-                top_p=0.9
+                top_p=0.9,
+                return_full_text=False # THIS IS THE KEY FIX
             )
 
-            # Extract and clean the generated text
-            response = outputs[0]["generated_text"]
-            if isinstance(response, str) and response.startswith(prompt):
-                response = response[len(prompt):].strip()
-
-            # Show the response
+            response = outputs[0]["generated_text"].strip()
             st.markdown(response)
 
-    # Save assistant response to history
     st.session_state.messages.append({"role": "assistant", "content": response})
 
-# Optional: Clear conversation button
 if st.button("Clear Conversation"):
     st.session_state.messages = []
     st.rerun()