Update app.py

app.py

@@ -1,22 +1,20 @@
 import os
 import re
-import sys
-import copy
 import time
+import msgpack
 import numpy as np
 import streamlit as st
-from typing import Optional
-from stqdm import stqdm 
 from numpy import ndarray
-from typing import Iterable
+from scipy.sparse import csr_matrix, save_npz, load_npz, vstack
 from qdrant_client import QdrantClient, models
 from fastembed.sparse.splade_pp import supported_splade_models
 from fastembed import SparseTextEmbedding, SparseEmbedding
-from langchain_community.llms.exllamav2 import ExLlamaV2
-from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
 from fastembed_ext import FastEmbedEmbeddingsLc
+from langchain_community.chat_models.ollama import ChatOllama
 from langchain_community.document_loaders.wikipedia import WikipediaLoader
 from langchain_community.document_loaders.unstructured import UnstructuredFileLoader
+from langchain_core.prompts import PromptTemplate
+from langchain.chains.summarize import load_summarize_chain
 from langchain_experimental.text_splitter import SemanticChunker
 from langchain_core.documents import Document
 from qdrant_client.models import (
@@ -27,10 +25,6 @@ from qdrant_client.models import (
     SearchRequest,
     ScoredPoint,
 )
-from langchain_core.prompts import PromptTemplate
-from langchain.chains.summarize import load_summarize_chain
-from huggingface_hub import snapshot_download
-from exllamav2.generator import ExLlamaV2Sampler
 
 MAP_PROMPT = """
 You will be given a single passage of a book. This section will be enclosed in triple backticks (```)
@@ -50,21 +44,11 @@ The reader should be able to grasp what happened in the book.
 VERBOSE SUMMARY:
 """
 
-supported_splade_models[0] = {
-        "model": "prithivida/Splade_PP_en_v2",
-        "vocab_size": 30522,
-        "description": "Implementation of SPLADE++ Model for English v2",
-        "size_in_GB": 0.532,
-        "sources": {
-            "hf": "devve1/Splade_PP_en_v2_onnx"
-        },
-        "model_file": "model.onnx"
-    }
 
-def make_points(chunks: list[str], dense: list[ndarray], sparse)-> Iterable[PointStruct]:
+def make_points(chunks: list[str], dense: list[ndarray], sparse: list[SparseEmbedding])-> list[PointStruct]:
     points = []
-    for idx, (sparse_vec, chunk, dense_vector) in enumerate(zip(sparse, chunks, dense)):
-        sparse_vector = SparseVector(indices=sparse_vec.indices.tolist(), values=sparse_vec.values.tolist())
+    for idx, (sparse_vector, chunk, dense_vector) in enumerate(zip(sparse, chunks, dense)):
+        sparse_vector = SparseVector(indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist())
         point = PointStruct(
             id=idx,
             vector={
@@ -131,7 +115,7 @@ def rrf(rank_lists, alpha=60, default_rank=1000):
     return sorted_items
 
 
-def main(query: str, client: QdrantClient, collection_name: str, llm, dense_model, sparse_model):
+def main(query: str, client: QdrantClient, collection_name: str, llm, dense_model: FastEmbedEmbeddingsLc, sparse_model: SparseTextEmbedding):
      # name = 'Kia_EV6'
     # filepath = os.path.join(os.getcwd(), name + '.pdf')
 
@@ -145,6 +129,7 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
     # )
 
     # docs = docs.load()
+
     
     dense_query = list(dense_model.embed_query(query, 32))
     sparse_query = list(sparse_model.embed(query, 32))
@@ -166,8 +151,6 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
     
     docs = [Document(record.payload['text']) for record in records_list[:3]]
 
-    print(docs)
-
     map_prompt = PromptTemplate(
         template=MAP_PROMPT,
         input_variables=['text']
@@ -198,39 +181,44 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
     output = reduce_chain.invoke([summaries])
     return output['output_text']
 
-@st.cache_resource
-def load_models_and_components(show_spinner="Loading models..."):
-    settings = ExLlamaV2Sampler.Settings()
-    settings.temperature = 0.75
-    settings.top_k = 50
-    settings.top_p = 0.8
-    settings.token_repetition_penalty = 1.05
-    
-    model_path = snapshot_download(repo_id='Zoyd/NousResearch_Hermes-2-Theta-Llama-3-8B-6_5bpw_exl2')
-    callbacks = [StreamingStdOutCallbackHandler()]
-
-    llm = ExLlamaV2(
-        model_path=model_path,
-        callbacks=callbacks,
-        settings=settings,
-        streaming=True,
-        max_new_tokens=3000
-    )
-        
-    provider = ['CPUExecutionProvider']
-
-    sparse_model = SparseTextEmbedding(
-        'Qdrant/bm42-all-minilm-l6-v2-attentions',
-        cache_dir=os.getenv('HF_HOME'),
-        providers=provider
-    )
+def load_models_and_documents():
+    supported_splade_models[0] = {
+        "model": "prithivida/Splade_PP_en_v2",
+        "vocab_size": 30522,
+        "description": "Implementation of SPLADE++ Model for English v2",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "devve1/Splade_PP_en_v2_onnx"
+        },
+        "model_file": "model.onnx"
+    }
 
-    dense_model = FastEmbedEmbeddingsLc(
-        model_name='mixedbread-ai/mxbai-embed-large-v1',
-        providers=provider,
-        cache_dir=os.getenv('HF_HOME'),
-        batch_size=32
-    )
+    with st.spinner('Load models...'):
+        settings = ExLlamaV2Sampler.Settings()
+        settings.temperature = 0.75
+        
+        model_path = snapshot_download(repo_id='Zoyd/NousResearch_Hermes-2-Theta-Llama-3-8B-6_5bpw_exl2')
+        
+        llm = ExLlamaV2(
+            model_path=model_path,
+            settings=settings,
+            max_new_tokens=3000
+        )
+        
+        provider = ['CPUExecutionProvider']
+        
+        dense_model = FastEmbedEmbeddingsLc(
+            model_name='mixedbread-ai/mxbai-embed-large-v1',
+            providers=provider,
+            cache_dir=os.getenv('HF_HOME'),
+            batch_size=32
+        )
+        
+        sparse_model = SparseTextEmbedding(
+            'Qdrant/bm42-all-minilm-l6-v2-attentions',
+            cache_dir=os.getenv('HF_HOME'),
+            providers=provider
+        )
 
     client = QdrantClient(path=os.getenv('HF_HOME'))
     collection_name = 'collection_demo'
@@ -262,7 +250,7 @@ def load_models_and_components(show_spinner="Loading models..."):
             on_disk_payload=True,
             optimizers_config=models.OptimizersConfigDiff(
                 memmap_threshold=10000,
-                indexing_treshold=0
+                indexing_threshold=0
             ),
             hnsw_config=models.HnswConfigDiff(
                 on_disk=True, 
@@ -270,18 +258,65 @@ def load_models_and_components(show_spinner="Loading models..."):
                 ef_construct=100
             )
         )
-
-    docs = WikipediaLoader(query='Action-RPG').load() 
-    chunks, dense, sparse = chunk_documents(docs, dense_model, sparse_model)
-
-    client.upsert(
-        collection_name,
-        make_points(
-            chunks, 
-            dense, 
-            sparse
+        
+    with st.spinner('Parse and chunk documents...'):
+        name = 'action_rpg'
+        embeddings_path = os.path.join(os.getenv('HF_HOME'), 'collection', 'embeddings')
+
+        chunks_path = os.path.join(embeddings_path, name + '_chunks.msgpack')
+        dense_path = os.path.join(embeddings_path, name + '_dense.npz')
+        sparse_path = os.path.join(embeddings_path, name + '_sparse.npz')
+
+        if not os.path.exists(embeddings_path):
+            os.mkdir(embeddings_path)
+
+            docs = WikipediaLoader(query='Action-RPG').load()
+            chunks, dense_embeddings, sparse_embeddings = chunk_documents(docs, dense_model, sparse_model)
+
+            with open(chunks_path, "wb") as outfile:
+                packed = msgpack.packb(chunks, use_bin_type=True)
+                outfile.write(packed)
+            
+            np.savez_compressed(dense_path, *dense_embeddings)
+            max_index = max(np.max(embedding.indices) for embedding in sparse_embeddings)
+        
+            sparse_matrices = []
+            for embedding in sparse_embeddings:
+                data = embedding.values
+                indices = embedding.indices
+                indptr = np.array([0, len(data)])
+                matrix = csr_matrix((data, indices, indptr), shape=(1, max_index + 1))
+                sparse_matrices.append(matrix)
+            
+            combined_sparse_matrix = vstack(sparse_matrices)
+            save_npz(sparse_path, combined_sparse_matrix)
+        else:
+            with open(chunks_path, "rb") as data_file:
+                byte_data = data_file.read()
+                
+            chunks = msgpack.unpackb(byte_data, raw=False)
+
+            dense_embeddings = list(np.load(dense_path).values())
+
+            sparse_embeddings = []
+            loaded_sparse_matrix = load_npz(sparse_path)
+        
+            for i in range(loaded_sparse_matrix.shape[0]):
+                row = loaded_sparse_matrix.getrow(i)
+                values = row.data
+                indices = row.indices
+                embedding = SparseEmbedding(values, indices)
+                sparse_embeddings.append(embedding)
+
+    with st.spinner('Save documents...'):
+        client.upsert(
+            collection_name,
+            make_points(
+                chunks, 
+                dense_embeddings, 
+                sparse_embeddings
+            )
         )
-    )
     client.update_collection(
         collection_name=collection_name,
         optimizer_config=models.OptimizersConfigDiff(indexing_threshold=20000)
@@ -289,7 +324,7 @@ def load_models_and_components(show_spinner="Loading models..."):
 
     return client, collection_name, llm, dense_model, sparse_model
 
-def chunk_documents(docs, dense_model, sparse_model, show_spinner="Parsing and chunking texts..."):
+def chunk_documents(docs, dense_model, sparse_model):
     text_splitter = SemanticChunker(
         dense_model,
         breakpoint_threshold_type='standard_deviation'
@@ -297,20 +332,20 @@ def chunk_documents(docs, dense_model, sparse_model, show_spinner="Parsing and c
 
     documents = [doc.page_content for doc in text_splitter.transform_documents(list(docs))]
 
-    dense_embeddings = dense_model.embed_documents(stqdm(documents,desc='Generate dense embeddings...', backend=True), 32)
-    sparse_embeddings = list(sparse_model.embed(stqdm(documents, desc='Generate sparse embeddings...', backend=True), 32))
+    dense_embeddings = dense_model.embed_documents(documents,32)
+    sparse_embeddings = list(sparse_model.embed(documents, 32))
     
     return documents, dense_embeddings, sparse_embeddings
     
 if __name__ == '__main__':
     st.set_page_config(page_title="Video Game Assistant",
                        layout="wide"
-                      )
+                       )
+    st.title("Video Game Assistant :sunglasses:")
+    
     if 'models_loaded' not in st.session_state:
-        st.session_state.client, st.session_state.collection_name, st.session_state.llm, st.session_state.dense_model, st.session_state.sparse_model = load_models_and_components()
+        st.session_state.client, st.session_state.collection_name, st.session_state.llm, st.session_state.dense_model, st.session_state.sparse_model = load_models_and_documents()
         st.session_state.models_loaded = True
-
-    st.title("Video Game Assistant")
     
     if "messages" not in st.session_state:
         st.session_state.messages = []
@@ -323,13 +358,7 @@ if __name__ == '__main__':
         st.chat_message("user").markdown(prompt)
         st.session_state.messages.append({"role": "user", "content": prompt})
 
-        client = st.session_state.client
-        collection_name = st.session_state.collection_name
-        llm = st.session_state.llm
-        dense_model = st.session_state.dense_model
-        sparse_model = st.session_state.sparse_model
-
-        ai_response = main(prompt, client, collection_name, llm, dense_model, sparse_model)
+        ai_response = main(prompt, st.session_state.client, st.session_state.collection_name, st.session_state.llm, st.session_state.dense_model, st.session_state.sparse_model)
         response = f"Echo: {ai_response}"
         with st.chat_message("assistant"):
             message_placeholder = st.empty()
@@ -338,5 +367,4 @@ if __name__ == '__main__':
                 full_response += chunk + " "
                 time.sleep(0.01)
                 message_placeholder.markdown(full_response + "▌")
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
-    
+        st.session_state.messages.append({"role": "assistant", "content": full_response})