feat: Add a CrossEncoderRerankerHook as a post-processing hook

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
+}

src/classifai/indexers/__init__.py

@@ -40,13 +40,15 @@
 )
 from .hooks import (
     CapitalisationStandardisingHook,
+    CrossEncoderRerankerHook,
     DeduplicationHook,
     HookBase,
 )
 from .main import VectorStore
 
 __all__ = [
     "CapitalisationStandardisingHook",
+    "CrossEncoderRerankerHook",
     "DeduplicationHook",
     "HookBase",
     "VectorStore",

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",

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",
+]