feat(cli): Stop hardcode the reranker of choice. Implements #68

Author: Zhe Yu

src/vectorcode/cli_utils.py

@@ -85,8 +85,12 @@ class Config:
     overlap_ratio: float = 0.2
     query_multiplier: int = -1
     query_exclude: list[PathLike] = field(default_factory=list)
-    reranker: Optional[str] = "cross-encoder/ms-marco-MiniLM-L-6-v2"
-    reranker_params: dict[str, Any] = field(default_factory=dict)
+    reranker: Optional[str] = "CrossEncoderReranker"
+    reranker_params: dict[str, Any] = field(
+        default_factory=lambda: {
+            "model_name_or_path": "cross-encoder/ms-marco-MiniLM-L-6-v2"
+        }
+    )
     check_item: Optional[str] = None
     use_absolute_path: bool = False
     include: list[QueryInclude] = field(
@@ -100,6 +104,7 @@ async def import_from(cls, config_dict: dict[str, Any]) -> "Config":
         """
         Raise IOError if db_path is not valid.
         """
+        default_config = Config()
         db_path = config_dict.get("db_path")
         host = config_dict.get("host") or "localhost"
         port = config_dict.get("port") or 8000
@@ -112,25 +117,35 @@ async def import_from(cls, config_dict: dict[str, Any]) -> "Config":
         return Config(
             **{
                 "embedding_function": config_dict.get(
-                    "embedding_function", "SentenceTransformerEmbeddingFunction"
+                    "embedding_function", default_config.embedding_function
+                ),
+                "embedding_params": config_dict.get(
+                    "embedding_params", default_config.embedding_params
                 ),
-                "embedding_params": config_dict.get("embedding_params", {}),
                 "host": host,
                 "port": port,
                 "db_path": db_path,
                 "db_log_path": os.path.expanduser(
-                    config_dict.get("db_log_path", "~/.local/share/vectorcode/")
+                    config_dict.get("db_log_path", default_config.db_log_path)
+                ),
+                "chunk_size": config_dict.get("chunk_size", default_config.chunk_size),
+                "overlap_ratio": config_dict.get(
+                    "overlap_ratio", default_config.overlap_ratio
+                ),
+                "query_multiplier": config_dict.get(
+                    "query_multiplier", default_config.query_multiplier
+                ),
+                "reranker": config_dict.get("reranker", default_config.reranker),
+                "reranker_params": config_dict.get(
+                    "reranker_params", default_config.reranker_params
+                ),
+                "db_settings": config_dict.get(
+                    "db_settings", default_config.db_settings
                 ),
-                "chunk_size": config_dict.get("chunk_size", 2500),
-                "overlap_ratio": config_dict.get("overlap_ratio", 0.2),
-                "query_multiplier": config_dict.get("query_multiplier", -1),
-                "reranker": config_dict.get(
-                    "reranker", "cross-encoder/ms-marco-MiniLM-L-6-v2"
+                "hnsw": config_dict.get("hnsw", default_config.hnsw),
+                "chunk_filters": config_dict.get(
+                    "chunk_filters", default_config.chunk_filters
                 ),
-                "reranker_params": config_dict.get("reranker_params", {}),
-                "db_settings": config_dict.get("db_settings", None),
-                "hnsw": config_dict.get("hnsw", {}),
-                "chunk_filters": config_dict.get("chunk_filters", {}),
             }
         )
 

src/vectorcode/subcommands/query/__init__.py

@@ -14,6 +14,7 @@
     get_collection,
     verify_ef,
 )
+from vectorcode.subcommands.query.reranker import get_reranker
 
 logger = logging.getLogger(name=__name__)
 
@@ -66,17 +67,8 @@ async def get_query_result_files(
         # no results found
         return []
 
-    if not configs.reranker:
-        from .reranker import NaiveReranker
-
-        aggregated_results = NaiveReranker(configs).rerank(results)
-    else:
-        from .reranker import CrossEncoderReranker
-
-        aggregated_results = CrossEncoderReranker(
-            configs, query_chunks, configs.reranker, **configs.reranker_params
-        ).rerank(results)
-    return aggregated_results
+    reranker = get_reranker(configs)
+    return reranker.rerank(results, query_chunks)
 
 
 async def build_query_results(

src/vectorcode/subcommands/query/reranker.py

@@ -0,0 +1,31 @@
+import logging
+
+from vectorcode.cli_utils import Config
+
+from .base import RerankerBase
+from .cross_encoder import CrossEncoderReranker
+from .naive import NaiveReranker
+
+__all__ = ["RerankerBase", "NaiveReranker", "CrossEncoderReranker"]
+
+logger = logging.getLogger(name=__name__)
+
+
+def get_reranker(configs: Config) -> RerankerBase:
+    if configs.reranker == "NaiveReranker" or not configs.reranker:
+        return NaiveReranker(configs)
+    elif configs.reranker == "CrossEncoderReranker":
+        return CrossEncoderReranker(configs)
+    else:
+        logger.warning(
+            f"""
+"reranker" option should be set to one of the following: {list(i for i in __all__ if i != "RerankerBase")}.
+To choose a custom reranker model, you can set the "model_name_or_path" key in the "reranker_params" option to the name/path of the model.
+The old configuration syntax will be deprecated in v0.6.0
+                """
+        )
+        configs.reranker_params.update({"model_name_or_path": configs.reranker})
+        configs.reranker = "CrossEncoderReranker"
+        return CrossEncoderReranker(
+            configs,
+        )

src/vectorcode/subcommands/query/reranker/__init__.py

@@ -0,0 +1,19 @@
+import logging
+from abc import abstractmethod
+from typing import Any
+
+from chromadb.api.types import QueryResult
+
+from vectorcode.cli_utils import Config
+
+logger = logging.getLogger(name=__name__)
+
+
+class RerankerBase:
+    def __init__(self, configs: Config, **kwargs: Any):
+        self.configs = configs
+        self.n_result = configs.n_result
+
+    @abstractmethod
+    def rerank(self, results: QueryResult, query_chunks: list[str]) -> list[str]:
+        raise NotImplementedError

src/vectorcode/subcommands/query/reranker/base.py

@@ -0,0 +1,54 @@
+import heapq
+import logging
+from collections import defaultdict
+from typing import Any, DefaultDict
+
+import numpy
+
+from vectorcode.cli_utils import Config, QueryInclude
+
+from .base import RerankerBase
+
+logger = logging.getLogger(name=__name__)
+
+
+class CrossEncoderReranker(RerankerBase):
+    def __init__(
+        self,
+        configs: Config,
+        **kwargs: Any,
+    ):
+        super().__init__(configs)
+        from sentence_transformers import CrossEncoder
+
+        self.model = CrossEncoder(**configs.reranker_params)
+
+    def rerank(self, results, query_chunks) -> list[str]:
+        self.query_chunks = query_chunks
+        assert results["metadatas"] is not None
+        assert results["documents"] is not None
+        documents: DefaultDict[str, list[float]] = defaultdict(list)
+        for query_chunk_idx in range(len(self.query_chunks)):
+            chunk_ids = results["ids"][query_chunk_idx]
+            chunk_metas = results["metadatas"][query_chunk_idx]
+            chunk_docs = results["documents"][query_chunk_idx]
+            ranks = self.model.rank(
+                self.query_chunks[query_chunk_idx], chunk_docs, apply_softmax=True
+            )
+            for rank in ranks:
+                if QueryInclude.chunk in self.configs.include:
+                    documents[chunk_ids[rank["corpus_id"]]].append(float(rank["score"]))
+                else:
+                    documents[chunk_metas[rank["corpus_id"]]["path"]].append(
+                        float(rank["score"])
+                    )
+        logger.debug("Document scores: %s", documents)
+        top_k = int(numpy.mean(tuple(len(i) for i in documents.values())))
+        for key in documents.keys():
+            documents[key] = heapq.nlargest(top_k, documents[key])
+
+        return heapq.nlargest(
+            self.n_result,
+            documents.keys(),
+            key=lambda x: float(numpy.mean(documents[x])),
+        )

src/vectorcode/subcommands/query/reranker/cross_encoder.py

@@ -0,0 +1,45 @@
+import heapq
+import logging
+from collections import defaultdict
+from typing import Any, DefaultDict
+
+import numpy
+
+from vectorcode.cli_utils import Config, QueryInclude
+
+from .base import RerankerBase
+
+logger = logging.getLogger(name=__name__)
+
+
+class NaiveReranker(RerankerBase):
+    def __init__(self, configs: Config, **kwargs: Any):
+        super().__init__(configs)
+
+    def rerank(self, results, query_chunks) -> list[str]:
+        assert results["metadatas"] is not None
+        assert results["distances"] is not None
+        documents: DefaultDict[str, list[float]] = defaultdict(list)
+        for query_chunk_idx in range(len(results["ids"])):
+            chunk_ids = results["ids"][query_chunk_idx]
+            chunk_metas = results["metadatas"][query_chunk_idx]
+            chunk_distances = results["distances"][query_chunk_idx]
+            # NOTE: distances, smaller is better.
+            paths = [str(meta["path"]) for meta in chunk_metas]
+            assert len(paths) == len(chunk_distances)
+            for distance, identifier in zip(
+                chunk_distances,
+                chunk_ids if QueryInclude.chunk in self.configs.include else paths,
+            ):
+                if identifier is None:  # pragma: nocover
+                    # so that vectorcode doesn't break on old collections.
+                    continue
+                documents[identifier].append(distance)
+        logger.debug("Document scores: %s", documents)
+        top_k = int(numpy.mean(tuple(len(i) for i in documents.values())))
+        for key in documents.keys():
+            documents[key] = heapq.nsmallest(top_k, documents[key])
+
+        return heapq.nsmallest(
+            self.n_result, documents.keys(), lambda x: float(numpy.mean(documents[x]))
+        )

src/vectorcode/subcommands/query/reranker/naive.py

@@ -62,14 +62,14 @@ def mock_config():
 @pytest.mark.asyncio
 async def test_get_query_result_files(mock_collection, mock_config):
     # Mock the reranker
-    with patch("vectorcode.subcommands.query.reranker.NaiveReranker") as MockReranker:
+    with patch("vectorcode.subcommands.query.get_reranker") as mock_get_reranker:
         mock_reranker_instance = MagicMock()
         mock_reranker_instance.rerank.return_value = [
             "file1.py",
             "file2.py",
             "file3.py",
         ]
-        MockReranker.return_value = mock_reranker_instance
+        mock_get_reranker.return_value = mock_reranker_instance
 
         # Call the function
         result = await get_query_result_files(mock_collection, mock_config)
@@ -87,9 +87,9 @@ async def test_get_query_result_files(mock_collection, mock_config):
         assert not kwargs["where"]  # Since query_exclude is empty
 
         # Check reranker was used correctly
-        MockReranker.assert_called_once_with(mock_config)
+        mock_get_reranker.assert_called_once_with(mock_config)
         mock_reranker_instance.rerank.assert_called_once_with(
-            mock_collection.query.return_value
+            mock_collection.query.return_value, mock_config.query
         )
 
         # Check the result