Update app.py

app.py

@@ -12,7 +12,7 @@ from typing import List
 from numpy import ndarray
 from llama_cpp import Llama
 from statistical_chunker import StatisticalChunker
-from fastembed_encoder import FastEmbedEncoder
+from semantic_router.encoders.huggingface import HuggingFaceEncoder
 from scipy.sparse import csr_matrix, save_npz, load_npz, vstack
 from qdrant_client import QdrantClient, models
 from langchain_community.document_loaders import WikipediaLoader, WebBaseLoader
@@ -30,27 +30,22 @@ from pymilvus import (
    RRFRanker
 )
 
+def transform_query(query: str) -> str:
+    """ For retrieval, add the prompt for query (not for documents).
+    """
+    return f'Represent this sentence for searching relevant passages: {query}'
+
 def query_hybrid_search(col: Collection, query: str):
-    query_embeddings = ef([query])
+    query_dense_embeddings = dense_model(transform_query(query))
+    query_sparse_embeddings = sparse_model.encode(query)
 
-    sparse_req = AnnSearchRequest(
-        query_embeddings["sparse"], "sparse_vector", {"metric_type": "IP"}, limit=2
-    )
-    dense_req = AnnSearchRequest(
-        query_embeddings["dense"], "dense_vector", {"metric_type": "COSINE"}, limit=2
-    )
+    sparse_req = AnnSearchRequest(query_sparse_embeddings, "sparse_vector", {"metric_type": "IP"}, limit=10)
+    dense_req = AnnSearchRequest(query_dense_embeddings, "dense_vector", {"metric_type": "COSINE"}, limit=10)
 
-    res = col.hybrid_search(
-        [sparse_req, dense_req], rerank=RRFRanker(), limit=3, output_fields=["text"]
-    )
+    res = col.hybrid_search([sparse_req, dense_req], rerank=RRFRanker(), limit=3, output_fields=["text"])
 
     return res
 
-def transform_query(query: str) -> str:
-    """ For retrieval, add the prompt for query (not for documents).
-    """
-    return f'Represent this sentence for searching relevant passages: {query}'
-
 def main(query: str, client: QdrantClient, collection_name: str, llm, dense_model: FastEmbedEncoder, sparse_model: SparseTextEmbedding):
     dense_query = list(dense_model(query,32))
     sparse_query = list(sparse_model.embed(query, 32))
@@ -133,7 +128,7 @@ def load_models_and_documents():
             n_gpu_layers=32
         )
 
-        dense_model = HuggingfaceEncoder(
+        dense_model = HuggingFaceEncoder(
             name='mixedbread-ai/mxbai-embed-large-v1',
             device='cuda'
         )
@@ -150,7 +145,6 @@ def load_models_and_documents():
     collection_name = 'collection_demo'
 
     fields = [
-        FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
         FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=8192),
         FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
         FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=1024)
@@ -159,49 +153,10 @@ def load_models_and_documents():
     schema = CollectionSchema(fields, "")
     col = Collection(collection_name, schema)
 
-    sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
+    sparse_index = {"index_type": "SPARSE_WAND", "metric_type": "IP"}
     dense_index = {"index_type": "FLAT", "metric_type": "COSINE"}
     col.create_index("sparse_vector", sparse_index)
     col.create_index("dense_vector", dense_index)
-
-    entities = [
-        docs,
-        docs_embeddings["sparse"],
-        docs_embeddings["dense"]
-    ]
-    
-    client.create_collection(
-        collection_name,
-        {
-            "text-dense": models.VectorParams(
-                size=1024,
-                distance=models.Distance.COSINE,
-                on_disk=False,
-                quantization_config=models.BinaryQuantization(
-                    binary=models.BinaryQuantizationConfig(
-                        always_ram=True
-                    )
-                )
-            )
-        },
-        {
-            "text-sparse": models.SparseVectorParams(
-                index=models.SparseIndexParams(
-                    on_disk=False
-                )
-            )
-        },
-        2,
-        optimizers_config=models.OptimizersConfigDiff(
-            indexing_threshold=0,
-            default_segment_number=4
-        ),
-        hnsw_config=models.HnswConfigDiff(
-            on_disk=False, 
-            m=64,
-            ef_construct=512
-        )
-    )
         
     with st.spinner('Parse and chunk documents...'):
         name = 'action_rpg'