Update app.py

app.py

@@ -8,14 +8,11 @@ from langchain_community.llms import HuggingFacePipeline
 from langchain.retrievers import ParentDocumentRetriever
 from langchain.storage import InMemoryStore
 from langchain_chroma import Chroma
-from langchain.llms import LlamaCpp
 from langchain_openai import ChatOpenAI
 from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnableLambda
 from datetime import date
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import threading
 import time
 llm_list = ['Mistral-7B-Instruct-v0.2','Mixtral-8x7B-Instruct-v0.1','LLAMA3']
 blablador_base = "https://helmholtz-blablador.fz-juelich.de:8000/v1"
@@ -24,13 +21,6 @@ os.environ['LANGCHAIN_TRACING_V2'] = 'true'
 os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
 os.environ['LANGCHAIN_API_KEY'] = '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'
 
-
-@st.cache_resource
-def load_model():
-    model_name = "EleutherAI/gpt-neo-125M"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    return model, tokenizer
 def load_from_pickle(filename):
     with open(filename, "rb") as file:
         return pickle.load(file)
@@ -85,20 +75,7 @@ def get_chain(temperature,selected_model):
                     openai_api_key=llm_api,
                     openai_api_base=blablador_base,
                     streaming=True)
-    # model, tokenizer = load_model()
 
-    # pipe = pipeline(
-    #     "text-generation",
-    #     model=model, 
-    #     tokenizer=tokenizer, 
-    #     max_length=1800,
-    #     max_new_tokens = 200,
-    #     temperature=temperature,
-    #     top_p=0.95,
-    #     repetition_penalty=1.15
-    # )
-
-    # llm = HuggingFacePipeline(pipeline=pipe)
 
     
     today = date.today()
@@ -134,24 +111,7 @@ def clear_chat_history():
     st.session_state.messages = []
     st.session_state.context_sources = []
     st.session_state.key = 0
-def run_with_timeout(func, args, timeout):
-    result = [None]
-    def worker():
-        result[0] = func(*args)
-    thread = threading.Thread(target=worker)
-    thread.start()
-    thread.join(timeout)
-    if thread.is_alive():
-        return None
-    return result[0]
-# In your Streamlit app
-def generate_response(chain, query, context):
-    timeout_seconds = 180
-    result = chain.invoke, ({"question": query, "chat_history": st.session_state.messages},)
-    if result is None:
-        return result
-        # return "I apologize, but I couldn't generate a response in time. The query might be too complex for me to process quickly. Could you try simplifying your question?"
-    return result
+
 # Sidebar
 with st.sidebar:
     st.title("OHW Assistant")
@@ -206,17 +166,18 @@ if prompt := st.chat_input("How may I assist you today?"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        query = st.session_state.messages[-1]['content']
+        query=st.session_state.messages[-1]['content']
         tab1, tab2 = st.tabs(["Answer", "Sources"])
         with tab1:
-            with st.spinner("Generating answer..."):
-                start_time = time.time()
-                full_answer = chain.invoke({"question": query, "chat_history":st.session_state.messages})# Context is handled within the chain
-                end_time = time.time()
-
-            st.markdown(full_answer,unsafe_allow_html=True)
+            start_time = time.time()
+            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)
+            end_time = time.time()
             st.caption(f"Response time: {end_time - start_time:.2f} seconds")
-
         with tab2:
             if st.session_state.context_sources:
                 for i, source in enumerate(st.session_state.context_sources):
@@ -226,6 +187,7 @@ if prompt := st.chat_input("How may I assist you today?"):
             else:
                 st.write("No sources available for this query.")
 
+
     st.session_state.messages.append({"role": "assistant", "content": full_answer})
     st.session_state.messages[-1]['sources'] = st.session_state.context_sources
     st.session_state.messages[-1]['context'] = st.session_state.context_content