Update app.py

app.py

@@ -1,27 +1,33 @@
-import streamlit as st
 import os
-import torch
+import random
 from openai import OpenAI
-import numpy as np
-import sys
+import streamlit as st
 from dotenv import load_dotenv
-import random
-from huggingface_hub import InferenceClient
+from huggingface_hub import get_token
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain.indexes import VectorstoreIndexCreator
+from langchain_community.document_loaders.hugging_face_dataset import HuggingFaceDatasetLoader
+from langchain_huggingface.embeddings.huggingface_endpoint import HuggingFaceEndpointEmbeddings
+from langchain.chains import RetrievalQA
+from langchain_community.vectorstores import FAISS
 
 # Load environment variables
 load_dotenv()
 
+api_key=os.environ.get('API_KEY')
+
+get_token()
+
 # Constants
 MAX_TOKENS = 4000
 DEFAULT_TEMPERATURE = 0.5
 
-# initialize the client
-
+# Initialize the OpenAI client
 client = OpenAI(
-  base_url="https://api-inference.huggingface.co/v1",
-  api_key=os.environ.get('API_KEY')  # Replace with your token
+    base_url="https://api-inference.huggingface.co/v1",
+    api_key=api_key
 )
-    
+
 # Create supported models
 model_links = {
     "Meta-Llama-3.1-8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
@@ -30,40 +36,50 @@ model_links = {
     "Falcon-7b-Instruct": "tiiuae/falcon-7b-instruct",
 }
 
+# Load documents and set up RAG pipeline
+@st.cache_resource
+def setup_rag_pipeline():
+    loader = HuggingFaceDatasetLoader(
+        path='Atreyu4EVR/General-BYUI-Data',
+        page_content_column='content'
+    )
+    documents = loader.load()
 
-# Random dog images for error message
-random_dog_images = ["broken_llama3.jpeg"]
+    hf_embeddings = HuggingFaceEndpointEmbeddings(
+        model="sentence-transformers/all-MiniLM-L12-v2",
+        task="feature-extraction",
+        huggingfacehub_api_token=api_key
+    )
+
+    vector_store = FAISS.from_documents(documents, hf_embeddings)
+    retriever = vector_store.as_retriever()
+
+    return retriever
 
 def reset_conversation():
-    '''
-    Resets Conversation
-    '''
     st.session_state.conversation = []
     st.session_state.messages = []
-    return None
 
-st.sidebar.button('Reset Chat', on_click=reset_conversation) #Reset button
-    
 def main():
-    st.header('Multi-Models')
-
+    st.header('Multi-Models with RAG')
 
     # Sidebar for model selection and temperature
     selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
     temperature = st.sidebar.slider('Select a temperature value', 0.0, 1.0, DEFAULT_TEMPERATURE)
 
+    st.sidebar.button('Reset Chat', on_click=reset_conversation)
+
     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.write(f"Changed to {selected_model}")
         st.session_state.prev_option = selected_model
         reset_conversation()
 
     st.markdown(f'_powered_ by ***:violet[{selected_model}]***')
 
-    # Display model info and logo
+    # Display model info
     st.sidebar.write(f"You're now chatting with **{selected_model}**")
     st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
@@ -76,15 +92,28 @@ def main():
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
+    # Set up RAG pipeline
+    retriever = setup_rag_pipeline()
+
     # Chat input and response
     if prompt := st.chat_input("Type message here..."):
-        process_user_input(client, prompt, selected_model, temperature)
+        process_user_input(client, prompt, selected_model, temperature, retriever)
 
-def process_user_input(client, prompt, selected_model, temperature):
+def process_user_input(client, prompt, selected_model, temperature, retriever):
     # Display user message
     with st.chat_message("user"):
         st.markdown(prompt)
-    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Retrieve relevant documents
+    relevant_docs = retriever.get_relevant_documents(prompt)
+    context = "\n".join([doc.page_content for doc in relevant_docs])
+    
+    # Prepare messages with context
+    messages = [
+        {"role": "system", "content": f"You are an AI assistant. Use the following context to answer the user's question: {context}"},
+        {"role": "user", "content": prompt}
+    ]
+    st.session_state.messages.extend(messages)
 
     # Generate and display assistant response
     with st.chat_message("assistant"):
@@ -110,7 +139,7 @@ def handle_error(error):
     response = """😵‍💫 Looks like someone unplugged something!
     \n Either the model space is being updated or something is down."""
     st.write(response)
-    random_dog_pick = random.choice(random_dog_images)
+    random_dog_pick = random.choice(["broken_llama3.jpeg"])
     st.image(random_dog_pick)
     st.write("This was the error message:")
     st.write(str(error))