Update app.py

app.py

@@ -1,41 +1,47 @@
 import streamlit as st
-from transformers import pipeline
 from langchain_community.llms import HuggingFacePipeline
+from langchain.schema import HumanMessage, SystemMessage
+from transformers import pipeline
 
-# Title
-st.title("🤖 Hugging Face Chatbot with LangChain")
+# Cache the model so Streamlit doesn't reload every time
+@st.cache_resource
+def load_model():
+    return pipeline(
+        "text-generation",
+        model="tiiuae/falcon-1b-instruct",
+        torch_dtype="float32",     # keep it stable
+        device_map="auto",         # use GPU if available
+        trust_remote_code=True
+    )
+
+# Load pipeline once
+pipe = load_model()
+chat = HuggingFacePipeline(pipeline=pipe)
 
-# Session state for conversation history
+# --- Session state for messages ---
 if "sessionMessages" not in st.session_state:
-    st.session_state.sessionMessages = []
-
-# Build HF pipeline in float32
-pipe = pipeline(
-    "text-generation",
-    model="tiiuae/falcon-7b-instruct",  # or any Hugging Face model
-    torch_dtype="float32",              # force safe float32 precision
-    device_map="auto"                   # auto-detect GPU/CPU
-)
+    st.session_state.sessionMessages = [
+        SystemMessage(content="You are a helpful assistant.")
+    ]
 
-# Wrap into LangChain
-chat = HuggingFacePipeline(pipeline=pipe)
-
-# Function to handle answers
-def load_answer(question):
-    st.session_state.sessionMessages.append({"role": "user", "content": question})
+def load_answer(user_input):
+    st.session_state.sessionMessages.append(HumanMessage(content=user_input))
     assistant_answer = chat.invoke(st.session_state.sessionMessages)
-    st.session_state.sessionMessages.append({"role": "assistant", "content": assistant_answer})
+    st.session_state.sessionMessages.append(SystemMessage(content=assistant_answer))
     return assistant_answer
 
-# Streamlit input
-user_input = st.text_input("Ask me anything:")
+# --- Streamlit UI ---
+st.title("Falcon-1B Chatbot 🦅")
+
+user_input = st.text_input("You:", key="input")
 
 if user_input:
-    response = load_answer(user_input)
+    with st.spinner("Thinking..."):
+        response = load_answer(user_input)
     st.write(f"**Assistant:** {response}")
 
-# Display chat history
-if st.session_state.sessionMessages:
-    st.subheader("Chat History")
-    for msg in st.session_state.sessionMessages:
-        st.write(f"**{msg['role'].capitalize()}:** {msg['content']}")
+if st.button("Clear Chat"):
+    st.session_state.sessionMessages = [
+        SystemMessage(content="You are a helpful assistant.")
+    ]
+    st.experimental_rerun()