Update app.py

app.py

@@ -2,26 +2,37 @@ import streamlit as st
 from transformers import 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")
 
+# Initialize session messages
 if "sessionMessages" not in st.session_state:
     st.session_state.sessionMessages = [
-        SystemMessage(content="You are a customer support chatbot for a website.")
+        SystemMessage(content="You are a helpful customer support chatbot for a website.")
     ]
 
-# Hugging Face model using transformers pipeline
+# ------------------------
+# Load Hugging Face model (CPU-friendly)
+# ------------------------
+# Using a smaller model for free Spaces CPU plan
 generator = pipeline(
     "text-generation",
-    model="OpenAssistant/oasst-sft-4-pythia-12b",
-    device=0,  # use GPU if available
-    max_new_tokens=512,
-    temperature=0
+    model="mosaicml/mpt-7b-instruct",  # smaller, instruction-tuned, CPU-compatible
+    device=-1,  # -1 = CPU
+    max_new_tokens=256,
+    temperature=0.3
 )
 
+# ------------------------
+# Helper functions
+# ------------------------
 def load_answer(question):
     st.session_state.sessionMessages.append(HumanMessage(content=question))
 
+    # Convert session messages to a single prompt string
     prompt = ""
     for msg in st.session_state.sessionMessages:
         if isinstance(msg, SystemMessage):
@@ -31,15 +42,21 @@ def load_answer(question):
         elif isinstance(msg, AIMessage):
             prompt += f"AI: {msg.content}\n"
 
-    output = generator(prompt, max_new_tokens=512, do_sample=True, temperature=0)
+    # Generate response
+    output = generator(prompt, max_new_tokens=256, do_sample=True, temperature=0.3)
     answer_text = output[0]["generated_text"][len(prompt):].strip()
 
+    # Add AI message to session
     st.session_state.sessionMessages.append(AIMessage(content=answer_text))
+
     return answer_text
 
 def get_text():
     return st.text_input("You: ", key="input")
 
+# ------------------------
+# Main App
+# ------------------------
 user_input = get_text()
 submit = st.button("Generate")