Update app.py

app.py

@@ -14,23 +14,6 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableLambda
 from datetime import date
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-# from setup import download_olmo_model, OLMO_MODEL
-
-# Ensure model is downloaded before proceeding
-# @st.cache_resource
-# def load_model():
-#     model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-#     tokenizer = AutoTokenizer.from_pretrained(model_name)
-#     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
-#     return model, tokenizer
-
-# model, tokenizer = load_model()
-
-# # Define the path to your bash script
-# script_path = "./start.sh"
-
-# # Run the bash script
-# subprocess.run([script_path], shell=True, check=True)
 
 # Environment variables
 os.environ['LANGCHAIN_TRACING_V2'] = 'true'
@@ -38,15 +21,7 @@ os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
 os.environ['LANGCHAIN_API_KEY'] = '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'
 
 
-# Load model and tokenizer
-# @st.cache_resource
-# def load_model():
-#     model_name = "allenai/OLMo-7B-Instruct"
-#     tokenizer = AutoTokenizer.from_pretrained(model_name)
-#     model = AutoModelForCausalLM.from_pretrained(model_name)
-#     return model, tokenizer
 
-# model, tokenizer = load_model()
 def load_from_pickle(filename):
     with open(filename, "rb") as file:
         return pickle.load(file)
@@ -185,13 +160,12 @@ if prompt := st.chat_input("How may I assist you today?"):
         query=st.session_state.messages[-1]['content']
         tab1, tab2 = st.tabs(["Answer", "Sources"])
         with tab1:
-            placeholder = st.empty()  # Create a placeholder in Streamlit
-            full_answer = ""
-            for chunk in chain.stream({"question": query, "chat_history":st.session_state.messages}):
-                
-                full_answer += chunk
-                placeholder.markdown(full_answer,unsafe_allow_html=True)
+            with st.spinner("Generating answer..."):
+                # Generate the full answer at once
+                full_answer = chain.invoke({"question": query, "chat_history": st.session_state.messages})
             
+            # Display the full answer
+            st.markdown(full_answer, unsafe_allow_html=True)
         with tab2:
             for i, source in enumerate(st.session_state.context_sources):
                 name = f'{source}'