Update app.py

app.py

@@ -1,17 +1,24 @@
 import streamlit as st
 from openai import OpenAI
-import torch
 import os
-import sys
-from dotenv import load_dotenv, dotenv_values
+from dotenv import load_dotenv
 import numpy as np
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.schema import Document
+from langchain_community.llms import HuggingFaceHub
+from langchain.chains import RetrievalQA, LLMChain
+from langchain.prompts import PromptTemplate
+from langchain_community.embeddings import HuggingFaceInstructEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import LLMChainExtractor
 
 load_dotenv()
 
-# Initialize the client
+# Initialize the OpenAI client for Hugging Face
 client = OpenAI(
     base_url="https://api-inference.huggingface.co/v1",
-    api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')  # Replace with your token
+    api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')
 )
 
 # Create supported models
@@ -21,7 +28,6 @@ model_links = {
     "gemma-2-2b": "google/gemma-2-2b",
 }
 
-# Pull info about the model to display
 model_info = {
     "Meta-Llama-3.1-8B": {
         'description': """The Llama (3.1) model is a **Large Language Model (LLM)** that's able to have question and answer interactions.
@@ -49,78 +55,117 @@ models = [key for key in model_links.keys()]
 # Create the sidebar with the dropdown for model selection
 selected_model = st.sidebar.selectbox("Select Model", models)
 
-# Create a temperature slider
+# Function to load and process documents
+def load_and_process_documents(file_path):
+    with open(file_path, 'r') as file:
+        content = file.read()
+    
+    doc = Document(page_content=content, metadata={"source": file_path})
+    
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=8192, chunk_overlap=200)
+    splits = text_splitter.split_documents([doc])
+    
+    return splits
+
+# Function to set up the advanced RAG pipeline
+@st.cache_resource
+def setup_advanced_rag_pipeline(model_name):
+    # Load and process documents
+    splits = load_and_process_documents("index_training.json")  # Replace with your document path
+    
+    # Set up InstructorEmbeddings
+    embeddings = HuggingFaceInstructEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    
+    # Create vectorstore
+    vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
+    
+    # Set up language model
+    llm = HuggingFaceHub(repo_id=model_links[model_name], model_kwargs={"temperature": 0.5, "max_length": 4000})
+    
+    # Set up HyDE
+    hyde_prompt = PromptTemplate(
+        input_variables=["question"],
+        template="Please write a passage to answer the question\nQuestion: {question}\nPassage:"
+    )
+    hyde_chain = LLMChain(llm=llm, prompt=hyde_prompt)
+    
+    def hyde_retriever(query):
+        hypothetical_doc = hyde_chain.run(query)
+        hyde_embedding = embeddings.embed_query(hypothetical_doc)
+        return vectorstore.similarity_search_by_vector(hyde_embedding, k=3)
+    
+    # Set up ContextualCompressionRetriever
+    compressor = LLMChainExtractor.from_llm(llm)
+    compression_retriever = ContextualCompressionRetriever(
+        base_compressor=compressor,
+        base_retriever=hyde_retriever
+    )
+    
+    # Create RetrievalQA chain
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=compression_retriever,
+        return_source_documents=True
+    )
+    
+    return qa_chain
+
+# Streamlit app
+st.header('Liahona.AI')
+
+# Sidebar for model selection
+selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
+st.markdown(f'_powered_ by ***:violet[{selected_model}]***')
+
+# Temperature slider
 temp_values = st.sidebar.slider('Select a temperature value', 0.0, 1.0, 0.5)
 
-# Create model description
+# Display model info
 st.sidebar.write(f"You're now chatting with **{selected_model}**")
 st.sidebar.markdown(model_info[selected_model]['description'])
 st.sidebar.image(model_info[selected_model]['logo'])
 st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
-if "prev_option" not in st.session_state:
-    st.session_state.prev_option = selected_model
-
-if st.session_state.prev_option != selected_model:
-    st.session_state.messages = []
-    st.session_state.prev_option = selected_model
-
-# Pull in the model we want to use
-repo_id = model_links[selected_model]
-
-st.header('Liahona.AI')
-st.markdown(f'_powered_ by ***:violet[{selected_model}]***')
-
-# Set a default model
-if selected_model not in st.session_state:
-    st.session_state[selected_model] = model_links[selected_model]
-
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display chat messages from history on app rerun
+# Display chat messages from history
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# Accept user input
-if prompt := st.chat_input("Type message here..."):
+# Set up advanced RAG pipeline
+qa_chain = setup_advanced_rag_pipeline(selected_model)
 
-    # Display user message in chat message container
+# Chat input
+if prompt := st.chat_input("Type message here..."):
+    # Display user message
     with st.chat_message("user"):
         st.markdown(prompt)
-    # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Display assistant response in chat message container
+    # Generate and display assistant response
     with st.chat_message("assistant"):
-
         try:
-            stream = client.chat.completions.create(
-                model=repo_id,
-                messages=[
-                    {"role": m["role"], "content": m["content"]}
-                    for m in st.session_state.messages
-                ],
-                temperature=temp_values,
-                stream=True,
-                max_tokens=4000,
-            )
-
-            response = st.write_stream(stream)
-
+            result = qa_chain({"query": prompt})
+            response = result["result"]
+            st.write(response)
+            
+            # Optionally, display source documents
+            st.expander("View Source Documents"):
+                for doc in result["source_documents"]:
+                    st.write(doc.page_content)
+                    st.write("---")
         except Exception as e:
             response = """😵‍💫 Looks like someone unplugged something!
             \n Either the model space is being updated or something is down.
-            \n
-            \n Try again later.
-            \n
-            \n Here's a random pic of a 🐶:"""
+            \n"""
             st.write(response)
             random_dog_pick = 'https://random.dog/' + random_dog[np.random.randint(len(random_dog))]
             st.image(random_dog_pick)
             st.write("This was the error message:")
-            st.write(e)
+            st.write(str(e))
 
     st.session_state.messages.append({"role": "assistant", "content": response})