Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 st.set_page_config(page_title="LangChain Demo", page_icon=":robot:")
@@ -10,17 +10,17 @@ if "sessionMessages" not in st.session_state:
         SystemMessage(content="You are a helpful customer support chatbot for a website.")
     ]
 
-# Load tokenizer and model
-model_name = "tiiuae/Falcon3-1B-Instruct"
+# Load Flan-T5-Small (CPU-friendly)
+model_name = "google/flan-t5-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Create a text-generation pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1, max_new_tokens=256, temperature=0.3)
+generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=-1, max_new_tokens=256)
 
 def load_answer(question):
     st.session_state.sessionMessages.append(HumanMessage(content=question))
 
+    # Concatenate messages into a single prompt
     prompt = ""
     for msg in st.session_state.sessionMessages:
         if isinstance(msg, SystemMessage):
@@ -30,8 +30,9 @@ def load_answer(question):
         elif isinstance(msg, AIMessage):
             prompt += f"AI: {msg.content}\n"
 
-    output = generator(prompt, max_new_tokens=256, do_sample=True, temperature=0.3)
-    answer_text = output[0]["generated_text"][len(prompt):].strip()
+    # Generate response
+    output = generator(prompt)
+    answer_text = output[0]["generated_text"].strip()
 
     st.session_state.sessionMessages.append(AIMessage(content=answer_text))
     return answer_text