Update app.py

app.py

@@ -3,7 +3,6 @@ import re
 import lz4
 import time
 import uuid
-import vllm
 import torch
 import spacy
 import base64
@@ -20,10 +19,10 @@ from typing import List, Dict
 from ppt_chunker import ppt_chunk
 from outlines import models, generate
 from qdrant_client import QdrantClient
-from optimum_encoder import OptimumEncoder
 from unstructured.cleaners.core import clean
 from streamlit_navigation_bar import st_navbar
 from vllm.sampling_params import SamplingParams
+from vllm import LLM, PoolingParams, PoolingType
 from fastembed import SparseTextEmbedding, SparseEmbedding
 from unstructured.nlp.tokenize import download_nltk_packages
 from huggingface_hub import snapshot_download, hf_hub_download
@@ -89,7 +88,7 @@ def transform_query(query: str) -> str:
     """
     return f'Represent this sentence for searching relevant passages: {query}'
 
-def query_hybrid_search(query: str, client: QdrantClient, collection_name: str, dense_model: OptimumEncoder, sparse_model: SparseTextEmbedding):
+def query_hybrid_search(query: str, client: QdrantClient, collection_name: str, dense_model: LLM, sparse_model: SparseTextEmbedding):
     dense_embeddings = dense_model.embed_query(transform_query(query))[0]
     sparse_embeddings = list(sparse_model.query_embed(query))[0]
 
@@ -103,7 +102,7 @@ def query_hybrid_search(query: str, client: QdrantClient, collection_name: str,
         with_vectors=False,
         with_payload=True,
         limit=10,
-        score_threshold=0.9
+        score_threshold=0.95
     )
 
 def build_prompt_conv():
@@ -304,11 +303,15 @@ def load_models_and_documents():
     container = st.empty()
     
     with container.status("Load AI Models and Prepare Documents...", expanded=True) as status:
-        st.write('Downloading and Loading MixedBread Mxbai Dense Embedding Model under ONNX with Nvidia CUDA as backend...')
+        st.write('Downloading and Loading MixedBread Mxbai Dense Embedding Model with vLLM as backend...')
 
-        dense_model = OptimumEncoder(
-            device="cuda",
-            cache_dir=os.getenv('HF_HOME')
+        dense_model = LLM(
+            model='mixedbread-ai/mxbai-embed-large-v1',
+            enforce_eager=True,
+            max_model_len=512,
+            max_num_seqs=32,
+            tensor_parallel_size=1,
+            dtype=torch.float16
         )
 
         st.write('Downloading and Loading Qdrant BM42 Sparse Embedding Model under ONNX using the CPU...')
@@ -319,15 +322,10 @@ def load_models_and_documents():
             providers=['CPUExecutionProvider']
         )
 
-        st.write('Downloading Mistral Nemo AI Model...')
-        
-        model_path = snapshot_download('casperhansen/mistral-nemo-instruct-2407-awq')
-
-        st.write('Loading Mistral Nemo AI Model quantized with AWQ and using Outlines + vLLM Engine as backend...')
+        st.write('Downloading and Loading Mistral Nemo AI Model quantized with AWQ and using Outlines + vLLM Engine as backend...')
 
-        llm = vllm.LLM(
-            model=model_path,
-            tokenizer=model_path,
+        llm = LLM(
+            model='casperhansen/mistral-nemo-instruct-2407-awq',
             tensor_parallel_size=1,
             trust_remote_code=True,
             enforce_eager=True,
@@ -516,7 +514,8 @@ def chunk_documents(texts: List[str], metadatas: List[dict], dense_model: Optimu
         documents.append(doc.page_content)
 
     start_dense = time.time()
-    dense_embeddings = dense_model.embed_documents(documents)
+    dense_embeddings = dense_model.encode(documents, pooling_params=PoolingParams(pooling_type=PoolingType.MEAN))
+    print(f'DENSE EMBED : {dense_embeddings}')
     end_dense = time.time()
     final_dense = end_dense - start_dense
     print(f'DENSE TIME: {final_dense}')
@@ -529,7 +528,7 @@ def chunk_documents(texts: List[str], metadatas: List[dict], dense_model: Optimu
     final_sparse = end_sparse - start_sparse
     print(f'SPARSE TIME: {final_sparse}')
     
-    return payload_docs, dense_embeddings, sparse_embeddings
+    return payload_docs, dense_embeddings[0].outputs.embedding, sparse_embeddings
 
 def on_change_documents_only():
     if st.session_state.documents_only: