Update app.py

app.py

@@ -1,26 +1,52 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
+# ------------------------
+# Streamlit UI
+# ------------------------
 st.set_page_config(page_title="LangChain Demo", page_icon=":robot:")
 st.header("MHRV Chatbot")
 
+# ------------------------
+# Session memory
+# ------------------------
 if "sessionMessages" not in st.session_state:
     st.session_state.sessionMessages = [
-        SystemMessage(content="You are a helpful customer support chatbot for a website.")
+        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."
+            )
+        )
     ]
 
-# Load Flan-T5-Small (CPU-friendly)
-model_name = "google/flan-t5-small"
+# ------------------------
+# Load model and tokenizer
+# ------------------------
+model_name = "bigscience/bloom-560m"  # CPU-compatible
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=-1, max_new_tokens=256)
+# Create text-generation pipeline
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=-1,  # CPU
+    max_new_tokens=256,
+    temperature=0.3
+)
 
+# ------------------------
+# Helper functions
+# ------------------------
 def load_answer(question):
     st.session_state.sessionMessages.append(HumanMessage(content=question))
 
-    # Concatenate messages into a single prompt
+    # Build prompt from session messages
     prompt = ""
     for msg in st.session_state.sessionMessages:
         if isinstance(msg, SystemMessage):
@@ -30,9 +56,9 @@ def load_answer(question):
         elif isinstance(msg, AIMessage):
             prompt += f"AI: {msg.content}\n"
 
-    # Generate response
-    output = generator(prompt)
-    answer_text = output[0]["generated_text"].strip()
+    # Generate answer
+    output = generator(prompt, max_new_tokens=256, do_sample=True, temperature=0.3)
+    answer_text = output[0]["generated_text"][len(prompt):].strip()
 
     st.session_state.sessionMessages.append(AIMessage(content=answer_text))
     return answer_text
@@ -40,6 +66,9 @@ def load_answer(question):
 def get_text():
     return st.text_input("You: ", key="input")
 
+# ------------------------
+# Main app
+# ------------------------
 user_input = get_text()
 submit = st.button("Generate")