Update app.py

app.py

@@ -2,9 +2,6 @@ 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")
 
@@ -13,25 +10,18 @@ if "sessionMessages" not in st.session_state:
         SystemMessage(content="You are a customer support chatbot for a website.")
     ]
 
-# ------------------------
-# Load Hugging Face pipeline
-# ------------------------
-# Using a free, instruction-tuned model
+# Hugging Face model using transformers pipeline
 generator = pipeline(
     "text-generation",
     model="OpenAssistant/oasst-sft-4-pythia-12b",
-    device=0,  # GPU if available
+    device=0,  # use GPU if available
     max_new_tokens=512,
     temperature=0
 )
 
-# ------------------------
-# Helper Functions
-# ------------------------
 def load_answer(question):
     st.session_state.sessionMessages.append(HumanMessage(content=question))
 
-    # Convert session messages to a single string prompt
     prompt = ""
     for msg in st.session_state.sessionMessages:
         if isinstance(msg, SystemMessage):
@@ -41,20 +31,15 @@ def load_answer(question):
         elif isinstance(msg, AIMessage):
             prompt += f"AI: {msg.content}\n"
 
-    # Generate response
     output = generator(prompt, max_new_tokens=512, do_sample=True, temperature=0)
-    answer_text = output[0]["generated_text"][len(prompt):].strip()  # remove prompt from output
+    answer_text = output[0]["generated_text"][len(prompt):].strip()
 
     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")