feat: upgrade query adapter algorithm

src/raglite/_bench.py

@@ -2,7 +2,7 @@
 
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Generator
+from collections.abc import Iterator
 from dataclasses import replace
 from functools import cached_property
 from pathlib import Path
@@ -60,13 +60,15 @@ def trec_run_filename(self) -> str:
     def trec_run_filepath(self) -> Path:
         return self.cwd / self.trec_run_filename
 
-    def score(self) -> Generator[ScoredDoc, None, None]:
+    def score(self) -> Iterator[ScoredDoc]:
         """Read or compute a TREC run."""
         if self.trec_run_filepath.exists():
             yield from read_trec_run(self.trec_run_filepath.as_posix())  # type: ignore[no-untyped-call]
             return
         if not self.search("q0", next(self.dataset.queries_iter()).text):
             self.insert_documents()
+        if hasattr(self, "prescore"):
+            self.prescore()
         with self.trec_run_filepath.open(mode="w") as trec_run_file:
             for query in tqdm(
                 self.dataset.queries_iter(),
@@ -113,23 +115,37 @@ def insert_documents(self, max_workers: int | None = None) -> None:
         ]
         insert_documents(documents, max_workers=max_workers, config=self.config)
 
-    def update_query_adapter(self, num_evals: int = 1024) -> None:
+    def prescore(self) -> None:
+        from sqlalchemy import func, select
+        from sqlmodel import Session
+
         from raglite import insert_evals, update_query_adapter
-        from raglite._database import IndexMetadata
+        from raglite._database import Eval, IndexMetadata, create_database_engine
+
+        if not self.config.vector_search_query_adapter:
+            return
 
-        if (
-            self.config.vector_search_query_adapter
-            and IndexMetadata.get(config=self.config).get("query_adapter") is None
-        ):
-            insert_evals(num_evals=num_evals, config=self.config)
+        required_evals = 1024
+        with Session(create_database_engine(self.config)) as session:
+            num_evals = session.execute(select(func.count()).select_from(Eval)).scalar_one()
+        if num_evals < required_evals:
+            insert_evals(num_evals=required_evals - num_evals, config=self.config)
+        if IndexMetadata.get(config=self.config).get("query_adapter") is None:
             update_query_adapter(config=self.config)
 
     def search(self, query_id: str, query: str, *, num_results: int = 10) -> list[ScoredDoc]:
-        from raglite import retrieve_chunks, vector_search
+        from raglite import retrieve_chunks, search_and_rerank_chunks, vector_search
 
-        self.update_query_adapter()
-        chunk_ids, scores = vector_search(query, num_results=2 * num_results, config=self.config)
-        chunks = retrieve_chunks(chunk_ids, config=self.config)
+        if self.config.reranker:
+            chunks = search_and_rerank_chunks(
+                query=query, num_results=2 * num_results, config=self.config
+            )
+            scores = [1 / rank for rank in range(1, len(chunks) + 1)]
+        else:
+            chunk_ids, scores = vector_search(
+                query, num_results=2 * num_results, config=self.config
+            )
+            chunks = retrieve_chunks(chunk_ids, config=self.config)
         scored_docs = [
             ScoredDoc(query_id=query_id, doc_id=chunk.document.id, score=score)
             for chunk, score in zip(chunks, scores, strict=True)

src/raglite/_cli.py

@@ -132,25 +132,31 @@ def bench(
     ),
 ) -> None:
     """Run benchmark."""
-    import ir_datasets
-    import ir_measures
-    import pandas as pd
-
-    from raglite._bench import (
-        IREvaluator,
-        LlamaIndexEvaluator,
-        OpenAIVectorStoreEvaluator,
-        RAGLiteEvaluator,
-    )
+    try:
+        import ir_datasets
+        import ir_measures
+        import pandas as pd
+        from rerankers import Reranker
+
+        from raglite._bench import (
+            IREvaluator,
+            LlamaIndexEvaluator,
+            OpenAIVectorStoreEvaluator,
+            RAGLiteEvaluator,
+        )
+    except ModuleNotFoundError as import_error:
+        error_message = "To use the `bench` command, please install the `bench` extra."
+        raise ModuleNotFoundError(error_message) from import_error
 
     # Initialise the benchmark.
     evaluator: IREvaluator
     measures = [ir_measures.parse_measure(measure)]
     index, results = [], []
