Update app.py

app.py

@@ -1,23 +1,23 @@
 import os
 import re
 import time
+import torch
 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 fastembed.sparse.splade_pp import supported_splade_models
-from fastembed import SparseTextEmbedding, SparseEmbedding
-from fastembed_ext import FastEmbedEmbeddingsLc
 from langchain_community.llms.llamacpp import LlamaCpp
 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_community.embeddings import HuggingFaceEmbeddings
 from langchain_core.documents import Document
-from huggingface_hub import snapshot_download
+from huggingface_hub import hf_hub_download
 from qdrant_client.models import (
     NamedSparseVector,
     NamedVector,
@@ -63,7 +63,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: ndarray, sparse: list[SparseEmbedding]):
+def search(client: QdrantClient, collection_name: str, dense, indices, values):
     search_results = client.search_batch(
         collection_name,
         [
@@ -78,8 +78,8 @@ def search(client: QdrantClient, collection_name: str, dense: ndarray, sparse: l
                 vector=NamedSparseVector(
                     name="text-sparse",
                     vector=SparseVector(
-                        indices=sparse[0].indices.tolist(),
-                        values=sparse[0].values.tolist(),
+                        indices=indices.tolist(),
+                        values=values.tolist(),
                     ),
                 ),
                 limit=10
@@ -132,7 +132,7 @@ def main(query: str, client: QdrantClient, collection_name: str, llm, dense_mode
     # docs = docs.load()
 
     
-    dense_query = compute_dense(query, dense_model)
+    dense_query = compute_dense_query(query, dense_model)
     sparse_query = compute_sparse(query, sparse_model, sparse_tokenizer)
 
     search_results = search(
@@ -197,28 +197,22 @@ def compute_sparse(sentence, model, tokenizer):
     max_val, _ = torch.max(weighted_log, dim=1)
     vector = max_val.squeeze()
     
-    cols = vector.nonzero().numpy().flatten()
-    weights = vector.detach().numpy()[cols]
+    cols = vector.nonzero().squeeze().tolist()
+    weights = vector[cols].tolist()
 
     return cols, weights
 
-def compute_dense(sentence, model):
-    
+def compute_dense_query(sentence, model):
+    return model.embed_query(f'Represent this sentence for searching relevant passages: {sentence}')
 
-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"
-    }
+def compute_dense_docs(docs, model):
+    return model.embed_documents(docs)
 
+def load_models_and_documents():
     with st.spinner('Load models...'):
-        model_path = snapshot_download(repo_id='Zoyd/NousResearch_Hermes-2-Theta-Llama-3-8B-6_5bpw_exl2')
+        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'
+                                     )
         
         llm = LlamaCpp(
             model_path=model_path,
@@ -233,6 +227,11 @@ def load_models_and_documents():
         reverse_voc = {v: k for k, v in 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}
+                                           )
+
     client = QdrantClient(path=os.getenv('HF_HOME'))
     collection_name = 'collection_demo'
     
@@ -344,7 +343,7 @@ 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(documents, dense_model)
+    dense_embeddings = compute_dense_docs(documents, dense_model)
     indices, values = compute_sparse(documents, sparse_model, sparse_tokenizer)
     
     return documents, dense_embeddings, indices, values