From 1ef667ec321a42cfd907465f1f2d91010609990a Mon Sep 17 00:00:00 2001 From: Jamie Milsom Date: Thu, 27 Aug 2026 13:01:05 +0100 Subject: [PATCH 1/3] feat: add reranker hook using huggingface transformers --- src/classifai/indexers/__init__.py | 2 + src/classifai/indexers/hooks/__init__.py | 8 +- .../indexers/hooks/default_hooks/__init__.py | 9 +- .../hooks/default_hooks/postprocessing.py | 173 +++++++++++++++++- src/classifai/vectorisers/huggingface.py | 2 + 5 files changed, 190 insertions(+), 4 deletions(-) diff --git a/src/classifai/indexers/__init__.py b/src/classifai/indexers/__init__.py index 2636ce0..86df146 100644 --- a/src/classifai/indexers/__init__.py +++ b/src/classifai/indexers/__init__.py @@ -40,6 +40,7 @@ ) from .hooks import ( CapitalisationStandardisingHook, + CrossEncoderRerankerHook, DeduplicationHook, HookBase, ) @@ -47,6 +48,7 @@ __all__ = [ "CapitalisationStandardisingHook", + "CrossEncoderRerankerHook", "DeduplicationHook", "HookBase", "VectorStore", diff --git a/src/classifai/indexers/hooks/__init__.py b/src/classifai/indexers/hooks/__init__.py index 712f742..c3d170c 100644 --- a/src/classifai/indexers/hooks/__init__.py +++ b/src/classifai/indexers/hooks/__init__.py @@ -5,11 +5,17 @@ This submodule exposes HookBase for building configurable hooks, plus a set of default hooks for common workflows. """ -from .default_hooks import CapitalisationStandardisingHook, DeduplicationHook, RagHook +from .default_hooks import ( + CapitalisationStandardisingHook, + CrossEncoderRerankerHook, + DeduplicationHook, + RagHook, +) from .hook_factory import HookBase __all__ = [ "CapitalisationStandardisingHook", + "CrossEncoderRerankerHook", "DeduplicationHook", "HookBase", "RagHook", diff --git a/src/classifai/indexers/hooks/default_hooks/__init__.py b/src/classifai/indexers/hooks/default_hooks/__init__.py index 89153ae..e0f65ee 100644 --- a/src/classifai/indexers/hooks/default_hooks/__init__.py +++ b/src/classifai/indexers/hooks/default_hooks/__init__.py @@ -1,6 +1,11 @@ """Submodule containing the prebuilt hooks for the service.""" -from .postprocessing import DeduplicationHook, RagHook +from .postprocessing import CrossEncoderRerankerHook, DeduplicationHook, RagHook from .preprocessing import CapitalisationStandardisingHook -__all__ = ["CapitalisationStandardisingHook", "DeduplicationHook", "RagHook"] +__all__ = [ + "CapitalisationStandardisingHook", + "CrossEncoderRerankerHook", + "DeduplicationHook", + "RagHook", +] diff --git a/src/classifai/indexers/hooks/default_hooks/postprocessing.py b/src/classifai/indexers/hooks/default_hooks/postprocessing.py index d608658..62fc931 100644 --- a/src/classifai/indexers/hooks/default_hooks/postprocessing.py +++ b/src/classifai/indexers/hooks/default_hooks/postprocessing.py @@ -6,7 +6,7 @@ import pandas as pd from classifai._optional import check_deps -from classifai.exceptions import ConfigurationError, HookError +from classifai.exceptions import ConfigurationError, ExternalServiceError, HookError from classifai.indexers.dataclasses import VectorStoreSearchOutput from classifai.indexers.hooks.hook_factory import HookBase @@ -325,3 +325,174 @@ def __call__(self, search_output: VectorStoreSearchOutput) -> VectorStoreSearchO """ processed_output = self._call_llm(search_output) return processed_output + + +class CrossEncoderRerankerHook(HookBase): + """A post-processing hook to rerank search results using a cross-encoder model. + + Designed to operate on a `VectorStoreSearchOutput`, i.e. the output of + the `VectorStore.search()` method. Takes the top retrieved results and + reranks them using a cross-encoder model that scores (query, document) + pairs jointly, often providing more accurate results than bi-encoder + similarity scores. + + Attributes: + model_name (str): The name of the cross-encoder model to use from huggingface. + device (torch.device): [optional] The device to use for + computation. Defaults to GPU if available, otherwise CPU. + model_revision (str): [optional] The specific model revision to + use. Defaults to "main". + tokenizer_kwargs (dict): [optional] Additional keyword arguments to + pass to the tokenizer. Defaults to None. + model_kwargs (dict): [optional] Additional keyword arguments to + pass to the model. Defaults to None. + + Raises: + `ExternalServiceError`: If the model or tokenizer cannot be loaded. + `ConfigurationError`: If the model cannot be initialised on the + specified device. + """ + + def __init__( + self, + model_name: str = "BAAI/bge-reranker-v2-m3", + device=None, + model_revision: str = "main", + tokenizer_kwargs: dict | None = None, + model_kwargs: dict | None = None, + ): + """Initialises the hook with the specified cross-encoder model. + + Args: + model_name (str): The name of the cross-encoder model from + Hugging Face Hub. Defaults to "BAAI/bge-reranker-v2-m3", + a high-performance reranker suitable for local deployment. + device (torch.device): [optional] The device to use for + computation. Defaults to MPS if available (Apple Silicon), + else GPU if available, otherwise CPU. + model_revision (str): [optional] The specific model revision to + use. Defaults to "main". + tokenizer_kwargs (dict): [optional] Additional keyword arguments to + pass to the tokenizer. Defaults to None. + model_kwargs (dict): [optional] Additional keyword arguments to + pass to the model. Defaults to None. + + Raises: + ExternalServiceError: If the model or tokenizer cannot be loaded. + ConfigurationError: If the model cannot be initialised on the + specified device. + """ + check_deps(["transformers", "torch"], extra="huggingface") + import torch + from transformers import AutoModelForSequenceClassification, AutoTokenizer # type: ignore + + self.model_name = model_name + + tokenizer_kwargs = dict(tokenizer_kwargs or {}) + model_kwargs = dict(model_kwargs or {}) + + # Ensure consistent behavior unless user overrides it + tokenizer_kwargs.setdefault("trust_remote_code", False) + model_kwargs.setdefault("trust_remote_code", False) + + try: + self.tokenizer = AutoTokenizer.from_pretrained(model_name, revision=model_revision, **tokenizer_kwargs) # nosec: B615 + self.model = AutoModelForSequenceClassification.from_pretrained( + model_name, revision=model_revision, **model_kwargs + ) # nosec: B615 + except Exception as e: + raise ExternalServiceError( + "Failed to load Hugging Face cross-encoder model/tokenizer.", + context={ + "hook": "CrossEncoderRerankerHook", + "model": model_name, + "revision": model_revision, + "cause": str(e), + "cause_type": type(e).__name__, + }, + ) from e + + # Device selection / model placement is local configuration/runtime. + try: + if device is not None: + self.device = device + elif torch.backends.mps.is_available(): + self.device = torch.device("mps") + else: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + self.model.to(self.device) + self.model.eval() + except Exception as e: + raise ConfigurationError( + "Failed to initialise cross-encoder model on device.", + context={ + "hook": "CrossEncoderRerankerHook", + "model": model_name, + "device": str(device) if device else "auto", + "cause": str(e), + "cause_type": type(e).__name__, + }, + ) from e + + super().__init__(hook_type="post_processing") + + def __call__(self, data: VectorStoreSearchOutput) -> VectorStoreSearchOutput: + """Reranks search results using the cross-encoder model. + + For each query in the search output, creates (query, document) pairs, + scores them with the cross-encoder model, sorts results by the new + scores, and reassigns ranks accordingly. + + Args: + data (VectorStoreSearchOutput): The search output data to rerank. + + Returns: + A new VectorStoreSearchOutput with results reranked by the + cross-encoder model scores and ranks reassigned. + + Raises: + HookError: If cross-encoder inference fails. + """ + import torch + + df = data.copy() + num_output_classes = 2 + + try: + pairs = df[["query_text", "doc_text"]].values.tolist() + + with torch.no_grad(): + inputs = self.tokenizer(pairs, padding=True, truncation=True, return_tensors="pt", max_length=512).to( + self.device + ) + logits = self.model(**inputs).logits + + if logits.shape[1] == num_output_classes: + # Binary classification: apply softmax per pair and use positive class + logits = torch.softmax(logits, dim=1) + new_scores = logits[:, 1].cpu().numpy() + else: + # Single score per pair: normalize scores across all pairs + logits = logits.squeeze() + new_scores = torch.softmax(logits, dim=0).cpu().numpy() + + except Exception as e: + raise HookError( + "Cross-encoder reranking inference failed.", + context={ + "postprocessing": "CrossEncoderRerankerHook", + "model": self.model_name, + "device": str(self.device), + "n_pairs": len(pairs), + "cause": str(e), + "cause_type": type(e).__name__, + }, + ) from e + + df["score"] = new_scores + df = df.sort_values(by=["query_id", "score"], ascending=[True, False]).reset_index(drop=True) + df["rank"] = df.groupby("query_id").cumcount() + 1 + + processed_output = data.__class__.validate(df) + return processed_output diff --git a/src/classifai/vectorisers/huggingface.py b/src/classifai/vectorisers/huggingface.py index 19d4598..ff171ee 100644 --- a/src/classifai/vectorisers/huggingface.py +++ b/src/classifai/vectorisers/huggingface.py @@ -84,6 +84,8 @@ def __init__( try: if device is not None: self.device = device + elif torch.backends.mps.is_available(): + self.device = torch.device("mps") else: self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") From b875e2df723f895702a2b052421f0cabfc75491f Mon Sep 17 00:00:00 2001 From: Jamie Milsom Date: Thu, 27 Aug 2026 13:02:18 +0100 Subject: [PATCH 2/3] feat: add demo notebook using mock soc data to show performance difference with cross embedding hook --- DEMO/reranker_hook.ipynb | 237 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 DEMO/reranker_hook.ipynb diff --git a/DEMO/reranker_hook.ipynb b/DEMO/reranker_hook.ipynb new file mode 100644 index 0000000..1531c77 --- /dev/null +++ b/DEMO/reranker_hook.ipynb @@ -0,0 +1,237 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# ✨ ClassifAI Demo - Cross-Encoder Reranking ✨\n", + "\n", + "---\n", + "\n", + "## Overview\n", + "\n", + "This notebook demonstrates how to use the `CrossEncoderRerankerHook` to improve search result quality.\n", + "\n", + "While the standard bi-encoder `VectorStore.search()` is fast, we can achieve much higher accuracy by passing retrieved results through a cross-encoder model that scores **(query, document) pairs jointly**.\n", + "\n", + "This demo uses:\n", + "- **Bi-encoder model**: `sentence-transformers/all-MiniLM-L6-v2` (fast retrieval)\n", + "- **Cross-encoder reranker**: `BAAI/bge-reranker-v2-m3` (accurate ranking)\n", + "\n", + "The cross-encoder runs efficiently on Apple Silicon (MPS), CUDA, or CPU." + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Setup: Import Libraries and Initialize Vectoriser\n", + "\n", + "First, we'll set up our bi-encoder vectoriser for initial retrieval." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from os import getenv\n", + "\n", + "import polars as pl\n", + "\n", + "from classifai.indexers import VectorStore\n", + "from classifai.indexers.dataclasses import VectorStoreSearchInput\n", + "from classifai.indexers.hooks import CrossEncoderRerankerHook\n", + "from classifai.vectorisers import GcpVectoriser, HuggingFaceVectoriser" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## Create VectorStore Without Reranking (Baseline)\n", + "\n", + "First, let's create a baseline VectorStore with just the bi-encoder scores, so we can compare results." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "vectoriser_hf = HuggingFaceVectoriser(model_name=\"sentence-transformers/all-MiniLM-L6-v2\")\n", + "\n", + "vectorstore_baseline_hf = VectorStore(\n", + " file_name=\"data/fake_soc_dataset.csv\",\n", + " data_type=\"csv\",\n", + " vectoriser=vectoriser_hf,\n", + " output_dir=\"testdata_baseline_hf\",\n", + " overwrite=True,\n", + " quiet_mode=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "Lets use a second baseline of the GCP embedding model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "vectoriser_gcp = GcpVectoriser(project_id=getenv(\"PROJECT_ID\"), location=\"europe-west2\", vertexai=True)\n", + "\n", + "vectorstore_baseline_gcp = VectorStore(\n", + " file_name=\"data/fake_soc_dataset.csv\",\n", + " data_type=\"csv\",\n", + " vectoriser=vectoriser_gcp,\n", + " output_dir=\"testdata_baseline_gcp\",\n", + " overwrite=True,\n", + " quiet_mode=False,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## Create VectorStore With Reranking Hook\n", + "\n", + "Now, let's create a VectorStore that uses the `CrossEncoderRerankerHook` to improve result ranking." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer_kwargs = {\"local_files_only\": True}\n", + "\n", + "model_kwargs = {\"local_files_only\": True}\n", + "\n", + "rerank_hook = CrossEncoderRerankerHook()\n", + "\n", + "vectorstore_reranked_hf = VectorStore(\n", + " file_name=\"data/fake_soc_dataset.csv\",\n", + " data_type=\"csv\",\n", + " vectoriser=vectoriser_hf,\n", + " output_dir=\"testdata_reranked_hf\",\n", + " overwrite=True,\n", + " quiet_mode=False,\n", + " hooks={\n", + " \"search_postprocess\": rerank_hook,\n", + " },\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer_kwargs = {\"local_files_only\": True}\n", + "\n", + "model_kwargs = {\"local_files_only\": True}\n", + "\n", + "rerank_hook = CrossEncoderRerankerHook()\n", + "\n", + "vectorstore_reranked_gcp = VectorStore(\n", + " file_name=\"data/fake_soc_dataset.csv\",\n", + " data_type=\"csv\",\n", + " vectoriser=vectoriser_gcp,\n", + " output_dir=\"testdata_reranked_gcp\",\n", + " overwrite=True,\n", + " quiet_mode=False,\n", + " hooks={\n", + " \"search_postprocess\": rerank_hook,\n", + " },\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "## Compare Results: Baseline vs. Reranked\n", + "\n", + "Let's run some searches and compare how the cross-encoder reranker improves result ordering." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "test_query_df = pl.read_csv(\"data/fake_soc_eval_queries.csv\")\n", + "test_query = VectorStoreSearchInput(test_query_df.rename({\"label\": \"id\", \"text\": \"query\"}).sample(1).to_pandas())\n", + "\n", + "baseline_results_hf = vectorstore_baseline_hf.search(test_query, n_results=5)\n", + "print(\"\\n\" + \"=\" * 80)\n", + "print(\"BASELINE RESULTS HF (Bi-encoder only)\")\n", + "print(\"=\" * 80)\n", + "print(baseline_results_hf[[\"query_text\", \"doc_text\", \"score\", \"rank\"]].to_string())\n", + "\n", + "reranked_results_hf = vectorstore_reranked_hf.search(test_query, n_results=100).head(5)\n", + "print(\"\\n\" + \"=\" * 80)\n", + "print(\"RERANKED RESULTS HF (Bi-encoder + Cross-encoder Reranker)\")\n", + "print(\"=\" * 80)\n", + "print(reranked_results_hf[[\"query_text\", \"doc_text\", \"score\", \"rank\"]].to_string())\n", + "\n", + "baseline_results_gcp = vectorstore_baseline_gcp.search(test_query, n_results=5)\n", + "print(\"\\n\" + \"=\" * 80)\n", + "print(\"BASELINE RESULTS GCP (Bi-encoder only)\")\n", + "print(\"=\" * 80)\n", + "print(baseline_results_gcp[[\"query_text\", \"doc_text\", \"score\", \"rank\"]].to_string())\n", + "\n", + "reranked_results_gcp = vectorstore_reranked_gcp.search(test_query, n_results=100).head(5)\n", + "print(\"\\n\" + \"=\" * 80)\n", + "print(\"RERANKED RESULTS GCP (Bi-encoder + Cross-encoder Reranker)\")\n", + "print(\"=\" * 80)\n", + "print(reranked_results_gcp[[\"query_text\", \"doc_text\", \"score\", \"rank\"]].to_string())" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "classifai", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From c36e0925a6b6e6a532040c9005a4fd00d1f67a70 Mon Sep 17 00:00:00 2001 From: Jamie Milsom Date: Thu, 27 Aug 2026 15:05:07 +0100 Subject: [PATCH 3/3] feat: updated default model to cross-encoder/ms-marco-MiniLM-L-12-v2 for improved performance --- src/classifai/indexers/hooks/default_hooks/postprocessing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/classifai/indexers/hooks/default_hooks/postprocessing.py b/src/classifai/indexers/hooks/default_hooks/postprocessing.py index 62fc931..e6ed060 100644 --- a/src/classifai/indexers/hooks/default_hooks/postprocessing.py +++ b/src/classifai/indexers/hooks/default_hooks/postprocessing.py @@ -355,7 +355,7 @@ class CrossEncoderRerankerHook(HookBase): def __init__( self, - model_name: str = "BAAI/bge-reranker-v2-m3", + model_name: str = "cross-encoder/ms-marco-MiniLM-L-12-v2", device=None, model_revision: str = "main", tokenizer_kwargs: dict | None = None, @@ -365,7 +365,7 @@ def __init__( Args: model_name (str): The name of the cross-encoder model from - Hugging Face Hub. Defaults to "BAAI/bge-reranker-v2-m3", + Hugging Face Hub. Defaults to "cross-encoder/ms-marco-MiniLM-L-12-v2", a high-performance reranker suitable for local deployment. device (torch.device): [optional] The device to use for computation. Defaults to MPS if available (Apple Silicon),