Update app.py

app.py

@@ -14,9 +14,10 @@ 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"
 # Environment variables
 os.environ['LANGCHAIN_TRACING_V2'] = 'true'
 os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
@@ -65,7 +66,7 @@ def retrieve_normal_context(retriever, question):
 # Your OLMOLLM class implementation here (adapted for the Hugging Face model)
 
 @st.cache_resource
-def get_chain(temperature):
+def get_chain(temperature,selected_model):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L12-v2")
     
     docstore_path = 'ohw_proj_chorma_db.pcl'
@@ -77,20 +78,26 @@ def get_chain(temperature):
     child_splitter = RecursiveCharacterTextSplitter(chunk_size=300,
                                                     chunk_overlap=50)
     retriever = load_retriever(docstore_path,chroma_path,embeddings,child_splitter,parent_splitter)
-    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)
+    llm_api = 'glpat-AMzMevbqaVjp4HbLcVum'
+    llm = ChatOpenAI(model_name=selected_model,
+                    temperature=temperature,
+                    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()
@@ -147,8 +154,20 @@ def generate_response(chain, query, context):
 # Sidebar
 with st.sidebar:
     st.title("OHW Assistant")
+    selected_model = st.sidebar.selectbox('Choose a LLM model',
+                                           llm_list,
+                                             key='selected_model',
+                                             index = None)
+
     temperature = st.slider("Temperature: ", 0.0, 1.0, 0.5, 0.1)
-    chain = get_chain(temperature)
+    if selected_model in ['Mistral-7B-Instruct-v0.2', 'Mixtral-8x7B-Instruct-v0.1','LLAMA3']:
+        if selected_model == 'Mistral-7B-Instruct-v0.2':
+            selected_model = 'alias-fast'
+        elif selected_model == 'Mixtral-8x7B-Instruct-v0.1':
+            selected_model = 'alias-large'
+        elif selected_model == 'LLAMA3':
+            selected_model = 'alias-experimental'
+    chain = get_chain(temperature,selected_model)
     st.button('Clear Chat History', on_click=clear_chat_history)
 
 # Main app