Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 # ------------------------
 # Streamlit UI
@@ -11,26 +10,17 @@ st.header("MHRV Chatbot")
 # ------------------------
 # Session memory
 # ------------------------
-if "sessionMessages" not in st.session_state:
-    st.session_state.sessionMessages = [
-        SystemMessage(
-            content=(
-                "You are a highly intelligent and helpful customer support assistant. "
-                "Answer user questions clearly, politely, and professionally. "
-                "If you don’t know the answer, say so instead of making things up. "
-                "Provide step-by-step instructions if relevant and helpful."
-            )
-        )
-    ]
+if "conversation" not in st.session_state:
+    st.session_state.conversation = []
 
 # ------------------------
 # Load model and tokenizer
 # ------------------------
-model_name = "bigscience/bloom-560m"  # CPU-compatible
+model_name = "bigscience/bloom-560m"  # CPU-friendly
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Create text-generation pipeline
+# Text-generation pipeline
 generator = pipeline(
     "text-generation",
     model=model,
@@ -43,36 +33,48 @@ generator = pipeline(
 # ------------------------
 # Helper functions
 # ------------------------
-def load_answer(question):
-    st.session_state.sessionMessages.append(HumanMessage(content=question))
+def load_answer(user_input):
+    # Add user input to session
+    st.session_state.conversation.append({"role": "user", "content": user_input})
 
-    # Build prompt from session messages
-    prompt = ""
-    for msg in st.session_state.sessionMessages:
-        if isinstance(msg, SystemMessage):
-            prompt += f"System: {msg.content}\n"
-        elif isinstance(msg, HumanMessage):
-            prompt += f"Human: {msg.content}\n"
-        elif isinstance(msg, AIMessage):
-            prompt += f"AI: {msg.content}\n"
+    # Build prompt for BLOOM
+    system_instruction = (
+        "You are a helpful, professional customer support assistant. "
+        "Answer questions clearly, politely, and accurately. "
+        "If the question is mathematical or factual, provide the correct answer. "
+        "Do not repeat the user's message."
+    )
 
-    # Generate answer
-    output = generator(prompt, max_new_tokens=256, do_sample=True, temperature=0.3)
-    answer_text = output[0]["generated_text"][len(prompt):].strip()
+    # Use last 3 messages (or fewer if new) to maintain context
+    prompt = system_instruction + "\n"
+    for msg in st.session_state.conversation[-3:]:
+        if msg["role"] == "user":
+            prompt += f"User: {msg['content']}\n"
+        elif msg["role"] == "assistant":
+            prompt += f"Assistant: {msg['content']}\n"
+    prompt += "Assistant:"
 
-    st.session_state.sessionMessages.append(AIMessage(content=answer_text))
-    return answer_text
+    # Generate answer
+    output = generator(prompt, max_new_tokens=128, do_sample=False)
+    answer = output[0]["generated_text"][len(prompt):].strip()
 
-def get_text():
-    return st.text_input("You: ", key="input")
+    # Save answer in session
+    st.session_state.conversation.append({"role": "assistant", "content": answer})
+    return answer
 
 # ------------------------
-# Main app
+# Streamlit input
 # ------------------------
-user_input = get_text()
+user_input = st.text_input("You: ", key="input")
 submit = st.button("Generate")
 
 if submit and user_input:
     response = load_answer(user_input)
     st.subheader("Answer:")
     st.write(response)
+
+# Optional: show conversation history
+if st.checkbox("Show conversation history"):
+    for msg in st.session_state.conversation:
+        role = "You" if msg["role"] == "user" else "Bot"
+        st.write(f"**{role}:** {msg['content']}")