colbert_indexing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a68780c1-d46a-49a7-936e-9cbec4002abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git -C ColBERT/ pull || git clone https://github.com/stanford-futuredata/ColBERT.git\n",
    "import sys; sys.path.insert(0, 'ColBERT/')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb267215-b70f-4c50-bb0e-98138398d040",
   "metadata": {},
   "outputs": [],
   "source": [
    "try: # When on google Colab, let's install all dependencies with pip.\n",
    "    import google.colab\n",
    "    !pip install -U pip\n",
    "    !pip install -e ColBERT/['faiss-gpu','torch']\n",
    "except Exception:\n",
    "  import sys; sys.path.insert(0, 'ColBERT/')\n",
    "  try:\n",
    "    from colbert import Indexer, Searcher\n",
    "  except Exception:\n",
    "    print(\"If you're running outside Colab, please make sure you install ColBERT in conda following the instructions in our README. You can also install (as above) with pip but it may install slower or less stable faiss or torch dependencies. Conda is recommended.\")\n",
    "    assert False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f803e8ee-8de1-41d3-8e00-e0c27fb3941b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import colbert\n",
    "from colbert import Indexer, Searcher\n",
    "from colbert.infra import Run, RunConfig, ColBERTConfig\n",
    "from colbert.data import Queries, Collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f164a853-a4a1-41f7-be9c-205b13539881",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "dataset = 'saf'\n",
    "datasplit = 'train'\n",
    "\n",
    "\n",
    "def create_entry(row):\n",
    "    return 'Input: {\"question\": \"' + str(row['question']) + '\", \"student_answer\": \"' + str(row['student']) + '\", \"reference_answer\": \"' + str(row['reference']) + '\"} Output: {\"label\": \"' + row['label'] + '\", \"numeric_score\": ' + str(row['score']) + ', \"feedback\": \"' + row['feedback'] + '\"}'\n",
    "\n",
    "df = pd.read_csv(\"data/train.csv\")\n",
    "\n",
    "df['Entry'] = df.apply(create_entry, axis=1)\n",
    "entries = df['Entry'].tolist()\n",
    "collection = entries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b310205e-98c8-4da9-9ec4-f5b7da970fec",
   "metadata": {},
   "outputs": [],
   "source": [
    "collection[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "702b7cf7-c249-4717-80df-6c6e40f2a755",
   "metadata": {},
   "source": [
    "## Indexing\n",
    "\n",
    "For an efficient search, we can pre-compute the ColBERT representation of each passage and index them.\n",
    "\n",
    "Below, the `Indexer` take a model checkpoint and writes a (compressed) index to disk. We then prepare a `Searcher` for retrieval from this index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91e7d746-a657-4b07-85c3-579d4822473b",
   "metadata": {},
   "outputs": [],
   "source": [
    "nbits = 2   # encode each dimension with 2 bits\n",
    "doc_maxlen = 300 # truncate passages at 300 tokens\n",
    "# max_id = 10000\n",
    "\n",
    "index_name = f'{dataset}.{datasplit}.{nbits}bits'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13984279-9765-48d9-aec8-ad8e2c87a640",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = 'colbert-ir/colbertv2.0'\n",
    "\n",
    "with Run().context(RunConfig(nranks=1, experiment='notebook')):  # nranks specifies the number of GPUs to use\n",
    "    config = ColBERTConfig(doc_maxlen=doc_maxlen, nbits=nbits, kmeans_niters=4) # kmeans_niters specifies the number of iterations of k-means clustering; 4 is a good and fast default.\n",
    "                                                                                # Consider larger numbers for small datasets.\n",
    "\n",
    "    indexer = Indexer(checkpoint=checkpoint, config=config)\n",
    "    indexer.index(name=index_name, collection=collection, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6929c3f4-a1d7-487a-8375-3140f5c8afe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "indexer.get_index() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "server-env",
   "language": "python",
   "name": "server-env"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}