Restore

app.py

@@ -1,164 +1,85 @@
 import streamlit as st
 from openai import OpenAI
 import os
-import json
+import numpy as np
 from dotenv import load_dotenv
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.schema import Document
-from langchain_community.llms import HuggingFaceHub
-from langchain.chains import RetrievalQA
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import Chroma
-from tqdm import tqdm
 import random
 
 # Load environment variables
 load_dotenv()
 
 # Constants
-CHUNK_SIZE = 8192
-CHUNK_OVERLAP = 200
-BATCH_SIZE = 100
-RETRIEVER_K = 4
-VECTORSTORE_PATH = "./vectorstore"
-
-# Model information
+MAX_TOKENS = 4000
+DEFAULT_TEMPERATURE = 0.5
+
+# Initialize the client
+def initialize_client():
+    api_key = os.environ.get('HUGGINGFACEHUB_API_TOKEN')
+    if not api_key:
+        st.error("HUGGINGFACEHUB_API_TOKEN not found in environment variables.")
+        st.stop()
+    return OpenAI(
+        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",
     "Mistral-7B-Instruct-v0.3": "mistralai/Mistral-7B-Instruct-v0.3",
+    "Gemma-7b-it": "google/gemma-7b-it",
 }
 
+# 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.
-        \nIt was created by the [**Meta's AI**](https://llama.meta.com/) team and has over **8 billion parameters.**\n""",
-        "logo": "llama_logo.gif",
+        'description': """The Llama (3.1) model is a **Large Language Model (LLM)** that's able to have question and answer interactions.
+        \nIt was created by the [**Meta's AI**](https://llama.meta.com/) team and has over **8 billion parameters.**\n"""
     },
     "Mistral-7B-Instruct-v0.3": {
-        "description": """The Mistral-7B-Instruct-v0.3 Large Language Model (LLM) is an instruct fine-tuned version of the Mistral-7B-v0.3.
-        \nIt was created by the [**Mistral AI**](https://mistral.ai/news/announcing-mistral-7b/) team as has over **7 billion parameters.**\n""",
-        "logo": "https://mistral.ai/images/logo_hubc88c4ece131b91c7cb753f40e9e1cc5_2589_256x0_resize_q97_h2_lanczos_3.webp",
+        'description': """The Mistral-7B-Instruct-v0.3 is an instruct-tuned version of Mistral-7B.
+        \nIt was created by [**Mistral AI**](https://mistral.ai/) and has **7 billion parameters.**\n"""
     },
+    "Gemma-7b-it": {
+        'description': """Gemma is a family of lightweight, state-of-the-art open models from Google.
+        \nThe 7B-it variant is instruction-tuned and has **7 billion parameters.**\n"""
+    }
 }
 
 # Random dog images for error message
-random_dogs = ["randomdog.jpg", "randomdog2.jpg", "randomdog3.jpg"]  # Add more as needed
+random_dog_images = ["BlueLogoBox.jpg", "RandomDog1.jpg", "RandomDog2.jpg"]
 
-# Set up embeddings
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+def main():
+    st.header('Liahona.AI')
 
-def load_and_process_documents(file_path):
-    """Load and process documents from a JSON file."""
-    try:
-        with open(file_path, "r") as file:
-            data = json.load(file)
+    # 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)
 
-        documents = data.get("documents", [])
+    st.markdown(f'_powered_ by ***:violet[{selected_model}]***')
 
-        if not documents:
-            raise ValueError("No valid documents found in JSON file.")
+    # Display model info
+    st.sidebar.write(f"You're now chatting with **{selected_model}**")
+    st.sidebar.markdown(model_info[selected_model]['description'])
+    st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
-        doc_objects = [
-            Document(
-                page_content=doc["content"],
-                metadata={"title": doc["title"], "id": doc["id"]},
-            )
-            for doc in documents
-        ]
-
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
-        )
-        splits = text_splitter.split_documents(doc_objects)
-
-        return splits
-    except Exception as e:
-        st.error(f"Error loading documents: {str(e)}")
-        return []
-
-def get_vectorstore(file_path):
-    """Get or create a vectorstore."""
-    try:
-        if os.path.exists(VECTORSTORE_PATH):
-            print("Loading existing vectorstore...")
-            return Chroma(
-                persist_directory=VECTORSTORE_PATH, embedding_function=embeddings
-            )
+    # Initialize chat history
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+
+    # Display chat messages from history on app rerun
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
 
