Update app.py

app.py

@@ -1,23 +1,22 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 st.set_page_config(page_title="LangChain Demo", page_icon=":robot:")
-st.header("Pequeno Chatbot - Matheus")
+st.header("MHRV Chatbot")
 
 if "sessionMessages" not in st.session_state:
     st.session_state.sessionMessages = [
         SystemMessage(content="You are a helpful customer support chatbot for a website.")
     ]
 
-# Load the Falcon3-1B-Instruct model
-generator = pipeline(
-    "text-generation",
-    model="tiiuae/Falcon3-1B-Instruct",
-    device=-1,  # CPU
-    max_new_tokens=256,
-    temperature=0.3
-)
+# Load tokenizer and model
+model_name = "tiiuae/Falcon3-1B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
+
+# Create a text-generation pipeline
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1, max_new_tokens=256, temperature=0.3)
 
 def load_answer(question):
     st.session_state.sessionMessages.append(HumanMessage(content=question))