Update app.py

app.py

@@ -6,7 +6,6 @@ import msgpack
 import numpy as np
 import streamlit as st
 from numpy import ndarray
-from transformers import AutoModelForMaskedLM, AutoTokenizer
 from scipy.sparse import csr_matrix, save_npz, load_npz, vstack
 from qdrant_client import QdrantClient, models
 from langchain_community.llms.llamacpp import LlamaCpp
@@ -15,7 +14,9 @@ from langchain_community.document_loaders.unstructured import UnstructuredFileLo
 from langchain_core.prompts import PromptTemplate
 from langchain.chains.summarize import load_summarize_chain
 from langchain_experimental.text_splitter import SemanticChunker
-from langchain_huggingface import HuggingFaceEmbeddings
+from fastembed.sparse.splade_pp import supported_splade_models
+from fastembed import SparseTextEmbedding, SparseEmbedding
+from fastembed_ext import FastEmbedEmbeddingsLc
 from langchain_core.documents import Document
 from huggingface_hub import hf_hub_download
 from qdrant_client.models import (
@@ -46,10 +47,10 @@ VERBOSE SUMMARY:
 """
 
 
-def make_points(chunks: list[str], dense: list[ndarray], indices, values)-> list[PointStruct]:
+def make_points(chunks: list[str], dense: list[ndarray], sparse)-> list[PointStruct]:
     points = []
-    for idx, (indice, value, chunk, dense_vector) in enumerate(zip(indices, values, chunks, dense)):
-        sparse_vector = SparseVector(indices=indice, values=value)
+    for idx, (sparse_vector, chunk, dense_vector) in enumerate(zip(sparse, chunks, dense)):
+        sparse_vec = SparseVector(indices=sparse_vector.indices.tolist(), values=sparse_vector.values.tolist())
         point = PointStruct(
             id=idx,
             vector={
@@ -63,7 +64,7 @@ def make_points(chunks: list[str], dense: list[ndarray], indices, values)-> list
         points.append(point)
     return points
 
-def search(client: QdrantClient, collection_name: str, dense, indices, values):
+def search(client: QdrantClient, collection_name: str, dense, sparse):
     search_results = client.search_batch(
         collection_name,
         [
@@ -78,8 +79,8 @@ def search(client: QdrantClient, collection_name: str, dense, indices, values):
                 vector=NamedSparseVector(
                     name="text-sparse",
                     vector=SparseVector(
-                        indices=indices,
-                        values=values,
+                        indices=sparse[0].indices.tolist(),
+                        values=sparse[0].values.tolist(),
                     ),
                 ),
                 limit=10
@@ -116,7 +117,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, sparse_tokenizer):
+def main(query: str, client: QdrantClient, collection_name: str, llm, dense_model, sparse_model):
      # name = 'Kia_EV6'
     # filepath = os.path.join(os.getcwd(), name + '.pdf')
 
@@ -130,17 +131,15 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
     # )
 
     # docs = docs.load()
-
     
-    dense_query = compute_dense_query(query, dense_model)
-    sparse_query_indices, sparse_query_values = compute_sparse(query, sparse_model, sparse_tokenizer)
+    dense_query = list(dense_model.embed_query(query, 32))
+    sparse_query = list(sparse_model.embed(query, 32))
 
     search_results = search(
         client,
         collection_name,
         dense_query,
-        sparse_query_indices,
-        sparse_query_values
+        sparse_query
     )
 
     dense_rank_list, sparse_rank_list = rank_list(search_results[0]), rank_list(search_results[1])
@@ -183,33 +182,18 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
     output = reduce_chain.invoke([summaries])
     return output['output_text']
 
-def compute_sparse(sentence, model, tokenizer):
-    inputs = tokenizer(sentence, return_tensors='pt')
-    inputs = {key: val.to(device) for key, val in inputs.items()}
-    input_ids = inputs['input_ids']
-    
-    attention_mask = inputs['attention_mask']
-    
-    outputs = model(**inputs)
-    
-    logits, attention_mask = outputs.logits, attention_mask
-    relu_log = torch.log(1 + torch.relu(logits))
-    weighted_log = relu_log * attention_mask.unsqueeze(-1)
-    max_val, _ = torch.max(weighted_log, dim=1)
-    vector = max_val.squeeze()
-    
-    cols = vector.nonzero().squeeze().tolist()
-    weights = vector[cols].tolist()
-
-    return cols, weights
-
-def compute_dense_query(sentence, model):
-    return model.embed_query(f'Represent this sentence for searching relevant passages: {sentence}')
-
-def compute_dense_docs(docs, model):
-    return model.embed_documents(docs)
-
 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"
+    }
+    
     with st.spinner('Load models...'):
         model_path = hf_hub_download(repo_id='NousResearch/Hermes-2-Theta-Llama-3-8B-GGUF',
                                      filename='Hermes-2-Pro-Llama-3-Instruct-Merged-DPO-Q8_0.gguf'
@@ -223,15 +207,21 @@ def load_models_and_documents():
             n_batch=512,
             f16_kv=True
         )
-        
-        sparse_tokenizer = AutoTokenizer.from_pretrained('prithivida/Splade_PP_en_v2')
-        reverse_voc = {v: k for k, v in sparse_tokenizer.vocab.items()}
-        sparse_model = AutoModelForMaskedLM.from_pretrained('prithivida/Splade_PP_en_v2')
 
-        dense_model = HuggingFaceEmbeddings(model_name='mixedbread-ai/mxbai-embed-large-v1', 
-                                            cache_folder=os.getenv('HF_HOME'), 
-                                            model_kwargs={'truncate_dim':512}
-                                           )
+        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(
+            'prithivida/Splade_PP_en_v2',
+            cache_dir=os.getenv('HF_HOME'),
+            providers=provider
+        )
 
     client = QdrantClient(path=os.getenv('HF_HOME'))
     collection_name = 'collection_demo'
@@ -283,21 +273,21 @@ def load_models_and_documents():
             os.mkdir(embeddings_path)
 
             docs = WikipediaLoader(query='Action-RPG').load()
-            chunks, dense_embeddings, indices, values = chunk_documents(docs, dense_model, sparse_model, sparse_tokenizer)
+            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(indice) for indice in indices)
+            max_index = max(np.max(embedding.indices) for embedding in sparse_embeddings)
         
             sparse_matrices = []
-            for indice, value in zip(indices, values):
-                data = value
-                indices = indice
+            for embedding in sparse_embeddings:
+                data = embedding.values
+                indices = embedding.indices
                 indptr = np.array([0, len(data)])
-                matrix = csr_matrix((data, indice, indptr), shape=(1, max_index + 1))
+                matrix = csr_matrix((data, indices, indptr), shape=(1, max_index + 1))
                 sparse_matrices.append(matrix)
             
             combined_sparse_matrix = vstack(sparse_matrices)
@@ -310,14 +300,15 @@ def load_models_and_documents():
 
             dense_embeddings = list(np.load(dense_path).values())
 
-            indices = []
-            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.append(row.data)
-                indices.append(row.indices)
+                values = row.data
+                indices = row.indices
+                embedding = SparseEmbedding(values, indices)
+                sparse_embeddings.append(embedding)
 
     with st.spinner('Save documents...'):
         client.upsert(
@@ -325,8 +316,7 @@ def load_models_and_documents():
             make_points(
                 chunks, 
                 dense_embeddings, 
-                indices,
-                values,
+                sparse_embeddings
             )
         )
     client.update_collection(
@@ -334,9 +324,9 @@ def load_models_and_documents():
         optimizer_config=models.OptimizersConfigDiff(indexing_threshold=20000)
     )
 
-    return client, collection_name, llm, dense_model, sparse_model, sparse_tokenizer
+    return client, collection_name, llm, dense_model, sparse_model
 
-def chunk_documents(docs, dense_model, sparse_model, sparse_tokenizer):
+def chunk_documents(docs, dense_model, sparse_model):
     text_splitter = SemanticChunker(
         dense_model,
         breakpoint_threshold_type='standard_deviation'
@@ -344,10 +334,10 @@ def chunk_documents(docs, dense_model, sparse_model, sparse_tokenizer):
 
     documents = [doc.page_content for doc in text_splitter.transform_documents(list(docs))]
 
-    dense_embeddings = compute_dense_docs(documents, dense_model)
-    indices, values = compute_sparse(documents, sparse_model, sparse_tokenizer)
+    dense_embeddings = dense_model.embed_documents(documents,32)
+    sparse_embeddings = list(sparse_model.embed(documents, 32))
     
-    return documents, dense_embeddings, indices, values
+    return documents, dense_embeddings, sparse_embeddings
     
 if __name__ == '__main__':
     st.set_page_config(page_title="Video Game Assistant",
@@ -356,7 +346,7 @@ if __name__ == '__main__':
     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, st.session_state.sparse_tokenizer = load_models_and_documents()
+        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
 
     if st.session.state.models_loaded:
@@ -371,7 +361,7 @@ if __name__ == '__main__':
             st.chat_message("user").markdown(prompt)
             st.session_state.messages.append({"role": "user", "content": prompt})
 
-            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, st.session_state.sparse_tokenizer)
+            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()