-        print("Creating new vectorstore...")
-        splits = load_and_process_documents(file_path)
-
-        vectorstore = None
-        for i in tqdm(range(0, len(splits), BATCH_SIZE), desc="Processing batches"):
-            batch = splits[i : i + BATCH_SIZE]
-            if vectorstore is None:
-                vectorstore = Chroma.from_documents(
-                    documents=batch,
-                    embedding=embeddings,
-                    persist_directory=VECTORSTORE_PATH,
-                )
-            else:
-                vectorstore.add_documents(documents=batch)
-
-        vectorstore.persist()
-        return vectorstore
-    except Exception as e:
-        st.error(f"Error creating vectorstore: {str(e)}")
-        return None
-
-@st.cache_resource(hash_funcs={"builtins.tuple": lambda _: None})
-def setup_rag_pipeline(file_path, model_name, temperature):
-    """Set up the RAG pipeline."""
-    try:
-        vectorstore = get_vectorstore(file_path)
-        if vectorstore is None:
-            raise ValueError("Failed to create or load vectorstore.")
-
-        llm = HuggingFaceHub(
-            repo_id=model_links[model_name],
-            model_kwargs={"temperature": temperature, "max_length": 4000},
-        )
-
-        return RetrievalQA.from_chain_type(
-            llm=llm,
-            chain_type="stuff",
-            retriever=vectorstore.as_retriever(search_kwargs={"k": RETRIEVER_K}),
-            return_source_documents=True,
-        )
-    except Exception as e:
-        st.error(f"Error setting up RAG pipeline: {str(e)}")
-        return None
-
-# 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
-temperature = st.sidebar.slider("Select a temperature value", 0.0, 1.0, 0.5)
-
-# 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.*")
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# Display chat messages from history
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Set up advanced RAG pipeline
-qa_chain = setup_rag_pipeline("index_training.json", selected_model, temperature)
-
-# Chat input
-if prompt := st.chat_input("Type message here..."):
+    # Initialize client
+    client = initialize_client()
+
+    # Chat input and response
+    if prompt := st.chat_input("Type message here..."):
+        process_user_input(client, prompt, selected_model, temperature)
+
+def process_user_input(client, prompt, selected_model, temperature):
     # Display user message
     with st.chat_message("user"):
         st.markdown(prompt)
@@ -167,21 +88,35 @@ if prompt := st.chat_input("Type message here..."):
     # Generate and display assistant response
     with st.chat_message("assistant"):
         try:
-            if qa_chain is None:
-                raise ValueError("RAG pipeline is not properly set up.")
-            
-            result = qa_chain({"query": prompt})
-            response = result["result"]
-            st.write(response)
-
+            stream = client.chat.completions.create(
+                model=model_links[selected_model],
+                messages=[
+                    {"role": m["role"], "content": m["content"]}
+                    for m in st.session_state.messages
+                ],
+                temperature=temperature,
+                stream=True,
+                max_tokens=MAX_TOKENS,
+            )
+            response = st.write_stream(stream)
         except Exception as e:
-            response = """😵‍💫 Looks like someone unplugged something!
-            \n Either the model space is being updated or something is down.
-            \n"""
-            st.write(response)
-            random_dog_pick = random.choice(random_dogs)
-            st.image(random_dog_pick)
-            st.write("This was the error message:")
-            st.write(str(e))
-
-    st.session_state.messages.append({"role": "assistant", "content": response})
+            handle_error(e)
+            return
+
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+def handle_error(error):
+    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 🐶:"""
+    st.write(response)
+    random_dog_pick = random.choice(random_dog_images)
+    st.image(random_dog_pick)
+    st.write("This was the error message:")
+    st.write(str(error))
+
+if __name__ == "__main__":
+    main()