update model names

main.py

@@ -31,7 +31,7 @@ embeddings = HuggingFaceInferenceAPIEmbeddings(
 
 vector_store = SupabaseVectorStore(supabase, embeddings, query_name='match_documents', table_name="documents")
 
-models = ["llama-2"]
+models = ["meta-llama/Llama-2-7b-chat-hf", "mistralai/Mixtral-8x7B-Instruct-v0.1"]
 
 if openai_api_key:
     models += ["gpt-3.5-turbo", "gpt-4"]
@@ -77,7 +77,7 @@ if st.session_state["authenticated"]:
 
     # Initialize session state variables
     if 'model' not in st.session_state:
-        st.session_state['model'] = "llama-2"
+        st.session_state['model'] = "meta-llama/Llama-2-7b-chat-hf"
     if 'temperature' not in st.session_state:
         st.session_state['temperature'] = 0.1
     if 'chunk_size' not in st.session_state:

question.py

@@ -9,8 +9,8 @@ from langchain.chat_models import ChatAnthropic
 from langchain.vectorstores import SupabaseVectorStore
 from stats import add_usage
 
-memory = ConversationBufferMemory(
-    memory_key="chat_history", return_messages=True)
+# memory = ConversationBufferMemory(memory_key="chat_history", input_key='question', output_key='answer', return_messages=True)
+memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
 openai_api_key = st.secrets.openai_api_key
 anthropic_api_key = st.secrets.anthropic_api_key
 hf_api_key = st.secrets.hf_api_key
@@ -62,10 +62,10 @@ def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db):
                 qa = ConversationalRetrievalChain.from_llm(
                     ChatAnthropic(
                         model=st.session_state['model'], anthropic_api_key=anthropic_api_key, temperature=st.session_state['temperature'], max_tokens_to_sample=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True, max_tokens_limit=102400)
-            elif hf_api_key and model.startswith("llama"):
-                logger.info('Using Llama model %s', model)
+            elif hf_api_key:
+                logger.info('Using HF model %s', model)
                 # print(st.session_state['max_tokens'])
-                endpoint_url = ("https://api-inference.huggingface.co/models/meta-llama/Llama-2-70b-chat-hf")
+                endpoint_url = ("https://api-inference.huggingface.co/models/"+ model)
                 model_kwargs = {"temperature" : st.session_state['temperature'],
                                 "max_new_tokens" : st.session_state['max_tokens'],
                                 "return_full_text" : False}
@@ -75,7 +75,7 @@ def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db):
                     huggingfacehub_api_token=hf_api_key,
                     model_kwargs=model_kwargs
                 )
-                qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(), memory=memory, verbose=True)
+                qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(), memory=memory, verbose=True, return_source_documents=True)
             
             st.session_state['chat_history'].append(("You", question))
 
@@ -84,6 +84,7 @@ def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db):
             logger.info('Result: %s', model_response)
 
             st.session_state['chat_history'].append(("meraKB", model_response["answer"]))
+            # logger.info('Sources: %s', model_response["source_documents"][0])
 
             # Display chat history
             st.empty()