-    # Evaluate RAGLite (single-vector) + DuckDB HNSW + text-embedding-3-large.
+    # Evaluate RAGLite (single-vector) + DuckDB + text-embedding-3-large.
     chunk_max_size = 2048
     config = RAGLiteConfig(
-        embedder="text-embedding-3-large",
+        embedder=(embedder := "text-embedding-3-large"),
+        reranker=None,
         chunk_max_size=chunk_max_size,
         vector_search_multivector=False,
         vector_search_query_adapter=False,
@@ -161,9 +167,10 @@ def bench(
     )
     index.append("RAGLite (single-vector)")
     results.append(ir_measures.calc_aggregate(measures, dataset.qrels_iter(), evaluator.score()))
-    # Evaluate RAGLite (multi-vector) + DuckDB HNSW + text-embedding-3-large.
+    # Evaluate RAGLite (multi-vector) + DuckDB + text-embedding-3-large.
     config = RAGLiteConfig(
-        embedder="text-embedding-3-large",
+        embedder=embedder,
+        reranker=None,
         chunk_max_size=chunk_max_size,
         vector_search_multivector=True,
         vector_search_query_adapter=False,
@@ -174,10 +181,11 @@ def bench(
     )
     index.append("RAGLite (multi-vector)")
     results.append(ir_measures.calc_aggregate(measures, dataset.qrels_iter(), evaluator.score()))
-    # Evaluate RAGLite (query adapter) + DuckDB HNSW + text-embedding-3-large.
+    # Evaluate RAGLite (multi-vector; query adapter) + DuckDB + text-embedding-3-large.
     config = RAGLiteConfig(
         llm=(llm := "gpt-4.1"),
-        embedder="text-embedding-3-large",
+        embedder=embedder,
+        reranker=None,
         chunk_max_size=chunk_max_size,
         vector_search_multivector=True,
         vector_search_query_adapter=True,
@@ -191,6 +199,29 @@ def bench(
     )
     index.append("RAGLite (query adapter)")
     results.append(ir_measures.calc_aggregate(measures, dataset.qrels_iter(), evaluator.score()))
+    # Evaluate RAGLite (multi-vector; query adapter; reranker) + DuckDB + text-embedding-3-large.
+    if os.environ.get("CO_API_KEY"):
+        config = RAGLiteConfig(
+            llm=llm,
+            embedder=embedder,
+            reranker=Reranker(
+                "rerank-v3.5", model_type="cohere", api_key=os.environ["CO_API_KEY"], verbose=0
+            ),
+            chunk_max_size=chunk_max_size,
+            vector_search_multivector=True,
+            vector_search_query_adapter=True,
+        )
+        dataset = ir_datasets.load(dataset_name)
+        evaluator = RAGLiteEvaluator(
+            dataset,
+            insert_variant=f"multi-vector-{chunk_max_size // 4}t",
+            search_variant=f"query-adapter-{llm}-cohere-rerank-3.5",
+            config=config,
+        )
+        index.append("RAGLite (Cohere Rerank 3.5)")
+        results.append(
+            ir_measures.calc_aggregate(measures, dataset.qrels_iter(), evaluator.score())
+        )
     # Evaluate LLamaIndex + FAISS HNSW + text-embedding-3-large.
     dataset = ir_datasets.load(dataset_name)
     evaluator = LlamaIndexEvaluator(dataset)

src/raglite/_insert.py

@@ -170,6 +170,7 @@ def insert_documents(  # noqa: C901
                 session.flush()  # Flush changes to the database.
                 session.expunge_all()  # Release memory of flushed changes.
                 num_unflushed_embeddings = 0
+            pbar.set_postfix({"id": document_record.id})
             pbar.update()
         session.commit()
         if engine.dialect.name == "duckdb":