Adding model and checkpoint

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.dnn filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,159 @@
+# ColBERT
+
+### ColBERT is a _fast_ and _accurate_ retrieval model, enabling scalable BERT-based search over large text collections in tens of milliseconds. 
+
+
+<p align="center">
+  <img align="center" src="docs/images/ColBERT-Framework-MaxSim-W370px.png" />
+</p>
+<p align="center">
+  <b>Figure 1:</b> ColBERT's late interaction, efficiently scoring the fine-grained similarity between a queries and a passage.
+</p>
+
+As Figure 1 illustrates, ColBERT relies on fine-grained **contextual late interaction**: it encodes each passage into a **matrix** of token-level embeddings (shown above in blue). Then at search time, it embeds every query into another matrix (shown in green) and efficiently finds passages that contextually match the query using scalable vector-similarity (`MaxSim`) operators.
+
+These rich interactions allow ColBERT to surpass the quality of _single-vector_ representation models, while scaling efficiently to large corpora. You can read more in our papers:
+
+* [**ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT**](https://arxiv.org/abs/2004.12832) (SIGIR'20).
+* [**Relevance-guided Supervision for OpenQA with ColBERT**](https://arxiv.org/abs/2007.00814) (TACL'21; to appear).
+
+
+----
+
+## Installation
+
+ColBERT (currently: [v0.2.0](#releases)) requires Python 3.7+ and Pytorch 1.6+ and uses the [HuggingFace Transformers](https://github.com/huggingface/transformers) library.
+
+We strongly recommend creating a conda environment using:
+
+```
+conda env create -f conda_env.yml
+conda activate colbert-v0.2
+```
+
+If you face any problems, please [open a new issue](https://github.com/stanford-futuredata/ColBERT/issues) and we'll help you promptly!
+
+
+## Overview
+
+Using ColBERT on a dataset typically involves the following steps.
+
+**Step 0: Preprocess your collection.** At its simplest, ColBERT works with tab-separated (TSV) files: a file (e.g., `collection.tsv`) will contain all passages and another (e.g., `queries.tsv`) will contain a set of queries for searching the collection.
+
+**Step 1: Train a ColBERT model.**  You can [train your own ColBERT model](#training) and [validate performance](#validation) on a suitable development set.
+
+**Step 2: Index your collection.** Once you're happy with your ColBERT model, you need to [index your collection](#indexing) to permit fast retrieval. This step encodes all passages into matrices, stores them on disk, and builds data structures for efficient search.
+
+**Step 3: Search the collection with your queries.** Given your model and index, you can [issue queries over the collection](#retrieval) to retrieve the top-k passages for each query.
+
+Below, we illustrate these steps via an example run on the MS MARCO Passage Ranking task.
+
+
+## Data
+
+This repository works directly with a simple **tab-separated file** format to store queries, passages, and top-k ranked lists.
+
+
+* Queries: each line is `qid \t query text`.
+* Collection: each line is `pid \t passage text`. 
+* Top-k Ranking: each line is `qid \t pid \t rank`.
+
+This works directly with the data format of the [MS MARCO Passage Ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking) dataset. You will need the training triples (`triples.train.small.tar.gz`), the official top-1000 ranked lists for the dev set queries (`top1000.dev`), and the dev set relevant passages (`qrels.dev.small.tsv`). For indexing the full collection, you will also need the list of passages (`collection.tar.gz`).
+
+
+
+## Training
+
+Training requires a list of _<query, positive passage, negative passage>_ tab-separated triples.
+
+You can supply **full-text** triples, where each line is `query text \t positive passage text \t negative passage text`. Alternatively, you can supply the query and passage **IDs** as a JSONL file `[qid, pid+, pid-]` per line, in which case you should specify `--collection path/to/collection.tsv` and `--queries path/to/queries.train.tsv`.
+
+
+```
+CUDA_VISIBLE_DEVICES="0,1,2,3" \
+python -m torch.distributed.launch --nproc_per_node=4 -m \
+colbert.train --amp --doc_maxlen 180 --mask-punctuation --bsize 32 --accum 1 \
+--triples /path/to/MSMARCO/triples.train.small.tsv \
+--root /root/to/experiments/ --experiment MSMARCO-psg --similarity l2 --run msmarco.psg.l2
+```
+
+You can use one or more GPUs by modifying `CUDA_VISIBLE_DEVICES` and `--nproc_per_node`.
+
+
+## Validation
+
+Before indexing into ColBERT, you can compare a few checkpoints by re-ranking a top-k set of documents per query. This will use ColBERT _on-the-fly_: it will compute document representations _during_ query evaluation.
+
+This script requires the top-k list per query, provided as a tab-separated file whose every line contains a tuple `queryID \t passageID \t rank`, where rank is {1, 2, 3, ...} for each query. The script also accepts the format of MS MARCO's `top1000.dev` and `top1000.eval` and you can optionally supply relevance judgements (qrels) for evaluation. This is a tab-separated file whose every line has a quadruple _<query ID, 0, passage ID, 1>_, like `qrels.dev.small.tsv`.
+
+Example command:
+
+```
+python -m colbert.test --amp --doc_maxlen 180 --mask-punctuation \
+--collection /path/to/MSMARCO/collection.tsv \
+--queries /path/to/MSMARCO/queries.dev.small.tsv \
+--topk /path/to/MSMARCO/top1000.dev  \
+--checkpoint /root/to/experiments/MSMARCO-psg/train.py/msmarco.psg.l2/checkpoints/colbert-200000.dnn \
+--root /root/to/experiments/ --experiment MSMARCO-psg  [--qrels path/to/qrels.dev.small.tsv]
+```
+
+
+## Indexing
+
+For fast retrieval, indexing precomputes the ColBERT representations of passages.
+
+Example command:
+
+```
+CUDA_VISIBLE_DEVICES="0,1,2,3" OMP_NUM_THREADS=6 \
+python -m torch.distributed.launch --nproc_per_node=4 -m \
+colbert.index --amp --doc_maxlen 180 --mask-punctuation --bsize 256 \
+--checkpoint /root/to/experiments/MSMARCO-psg/train.py/msmarco.psg.l2/checkpoints/colbert-200000.dnn \
+--collection /path/to/MSMARCO/collection.tsv \
+--index_root /root/to/indexes/ --index_name MSMARCO.L2.32x200k \
+--root /root/to/experiments/ --experiment MSMARCO-psg
+```
+
+The index created here allows you to re-rank the top-k passages retrieved by another method (e.g., BM25).
+
+We typically recommend that you use ColBERT for **end-to-end** retrieval, where it directly finds its top-k passages from the full collection. For this, you need FAISS indexing.
+
+
+#### FAISS Indexing for end-to-end retrieval
+
+For end-to-end retrieval, you should index the document representations into [FAISS](https://github.com/facebookresearch/faiss).
+
+```
+python -m colbert.index_faiss \
+--index_root /root/to/indexes/ --index_name MSMARCO.L2.32x200k \
+--partitions 32768 --sample 0.3 \
+--root /root/to/experiments/ --experiment MSMARCO-psg
+```
+
+
+## Retrieval
+
+In the simplest case, you want to retrieve from the full collection:
+
+```
+python -m colbert.retrieve \
+--amp --doc_maxlen 180 --mask-punctuation --bsize 256 \
+--queries /path/to/MSMARCO/queries.dev.small.tsv
+--nprobe 32 --partitions 32768 --faiss_depth 1024 \
+--index_root /root/to/indexes/ --index_name MSMARCO.L2.32x200k \
+--checkpoint /root/to/experiments/MSMARCO-psg/train.py/msmarco.psg.l2/checkpoints/colbert-200000.dnn \
+--root /root/to/experiments/ --experiment MSMARCO-psg
+```
+
+You may also want to re-rank a top-k set that you've retrieved before with ColBERT or with another model. For this, use `colbert.rerank` similarly and additionally pass `--topk`.
+
+If you have a large set of queries (or want to reduce memory usage), use **batch-mode** retrieval and/or re-ranking. This can be done by passing `--batch --only_retrieval` to `colbert.retrieve` and passing `--batch --log-scores` to colbert.rerank alongside `--topk` with the `unordered.tsv` output of this retrieval run.
+
+Some use cases (e.g., building a user-facing search engines) require more control over retrieval. For those, you typically don't want to use the command line for retrieval. Instead, you want to import our retrieval API from Python and directly work with that (e.g., to build a simple REST API). Instructions for this are coming soon, but you will just need to adapt/modify the retrieval loop in [`colbert/ranking/retrieval.py#L33`](colbert/ranking/retrieval.py#L33).
+
+
+## Releases
+
+* v0.2.0: Sep 2020
+* v0.1.0: June 2020
+

colbert/evaluation/load_model.py

@@ -0,0 +1,28 @@
+import os
+import ujson
+import torch
+import random
+
+from collections import defaultdict, OrderedDict
+
+from colbert.parameters import DEVICE
+from colbert.modeling.colbert import ColBERT
+from colbert.utils.utils import print_message, load_checkpoint
+
+
+def load_model(args, do_print=True):
+    colbert = ColBERT.from_pretrained('bert-base-multilingual-uncased',
+                                      query_maxlen=args.query_maxlen,
+                                      doc_maxlen=args.doc_maxlen,
+                                      dim=args.dim,
+                                      similarity_metric=args.similarity,
+                                      mask_punctuation=args.mask_punctuation)
+    colbert = colbert.to(DEVICE)
+
+    print_message("#> Loading model checkpoint.", condition=do_print)
+
+    checkpoint = load_checkpoint(args.checkpoint, colbert, do_print=do_print)
+
+    colbert.eval()
+
+    return colbert, checkpoint

colbert/evaluation/loaders.py

@@ -0,0 +1,196 @@
+import os
+import ujson
+import torch
+import random
+
+from collections import defaultdict, OrderedDict
+
+from colbert.parameters import DEVICE
+from colbert.modeling.colbert import ColBERT
+from colbert.utils.utils import print_message, load_checkpoint
+from colbert.evaluation.load_model import load_model
+from colbert.utils.runs import Run
+
+
+def load_queries(queries_path):
+    queries = OrderedDict()
+
+    print_message("#> Loading the queries from", queries_path, "...")
+
+    with open(queries_path) as f:
+        for line in f:
+            qid, query, *_ = line.strip().split('\t')
+            qid = int(qid)
+
+            assert (qid not in queries), ("Query QID", qid, "is repeated!")
+            queries[qid] = query
+
+    print_message("#> Got", len(queries), "queries. All QIDs are unique.\n")
+
+    return queries
+
+
+def load_qrels(qrels_path):
+    if qrels_path is None:
+        return None
+
+    print_message("#> Loading qrels from", qrels_path, "...")
+
+    qrels = OrderedDict()
+    with open(qrels_path, mode='r', encoding="utf-8") as f:
+        for line in f:
+            qid, x, pid, y = map(int, line.strip().split('\t'))
+            assert x == 0 and y == 1
+            qrels[qid] = qrels.get(qid, [])
+            qrels[qid].append(pid)
+
+    assert all(len(qrels[qid]) == len(set(qrels[qid])) for qid in qrels)
+
+    avg_positive = round(sum(len(qrels[qid]) for qid in qrels) / len(qrels), 2)
+
+    print_message("#> Loaded qrels for", len(qrels), "unique queries with",
+                  avg_positive, "positives per query on average.\n")
+
+    return qrels
+
+
+def load_topK(topK_path):
+    queries = OrderedDict()
+    topK_docs = OrderedDict()
+    topK_pids = OrderedDict()
+
+    print_message("#> Loading the top-k per query from", topK_path, "...")
+
+    with open(topK_path) as f:
+        for line_idx, line in enumerate(f):
+            if line_idx and line_idx % (10*1000*1000) == 0:
+                print(line_idx, end=' ', flush=True)
+
+            qid, pid, query, passage = line.split('\t')
+            qid, pid = int(qid), int(pid)
+
+            assert (qid not in queries) or (queries[qid] == query)
+            queries[qid] = query
+            topK_docs[qid] = topK_docs.get(qid, [])
+            topK_docs[qid].append(passage)
+            topK_pids[qid] = topK_pids.get(qid, [])
+            topK_pids[qid].append(pid)
+
+        print()
+
+    assert all(len(topK_pids[qid]) == len(set(topK_pids[qid])) for qid in topK_pids)
+
+    Ks = [len(topK_pids[qid]) for qid in topK_pids]
+
+    print_message("#> max(Ks) =", max(Ks), ", avg(Ks) =", round(sum(Ks) / len(Ks), 2))
+    print_message("#> Loaded the top-k per query for", len(queries), "unique queries.\n")
+
+    return queries, topK_docs, topK_pids
+
+
+def load_topK_pids(topK_path, qrels):
+    topK_pids = defaultdict(list)
+    topK_positives = defaultdict(list)
+
+    print_message("#> Loading the top-k PIDs per query from", topK_path, "...")
+
+    with open(topK_path) as f:
+        for line_idx, line in enumerate(f):
+            if line_idx and line_idx % (10*1000*1000) == 0:
+                print(line_idx, end=' ', flush=True)
+
+            qid, pid, *rest = line.strip().split('\t')
+            qid, pid = int(qid), int(pid)
+
+            topK_pids[qid].append(pid)
+
+            assert len(rest) in [1, 2, 3]
+
+            if len(rest) > 1:
+                *_, label = rest
+                label = int(label)
+                assert label in [0, 1]
+
+                if label >= 1:
+                    topK_positives[qid].append(pid)
+
+        print()
+
+    assert all(len(topK_pids[qid]) == len(set(topK_pids[qid])) for qid in topK_pids)
+    assert all(len(topK_positives[qid]) == len(set(topK_positives[qid])) for qid in topK_positives)
+
+    # Make them sets for fast lookups later
+    topK_positives = {qid: set(topK_positives[qid]) for qid in topK_positives}
+
+    Ks = [len(topK_pids[qid]) for qid in topK_pids]
+
+    print_message("#> max(Ks) =", max(Ks), ", avg(Ks) =", round(sum(Ks) / len(Ks), 2))
+    print_message("#> Loaded the top-k per query for", len(topK_pids), "unique queries.\n")
+
+    if len(topK_positives) == 0:
+        topK_positives = None
+    else:
+        assert len(topK_pids) >= len(topK_positives)
+
+        for qid in set.difference(set(topK_pids.keys()), set(topK_positives.keys())):
+            topK_positives[qid] = []
+
+        assert len(topK_pids) == len(topK_positives)
+
+        avg_positive = round(sum(len(topK_positives[qid]) for qid in topK_positives) / len(topK_pids), 2)
+
+        print_message("#> Concurrently got annotations for", len(topK_positives), "unique queries with",
+                      avg_positive, "positives per query on average.\n")
+
+    assert qrels is None or topK_positives is None, "Cannot have both qrels and an annotated top-K file!"
+
+    if topK_positives is None:
+        topK_positives = qrels
+
+    return topK_pids, topK_positives
+
+
+def load_collection(collection_path):
+    print_message("#> Loading collection...")
+
+    collection = []
+
+    with open(collection_path) as f:
+        for line_idx, line in enumerate(f):
+            if line_idx % (1000*1000) == 0:
+                print(f'{line_idx // 1000 // 1000}M', end=' ', flush=True)
+
+            pid, passage, *rest = line.strip().split('\t')
+            assert pid == 'id' or int(pid) == line_idx
+
+            if len(rest) >= 1:
+                title = rest[0]
+                passage = title + ' | ' + passage
+
+            collection.append(passage)
+
+    print()
+
+    return collection
+
+
+def load_colbert(args, do_print=True):
+    colbert, checkpoint = load_model(args, do_print)
+
+    # TODO: If the parameters below were not specified on the command line, their *checkpoint* values should be used.
+    # I.e., not their purely (i.e., training) default values.
+
+    for k in ['query_maxlen', 'doc_maxlen', 'dim', 'similarity', 'amp']:
+        if 'arguments' in checkpoint and hasattr(args, k):
+            if k in checkpoint['arguments'] and checkpoint['arguments'][k] != getattr(args, k):
+                a, b = checkpoint['arguments'][k], getattr(args, k)
+                Run.warn(f"Got checkpoint['arguments']['{k}'] != args.{k} (i.e., {a} != {b})")
+
+    if 'arguments' in checkpoint:
+        if args.rank < 1:
+            print(ujson.dumps(checkpoint['arguments'], indent=4))
+
+    if do_print:
+        print('\n')
+
+    return colbert, checkpoint

colbert/evaluation/metrics.py

@@ -0,0 +1,114 @@
+import ujson
+
+from collections import defaultdict
+from colbert.utils.runs import Run
+
+
+class Metrics:
+    def __init__(self, mrr_depths: set, recall_depths: set, success_depths: set, total_queries=None):
+        self.results = {}
+        self.mrr_sums = {depth: 0.0 for depth in mrr_depths}
+        self.recall_sums = {depth: 0.0 for depth in recall_depths}
+        self.success_sums = {depth: 0.0 for depth in success_depths}
+        self.total_queries = total_queries
+
+        self.max_query_idx = -1
+        self.num_queries_added = 0
+
+    def add(self, query_idx, query_key, ranking, gold_positives):
+        self.num_queries_added += 1
+
+        assert query_key not in self.results
+        assert len(self.results) <= query_idx
+        assert len(set(gold_positives)) == len(gold_positives)
+        assert len(set([pid for _, pid, _ in ranking])) == len(ranking)
+
+        self.results[query_key] = ranking
+
+        positives = [i for i, (_, pid, _) in enumerate(ranking) if pid in gold_positives]
+
+        if len(positives) == 0:
+            return
+
+        for depth in self.mrr_sums:
+            first_positive = positives[0]
+            self.mrr_sums[depth] += (1.0 / (first_positive+1.0)) if first_positive < depth else 0.0
+
+        for depth in self.success_sums:
+            first_positive = positives[0]
+            self.success_sums[depth] += 1.0 if first_positive < depth else 0.0
+
+        for depth in self.recall_sums:
+            num_positives_up_to_depth = len([pos for pos in positives if pos < depth])
+            self.recall_sums[depth] += num_positives_up_to_depth / len(gold_positives)
+
+    def print_metrics(self, query_idx):
+        for depth in sorted(self.mrr_sums):
+            print("MRR@" + str(depth), "=", self.mrr_sums[depth] / (query_idx+1.0))
+
+        for depth in sorted(self.success_sums):
+            print("Success@" + str(depth), "=", self.success_sums[depth] / (query_idx+1.0))
+
+        for depth in sorted(self.recall_sums):
+            print("Recall@" + str(depth), "=", self.recall_sums[depth] / (query_idx+1.0))
+
+    def log(self, query_idx):
+        assert query_idx >= self.max_query_idx
+        self.max_query_idx = query_idx
+
+        Run.log_metric("ranking/max_query_idx", query_idx, query_idx)
+        Run.log_metric("ranking/num_queries_added", self.num_queries_added, query_idx)
+
+        for depth in sorted(self.mrr_sums):
+            score = self.mrr_sums[depth] / (query_idx+1.0)
+            Run.log_metric("ranking/MRR." + str(depth), score, query_idx)
+
+        for depth in sorted(self.success_sums):
+            score = self.success_sums[depth] / (query_idx+1.0)
+            Run.log_metric("ranking/Success." + str(depth), score, query_idx)
+
+        for depth in sorted(self.recall_sums):
+            score = self.recall_sums[depth] / (query_idx+1.0)
+            Run.log_metric("ranking/Recall." + str(depth), score, query_idx)
+
+    def output_final_metrics(self, path, query_idx, num_queries):
+        assert query_idx + 1 == num_queries
+        assert num_queries == self.total_queries
+
+        if self.max_query_idx < query_idx:
+            self.log(query_idx)
+
+        self.print_metrics(query_idx)
+
+        output = defaultdict(dict)
+
+        for depth in sorted(self.mrr_sums):
+            score = self.mrr_sums[depth] / (query_idx+1.0)
+            output['mrr'][depth] = score
+
+        for depth in sorted(self.success_sums):
+            score = self.success_sums[depth] / (query_idx+1.0)
+            output['success'][depth] = score
+
+        for depth in sorted(self.recall_sums):
+            score = self.recall_sums[depth] / (query_idx+1.0)
+            output['recall'][depth] = score
+
+        with open(path, 'w') as f:
+            ujson.dump(output, f, indent=4)
+            f.write('\n')
+
+
+def evaluate_recall(qrels, queries, topK_pids):
+    if qrels is None:
+        return
+
+    assert set(qrels.keys()) == set(queries.keys())
+    recall_at_k = [len(set.intersection(set(qrels[qid]), set(topK_pids[qid]))) / max(1.0, len(qrels[qid]))
+                   for qid in qrels]
+    recall_at_k = sum(recall_at_k) / len(qrels)
+    recall_at_k = round(recall_at_k, 3)
+    print("Recall @ maximum depth =", recall_at_k)
+
+
+# TODO: If implicit qrels are used (for re-ranking), warn if a recall metric is requested + add an asterisk to output.

colbert/evaluation/ranking.py

@@ -0,0 +1,88 @@
+import os
+import random
+import time
+import torch
+import torch.nn as nn
+
+from itertools import accumulate
+from math import ceil
+
+from colbert.utils.runs import Run
+from colbert.utils.utils import print_message
+
+from colbert.evaluation.metrics import Metrics
+from colbert.evaluation.ranking_logger import RankingLogger
+from colbert.modeling.inference import ModelInference
+
+from colbert.evaluation.slow import slow_rerank
+
+
+def evaluate(args):
+    args.inference = ModelInference(args.colbert, amp=args.amp)
+    qrels, queries, topK_pids = args.qrels, args.queries, args.topK_pids
+
+    depth = args.depth
+    collection = args.collection
+    if collection is None:
+        topK_docs = args.topK_docs
+
+    def qid2passages(qid):
+        if collection is not None:
+            return [collection[pid] for pid in topK_pids[qid][:depth]]
+        else:
+            return topK_docs[qid][:depth]
+
+    metrics = Metrics(mrr_depths={10, 100}, recall_depths={50, 200, 1000},
+                      success_depths={5, 10, 20, 50, 100, 1000},
+                      total_queries=len(queries))
+
+    ranking_logger = RankingLogger(Run.path, qrels=qrels)
+
+    args.milliseconds = []
+
+    with ranking_logger.context('ranking.tsv', also_save_annotations=(qrels is not None)) as rlogger:
+        with torch.no_grad():
+            keys = sorted(list(queries.keys()))
+            random.shuffle(keys)
+
+            for query_idx, qid in enumerate(keys):
+                query = queries[qid]
+
+                print_message(query_idx, qid, query, '\n')
+
+                if qrels and args.shortcircuit and len(set.intersection(set(qrels[qid]), set(topK_pids[qid]))) == 0:
+                    continue
+
+                ranking = slow_rerank(args, query, topK_pids[qid], qid2passages(qid))
+
+                rlogger.log(qid, ranking, [0, 1])
+
+                if qrels:
+                    metrics.add(query_idx, qid, ranking, qrels[qid])
+
+                    for i, (score, pid, passage) in enumerate(ranking):
+                        if pid in qrels[qid]:
+                            print("\n#> Found", pid, "at position", i+1, "with score", score)
+                            print(passage)
+                            break
+
+                    metrics.print_metrics(query_idx)
+                    metrics.log(query_idx)
+
+                print_message("#> checkpoint['batch'] =", args.checkpoint['batch'], '\n')
+                print("rlogger.filename =", rlogger.filename)
+
+                if len(args.milliseconds) > 1:
+                    print('Slow-Ranking Avg Latency =', sum(args.milliseconds[1:]) / len(args.milliseconds[1:]))
+
+                print("\n\n")
+
+        print("\n\n")
+        # print('Avg Latency =', sum(args.milliseconds[1:]) / len(args.milliseconds[1:]))
+        print("\n\n")
+
+    print('\n\n')
+    if qrels:
+        assert query_idx + 1 == len(keys) == len(set(keys))
+        metrics.output_final_metrics(os.path.join(Run.path, 'ranking.metrics'), query_idx, len(queries))
+    print('\n\n')

colbert/evaluation/ranking_logger.py

@@ -0,0 +1,57 @@
+import os
+
+from contextlib import contextmanager
+from colbert.utils.utils import print_message, NullContextManager
+from colbert.utils.runs import Run
+
+
+class RankingLogger():
+    def __init__(self, directory, qrels=None, log_scores=False):
+        self.directory = directory
+        self.qrels = qrels
+        self.filename, self.also_save_annotations = None, None
+        self.log_scores = log_scores
+
+    @contextmanager
+    def context(self, filename, also_save_annotations=False):
+        assert self.filename is None
+        assert self.also_save_annotations is None
+
+        filename = os.path.join(self.directory, filename)
+        self.filename, self.also_save_annotations = filename, also_save_annotations
+
+        print_message("#> Logging ranked lists to {}".format(self.filename))
+
+        with open(filename, 'w') as f:
+            self.f = f
+            with (open(filename + '.annotated', 'w') if also_save_annotations else NullContextManager()) as g:
+                self.g = g
+                try:
+                    yield self
+                finally:
+                    pass
+
+    def log(self, qid, ranking, is_ranked=True, print_positions=[]):
+        print_positions = set(print_positions)
+
+        f_buffer = []
+        g_buffer = []
+
+        for rank, (score, pid, passage) in enumerate(ranking):
+            is_relevant = self.qrels and int(pid in self.qrels[qid])
+            rank = rank+1 if is_ranked else -1
+
+            possibly_score = [score] if self.log_scores else []
+
+            f_buffer.append('\t'.join([str(x) for x in [qid, pid, rank] + possibly_score]) + "\n")
+            if self.g:
+                g_buffer.append('\t'.join([str(x) for x in [qid, pid, rank, is_relevant]]) + "\n")
+
+            if rank in print_positions:
+                prefix = "** " if is_relevant else ""
+                prefix += str(rank)
+                print("#> ( QID {} ) ".format(qid) + prefix + ") ", pid, ":", score, '    ', passage)
+
+        self.f.write(''.join(f_buffer))
+        if self.g:
+            self.g.write(''.join(g_buffer))

colbert/evaluation/slow.py

@@ -0,0 +1,21 @@
+import os
+
+def slow_rerank(args, query, pids, passages):
+    colbert = args.colbert
+    inference = args.inference
+
+    Q = inference.queryFromText([query])
+
+    D_ = inference.docFromText(passages, bsize=args.bsize)
+    scores = colbert.score(Q, D_).cpu()
+
+    scores = scores.sort(descending=True)
+    ranked = scores.indices.tolist()
+
+    ranked_scores = scores.values.tolist()
+    ranked_pids = [pids[position] for position in ranked]
+    ranked_passages = [passages[position] for position in ranked]
+
+    assert len(ranked_pids) == len(set(ranked_pids))
+
+    return list(zip(ranked_scores, ranked_pids, ranked_passages))

colbert/index.py

@@ -0,0 +1,59 @@
+import os
+import ujson
+import random
+
+from colbert.utils.runs import Run
+from colbert.utils.parser import Arguments
+import colbert.utils.distributed as distributed
+
+from colbert.utils.utils import print_message, create_directory
+from colbert.indexing.encoder import CollectionEncoder
+
+
+def main():
+    random.seed(12345)
+
+    parser = Arguments(description='Precomputing document representations with ColBERT.')
+
+    parser.add_model_parameters()
+    parser.add_model_inference_parameters()
+    parser.add_indexing_input()
+
+    parser.add_argument('--chunksize', dest='chunksize', default=6.0, required=False, type=float)   # in GiBs
+
+    args = parser.parse()
+
+    with Run.context():
+        args.index_path = os.path.join(args.index_root, args.index_name)
+        assert not os.path.exists(args.index_path), args.index_path
+
+        distributed.barrier(args.rank)
+
+        if args.rank < 1:
+            create_directory(args.index_root)
+            create_directory(args.index_path)
+
+        distributed.barrier(args.rank)
+
+        process_idx = max(0, args.rank)
+        encoder = CollectionEncoder(args, process_idx=process_idx, num_processes=args.nranks)
+        encoder.encode()
+
+        distributed.barrier(args.rank)
+
+        # Save metadata.
+        if args.rank < 1:
+            metadata_path = os.path.join(args.index_path, 'metadata.json')
+            print_message("Saving (the following) metadata to", metadata_path, "..")
+            print(args.input_arguments)
+
+            with open(metadata_path, 'w') as output_metadata:
+                ujson.dump(args.input_arguments.__dict__, output_metadata)
+
+        distributed.barrier(args.rank)
+
+
+if __name__ == "__main__":
+    main()
+
+# TODO: Add resume functionality

colbert/index_faiss.py

@@ -0,0 +1,43 @@
+import os
+import random
+import math
+
+from colbert.utils.runs import Run
+from colbert.utils.parser import Arguments
+from colbert.indexing.faiss import index_faiss
+from colbert.indexing.loaders import load_doclens
+
+
+def main():
+    random.seed(12345)
+
+    parser = Arguments(description='Faiss indexing for end-to-end retrieval with ColBERT.')
+    parser.add_index_use_input()
+
+    parser.add_argument('--sample', dest='sample', default=None, type=float)
+    parser.add_argument('--slices', dest='slices', default=1, type=int)
+
+    args = parser.parse()
+    assert args.slices >= 1
+    assert args.sample is None or (0.0 < args.sample < 1.0), args.sample
+
+    with Run.context():
+        args.index_path = os.path.join(args.index_root, args.index_name)
+        assert os.path.exists(args.index_path), args.index_path
+
+        num_embeddings = sum(load_doclens(args.index_path))
+        print("#> num_embeddings =", num_embeddings)
+
+        if args.partitions is None:
+            args.partitions = 1 << math.ceil(math.log2(8 * math.sqrt(num_embeddings)))
+            print('\n\n')
+            Run.warn("You did not specify --partitions!")
+            Run.warn("Default computation chooses", args.partitions,
+                     "partitions (for {} embeddings)".format(num_embeddings))
+            print('\n\n')
+
+        index_faiss(args)
+
+
+if __name__ == "__main__":
+    main()

colbert/indexing/encoder.py

@@ -0,0 +1,187 @@
+import os
+import time
+import torch
+import ujson
+import numpy as np
+
+import itertools
+import threading
+import queue
+
+from colbert.modeling.inference import ModelInference
+from colbert.evaluation.loaders import load_colbert
+from colbert.utils.utils import print_message
+
+from colbert.indexing.index_manager import IndexManager
+
+
+class CollectionEncoder():
+    def __init__(self, args, process_idx, num_processes):
+        self.args = args
+        self.collection = args.collection
+        self.process_idx = process_idx
+        self.num_processes = num_processes
+
+        assert 0.5 <= args.chunksize <= 128.0
+        max_bytes_per_file = args.chunksize * (1024*1024*1024)
+
+        max_bytes_per_doc = (self.args.doc_maxlen * self.args.dim * 2.0)
+
+        # Determine subset sizes for output
+        minimum_subset_size = 10_000
+        maximum_subset_size = max_bytes_per_file / max_bytes_per_doc
+        maximum_subset_size = max(minimum_subset_size, maximum_subset_size)
+        self.possible_subset_sizes = [int(maximum_subset_size)]
+
+        self.print_main("#> Local args.bsize =", args.bsize)
+        self.print_main("#> args.index_root =", args.index_root)
+        self.print_main(f"#> self.possible_subset_sizes = {self.possible_subset_sizes}")
+
+        self._load_model()
+        self.indexmgr = IndexManager(args.dim)
+        self.iterator = self._initialize_iterator()
+
+    def _initialize_iterator(self):
+        return open(self.collection)
+
+    def _saver_thread(self):
+        for args in iter(self.saver_queue.get, None):
+            self._save_batch(*args)
+
+    def _load_model(self):
+        self.colbert, self.checkpoint = load_colbert(self.args, do_print=(self.process_idx == 0))
+        self.colbert = self.colbert.cuda()
+        self.colbert.eval()
+
+        self.inference = ModelInference(self.colbert, amp=self.args.amp)
+
+    def encode(self):
+        self.saver_queue = queue.Queue(maxsize=3)
+        thread = threading.Thread(target=self._saver_thread)
+        thread.start()
+
+        t0 = time.time()
+        local_docs_processed = 0
+
+        for batch_idx, (offset, lines, owner) in enumerate(self._batch_passages(self.iterator)):
+            if owner != self.process_idx:
+                continue
+
+            t1 = time.time()
+            batch = self._preprocess_batch(offset, lines)
+            embs, doclens = self._encode_batch(batch_idx, batch)
+
+            t2 = time.time()
+            self.saver_queue.put((batch_idx, embs, offset, doclens))
+
+            t3 = time.time()
+            local_docs_processed += len(lines)
+            overall_throughput = compute_throughput(local_docs_processed, t0, t3)
+            this_encoding_throughput = compute_throughput(len(lines), t1, t2)
+            this_saving_throughput = compute_throughput(len(lines), t2, t3)
+
+            self.print(f'#> Completed batch #{batch_idx} (starting at passage #{offset}) \t\t'
+                          f'Passages/min: {overall_throughput} (overall), ',
+                          f'{this_encoding_throughput} (this encoding), ',
+                          f'{this_saving_throughput} (this saving)')
+        self.saver_queue.put(None)
+
+        self.print("#> Joining saver thread.")
+        thread.join()
+
+    def _batch_passages(self, fi):
+        """
+        Must use the same seed across processes!
+        """
+        np.random.seed(0)
+
+        offset = 0
+        for owner in itertools.cycle(range(self.num_processes)):
+            batch_size = np.random.choice(self.possible_subset_sizes)
+
+            L = [line for _, line in zip(range(batch_size), fi)]
+
+            if len(L) == 0:
+                break  # EOF
+
+            yield (offset, L, owner)
+            offset += len(L)
+
+            if len(L) < batch_size:
+                break  # EOF
+
+        self.print("[NOTE] Done with local share.")
+
+        return
+
+    def _preprocess_batch(self, offset, lines):
+        endpos = offset + len(lines)
+
+        batch = []
+
+        for line_idx, line in zip(range(offset, endpos), lines):
+            line_parts = line.strip().split('\t')
+
+            pid, passage, *other = line_parts
+
+            assert len(passage) >= 1
+
+            if len(other) >= 1:
+                title, *_ = other
+                passage = title + ' | ' + passage
+
+            batch.append(passage)
+
+            # assert pid == 'id' or int(pid) == line_idx
+
+        return batch
+
+    def _encode_batch(self, batch_idx, batch):
+        with torch.no_grad():
+            embs = self.inference.docFromText(batch, bsize=self.args.bsize, keep_dims=False)
+            assert type(embs) is list
+            assert len(embs) == len(batch)
+
+            local_doclens = [d.size(0) for d in embs]
+            embs = torch.cat(embs)
+
+        return embs, local_doclens
+
+    def _save_batch(self, batch_idx, embs, offset, doclens):
+        start_time = time.time()
+
+        output_path = os.path.join(self.args.index_path, "{}.pt".format(batch_idx))
+        output_sample_path = os.path.join(self.args.index_path, "{}.sample".format(batch_idx))
+        doclens_path = os.path.join(self.args.index_path, 'doclens.{}.json'.format(batch_idx))
+
+        # Save the embeddings.
+        self.indexmgr.save(embs, output_path)
+        self.indexmgr.save(embs[torch.randint(0, high=embs.size(0), size=(embs.size(0) // 20,))], output_sample_path)
+
+        # Save the doclens.
+        with open(doclens_path, 'w') as output_doclens:
+            ujson.dump(doclens, output_doclens)
+
+        throughput = compute_throughput(len(doclens), start_time, time.time())
+        self.print_main("#> Saved batch #{} to {} \t\t".format(batch_idx, output_path),
+                        "Saving Throughput =", throughput, "passages per minute.\n")
+
+    def print(self, *args):
+        print_message("[" + str(self.process_idx) + "]", "\t\t", *args)
+
+    def print_main(self, *args):
+        if self.process_idx == 0:
+            self.print(*args)
+
+
+def compute_throughput(size, t0, t1):
+    throughput = size / (t1 - t0) * 60
+
+    if throughput > 1000 * 1000:
+        throughput = throughput / (1000*1000)
+        throughput = round(throughput, 1)
+        return '{}M'.format(throughput)
+
+    throughput = throughput / (1000)
+    throughput = round(throughput, 1)
+    return '{}k'.format(throughput)

colbert/indexing/faiss.py

@@ -0,0 +1,116 @@
+import os
+import math
+import faiss
+import torch
+import numpy as np
+
+import threading
+import queue
+
+from colbert.utils.utils import print_message, grouper
+from colbert.indexing.loaders import get_parts
+from colbert.indexing.index_manager import load_index_part
+from colbert.indexing.faiss_index import FaissIndex
+
+
+def get_faiss_index_name(args, offset=None, endpos=None):
+    partitions_info = '' if args.partitions is None else f'.{args.partitions}'
+    range_info = '' if offset is None else f'.{offset}-{endpos}'
+
+    return f'ivfpq{partitions_info}{range_info}.faiss'
+
+
+def load_sample(samples_paths, sample_fraction=None):
+    sample = []
+
+    for filename in samples_paths:
+        print_message(f"#> Loading {filename} ...")
+        part = load_index_part(filename)
+        if sample_fraction:
+            part = part[torch.randint(0, high=part.size(0), size=(int(part.size(0) * sample_fraction),))]
+        sample.append(part)
+
+    sample = torch.cat(sample).float().numpy()
+
+    print("#> Sample has shape", sample.shape)
+
+    return sample
+
+
+def prepare_faiss_index(slice_samples_paths, partitions, sample_fraction=None):
+    training_sample = load_sample(slice_samples_paths, sample_fraction=sample_fraction)
+
+    dim = training_sample.shape[-1]
+    index = FaissIndex(dim, partitions)
+
+    print_message("#> Training with the vectors...")
+
+    index.train(training_sample)
+
+    print_message("Done training!\n")
+
+    return index
+
+
+SPAN = 3
+
+
+def index_faiss(args):
+    print_message("#> Starting..")
+
+    parts, parts_paths, samples_paths = get_parts(args.index_path)
+
+    if args.sample is not None:
+        assert args.sample, args.sample
+        print_message(f"#> Training with {round(args.sample * 100.0, 1)}% of *all* embeddings (provided --sample).")
+        samples_paths = parts_paths
+
+    num_parts_per_slice = math.ceil(len(parts) / args.slices)
+
+    for slice_idx, part_offset in enumerate(range(0, len(parts), num_parts_per_slice)):
+        part_endpos = min(part_offset + num_parts_per_slice, len(parts))
+
+        slice_parts_paths = parts_paths[part_offset:part_endpos]
+        slice_samples_paths = samples_paths[part_offset:part_endpos]
+
+        if args.slices == 1:
+            faiss_index_name = get_faiss_index_name(args)
+        else:
+            faiss_index_name = get_faiss_index_name(args, offset=part_offset, endpos=part_endpos)
+
+        output_path = os.path.join(args.index_path, faiss_index_name)
+        print_message(f"#> Processing slice #{slice_idx+1} of {args.slices} (range {part_offset}..{part_endpos}).")
+        print_message(f"#> Will write to {output_path}.")
+
+        assert not os.path.exists(output_path), output_path
+
+        index = prepare_faiss_index(slice_samples_paths, args.partitions, args.sample)
+
+        loaded_parts = queue.Queue(maxsize=1)
+
+        def _loader_thread(thread_parts_paths):
+            for filenames in grouper(thread_parts_paths, SPAN, fillvalue=None):
+                sub_collection = [load_index_part(filename) for filename in filenames if filename is not None]
+                sub_collection = torch.cat(sub_collection)
+                sub_collection = sub_collection.float().numpy()
+                loaded_parts.put(sub_collection)
+
+        thread = threading.Thread(target=_loader_thread, args=(slice_parts_paths,))
+        thread.start()
+
+        print_message("#> Indexing the vectors...")
+
+        for filenames in grouper(slice_parts_paths, SPAN, fillvalue=None):
+            print_message("#> Loading", filenames, "(from queue)...")
+            sub_collection = loaded_parts.get()
+
+            print_message("#> Processing a sub_collection with shape", sub_collection.shape)
+            index.add(sub_collection)
+
+        print_message("Done indexing!")
+
+        index.save(output_path)
+
+        print_message(f"\n\nDone! All complete (for slice #{slice_idx+1} of {args.slices})!")
+
+        thread.join()

colbert/indexing/faiss_index.py

@@ -0,0 +1,58 @@
+import sys
+import time
+import math
+import faiss
+import torch
+
+import numpy as np
+
+from colbert.indexing.faiss_index_gpu import FaissIndexGPU
+from colbert.utils.utils import print_message
+
+
+class FaissIndex():
+    def __init__(self, dim, partitions):
+        self.dim = dim
+        self.partitions = partitions
+
+        self.gpu = FaissIndexGPU()
+        self.quantizer, self.index = self._create_index()
+        self.offset = 0
+
+    def _create_index(self):
+        quantizer = faiss.IndexFlatL2(self.dim)  # faiss.IndexHNSWFlat(dim, 32)
+        index = faiss.IndexIVFPQ(quantizer, self.dim, self.partitions, 16, 8)
+
+        return quantizer, index
+
+    def train(self, train_data):
+        print_message(f"#> Training now (using {self.gpu.ngpu} GPUs)...")
+
+        if self.gpu.ngpu > 0:
+            self.gpu.training_initialize(self.index, self.quantizer)
+
+        s = time.time()
+        self.index.train(train_data)
+        print(time.time() - s)
+
+        if self.gpu.ngpu > 0:
+            self.gpu.training_finalize()
+
+    def add(self, data):
+        print_message(f"Add data with shape {data.shape} (offset = {self.offset})..")
+
+        if self.gpu.ngpu > 0 and self.offset == 0:
+            self.gpu.adding_initialize(self.index)
+
+        if self.gpu.ngpu > 0:
+            self.gpu.add(self.index, data, self.offset)
+        else:
+            self.index.add(data)
+
+        self.offset += data.shape[0]
+
+    def save(self, output_path):
+        print_message(f"Writing index to {output_path} ...")
+
+        self.index.nprobe = 10  # just a default
+        faiss.write_index(self.index, output_path)

colbert/indexing/faiss_index_gpu.py

@@ -0,0 +1,138 @@
+"""
+    Heavily based on: https://github.com/facebookresearch/faiss/blob/master/benchs/bench_gpu_1bn.py
+"""
+
+
+import sys
+import time
+import math
+import faiss
+import torch
+
+import numpy as np
+from colbert.utils.utils import print_message
+
+
+class FaissIndexGPU():
+    def __init__(self):
+        self.ngpu = faiss.get_num_gpus()
+
+        if self.ngpu == 0:
+            return
+
+        self.tempmem = 1 << 33
+        self.max_add_per_gpu = 1 << 25
+        self.max_add = self.max_add_per_gpu * self.ngpu
+        self.add_batch_size = 65536
+
+        self.gpu_resources = self._prepare_gpu_resources()
+
+    def _prepare_gpu_resources(self):
+        print_message(f"Preparing resources for {self.ngpu} GPUs.")
+
+        gpu_resources = []
+
+        for _ in range(self.ngpu):
+            res = faiss.StandardGpuResources()
+            if self.tempmem >= 0:
+                res.setTempMemory(self.tempmem)
+            gpu_resources.append(res)
+
+        return gpu_resources
+
+    def _make_vres_vdev(self):
+        """
+        return vectors of device ids and resources useful for gpu_multiple
+        """
+
+        assert self.ngpu > 0
+
+        vres = faiss.GpuResourcesVector()
+        vdev = faiss.IntVector()
+
+        for i in range(self.ngpu):
+            vdev.push_back(i)
+            vres.push_back(self.gpu_resources[i])
+
+        return vres, vdev
+
+    def training_initialize(self, index, quantizer):
+        """
+        The index and quantizer should be owned by caller.
+        """
+
+        assert self.ngpu > 0
+
+        s = time.time()
+        self.index_ivf = faiss.extract_index_ivf(index)
+        self.clustering_index = faiss.index_cpu_to_all_gpus(quantizer)
+        self.index_ivf.clustering_index = self.clustering_index
+        print(time.time() - s)
+
+    def training_finalize(self):
+        assert self.ngpu > 0
+
+        s = time.time()
+        self.index_ivf.clustering_index = faiss.index_gpu_to_cpu(self.index_ivf.clustering_index)
+        print(time.time() - s)
+
+    def adding_initialize(self, index):
+        """
+        The index should be owned by caller.
+        """
+
+        assert self.ngpu > 0
+
+        self.co = faiss.GpuMultipleClonerOptions()
+        self.co.useFloat16 = True
+        self.co.useFloat16CoarseQuantizer = False
+        self.co.usePrecomputed = False
+        self.co.indicesOptions = faiss.INDICES_CPU
+        self.co.verbose = True
+        self.co.reserveVecs = self.max_add
+        self.co.shard = True
+        assert self.co.shard_type in (0, 1, 2)
+
+        self.vres, self.vdev = self._make_vres_vdev()
+        self.gpu_index = faiss.index_cpu_to_gpu_multiple(self.vres, self.vdev, index, self.co)
+
+    def add(self, index, data, offset):
+        assert self.ngpu > 0
+
+        t0 = time.time()
+        nb = data.shape[0]
+
+        for i0 in range(0, nb, self.add_batch_size):
+            i1 = min(i0 + self.add_batch_size, nb)
+            xs = data[i0:i1]
+
+            self.gpu_index.add_with_ids(xs, np.arange(offset+i0, offset+i1))
+
+            if self.max_add > 0 and self.gpu_index.ntotal > self.max_add:
+                self._flush_to_cpu(index, nb, offset)
+
+            print('\r%d/%d (%.3f s)  ' % (i0, nb, time.time() - t0), end=' ')
+            sys.stdout.flush()
+
+        if self.gpu_index.ntotal > 0:
+            self._flush_to_cpu(index, nb, offset)
+
+        assert index.ntotal == offset+nb, (index.ntotal, offset+nb, offset, nb)
+        print(f"add(.) time: %.3f s \t\t--\t\t index.ntotal = {index.ntotal}" % (time.time() - t0))
+
+    def _flush_to_cpu(self, index, nb, offset):
+        print("Flush indexes to CPU")
+
+        for i in range(self.ngpu):
+            index_src_gpu = faiss.downcast_index(self.gpu_index if self.ngpu == 1 else self.gpu_index.at(i))
+            index_src = faiss.index_gpu_to_cpu(index_src_gpu)
+
+            index_src.copy_subset_to(index, 0, offset, offset+nb)
+            index_src_gpu.reset()
+            index_src_gpu.reserveMemory(self.max_add)
+
+        if self.ngpu > 1:
+            try:
+                self.gpu_index.sync_with_shard_indexes()
+            except:
+                self.gpu_index.syncWithSubIndexes()

colbert/indexing/index_manager.py

@@ -0,0 +1,22 @@
+import torch
+import faiss
+import numpy as np
+
+from colbert.utils.utils import print_message
+
+
+class IndexManager():
+    def __init__(self, dim):
+        self.dim = dim
+
+    def save(self, tensor, path_prefix):
+        torch.save(tensor, path_prefix)
+
+
+def load_index_part(filename, verbose=True):
+    part = torch.load(filename)
+
+    if type(part) == list:  # for backward compatibility
+        part = torch.cat(part)
+
+    return part

colbert/indexing/loaders.py

@@ -0,0 +1,34 @@
+import os
+import torch
+import ujson
+
+from math import ceil
+from itertools import accumulate
+from colbert.utils.utils import print_message
+
+
+def get_parts(directory):
+    extension = '.pt'
+
+    parts = sorted([int(filename[: -1 * len(extension)]) for filename in os.listdir(directory)
+                    if filename.endswith(extension)])
+
+    assert list(range(len(parts))) == parts, parts
+
+    # Integer-sortedness matters.
+    parts_paths = [os.path.join(directory, '{}{}'.format(filename, extension)) for filename in parts]
+    samples_paths = [os.path.join(directory, '{}.sample'.format(filename)) for filename in parts]
+
+    return parts, parts_paths, samples_paths
+
+
+def load_doclens(directory, flatten=True):
+    parts, _, _ = get_parts(directory)
+
+    doclens_filenames = [os.path.join(directory, 'doclens.{}.json'.format(filename)) for filename in parts]
+    all_doclens = [ujson.load(open(filename)) for filename in doclens_filenames]
+
+    if flatten:
+        all_doclens = [x for sub_doclens in all_doclens for x in sub_doclens]
+
+    return all_doclens

colbert/modeling/colbert.py

@@ -0,0 +1,73 @@
+import string
+import torch
+import torch.nn as nn
+
+from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
+from colbert.parameters import DEVICE
+
+
+class ColBERT(BertPreTrainedModel):
+    def __init__(self, config, query_maxlen, doc_maxlen, mask_punctuation, dim=128, similarity_metric='cosine'):
+
+        super(ColBERT, self).__init__(config)
+
+        self.query_maxlen = query_maxlen
+        self.doc_maxlen = doc_maxlen
+        self.similarity_metric = similarity_metric
+        self.dim = dim
+
+        self.mask_punctuation = mask_punctuation
+        self.skiplist = {}
+
+        if self.mask_punctuation:
+            self.tokenizer = BertTokenizerFast.from_pretrained('bert-base-multilingual-uncased')
+            self.skiplist = {w: True
+                             for symbol in string.punctuation
+                             for w in [symbol, self.tokenizer.encode(symbol, add_special_tokens=False)[0]]}
+
+        self.bert = BertModel(config)
+        self.linear = nn.Linear(config.hidden_size, dim * 2, bias=False)
+
+        self.init_weights()
+
+    def forward(self, Q, D):
+        return self.score(self.query(*Q), self.doc(*D))
+
+    def query(self, input_ids, attention_mask):
+        input_ids, attention_mask = input_ids.to(DEVICE), attention_mask.to(DEVICE)
+        Q = self.bert(input_ids, attention_mask=attention_mask)[0]
+        Q = self.linear(Q)
+        Q = Q.split(int(Q.size(2)/2),2)
+        Q = torch.cat(Q,1)
+
+        return torch.nn.functional.normalize(Q, p=2, dim=2)
+
+    def doc(self, input_ids, attention_mask, keep_dims=True):
+        input_ids, attention_mask = input_ids.to(DEVICE), attention_mask.to(DEVICE)
+        D = self.bert(input_ids, attention_mask=attention_mask)[0]
+        D = self.linear(D)
+        D = D.split(int(D.size(2)/2),2)
+        D = torch.cat(D,1)
+
+        mask = torch.tensor(self.mask(input_ids), device=DEVICE).unsqueeze(2).float()
+        mask = torch.cat(2*[mask],1)
+        D = D * mask
+
+        D = torch.nn.functional.normalize(D, p=2, dim=2)
+
+        if not keep_dims:
+            D, mask = D.cpu().to(dtype=torch.float16), mask.cpu().bool().squeeze(-1)
+            D = [d[mask[idx]] for idx, d in enumerate(D)]
+
+        return D
+
+    def score(self, Q, D):
+        if self.similarity_metric == 'cosine':
+            return (Q @ D.permute(0, 2, 1)).max(2).values.sum(1)
+
+        assert self.similarity_metric == 'l2'
+        return (-1.0 * ((Q.unsqueeze(2) - D.unsqueeze(1))**2).sum(-1)).max(-1).values.sum(-1)
+
+    def mask(self, input_ids):
+        mask = [[(x not in self.skiplist) and (x != 0) for x in d] for d in input_ids.cpu().tolist()]
+        return mask

colbert/modeling/inference.py

@@ -0,0 +1,87 @@
+import torch
+
+from colbert.modeling.colbert import ColBERT
+from colbert.modeling.tokenization import QueryTokenizer, DocTokenizer
+from colbert.utils.amp import MixedPrecisionManager
+from colbert.parameters import DEVICE
+
+
+class ModelInference():
+    def __init__(self, colbert: ColBERT, amp=False):
+        assert colbert.training is False
+
+        self.colbert = colbert
+        self.query_tokenizer = QueryTokenizer(colbert.query_maxlen)
+        self.doc_tokenizer = DocTokenizer(colbert.doc_maxlen)
+
+        self.amp_manager = MixedPrecisionManager(amp)
+
+    def query(self, *args, to_cpu=False, **kw_args):
+        with torch.no_grad():
+            with self.amp_manager.context():
+                Q = self.colbert.query(*args, **kw_args)
+                return Q.cpu() if to_cpu else Q
+
+    def doc(self, *args, to_cpu=False, **kw_args):
+        with torch.no_grad():
+            with self.amp_manager.context():
+                D = self.colbert.doc(*args, **kw_args)
+                return D.cpu() if to_cpu else D
+
+    def queryFromText(self, queries, bsize=None, to_cpu=False):
+        if bsize:
+            batches = self.query_tokenizer.tensorize(queries, bsize=bsize)
+            batches = [self.query(input_ids, attention_mask, to_cpu=to_cpu) for input_ids, attention_mask in batches]
+            return torch.cat(batches)
+
+        input_ids, attention_mask = self.query_tokenizer.tensorize(queries)
+        return self.query(input_ids, attention_mask)
+
+    def docFromText(self, docs, bsize=None, keep_dims=True, to_cpu=False):
+        if bsize:
+            batches, reverse_indices = self.doc_tokenizer.tensorize(docs, bsize=bsize)
+
+            batches = [self.doc(input_ids, attention_mask, keep_dims=keep_dims, to_cpu=to_cpu)
+                       for input_ids, attention_mask in batches]
+
+            if keep_dims:
+                D = _stack_3D_tensors(batches)
+                return D[reverse_indices]
+
+            D = [d for batch in batches for d in batch]
+            return [D[idx] for idx in reverse_indices.tolist()]
+
+        input_ids, attention_mask = self.doc_tokenizer.tensorize(docs)
+        return self.doc(input_ids, attention_mask, keep_dims=keep_dims)
+
+    def score(self, Q, D, mask=None, lengths=None, explain=False):
+        if lengths is not None:
+            assert mask is None, "don't supply both mask and lengths"
+
+            mask = torch.arange(D.size(1), device=DEVICE) + 1
+            mask = mask.unsqueeze(0) <= lengths.to(DEVICE).unsqueeze(-1)
+
+        scores = (D @ Q)
+        scores = scores if mask is None else scores * mask.unsqueeze(-1)
+        scores = scores.max(1)
+
+        if explain:
+            assert False, "TODO"
+
+        return scores.values.sum(-1).cpu()
+
+
+def _stack_3D_tensors(groups):
+    bsize = sum([x.size(0) for x in groups])
+    maxlen = max([x.size(1) for x in groups])
+    hdim = groups[0].size(2)
+
+    output = torch.zeros(bsize, maxlen, hdim, device=groups[0].device, dtype=groups[0].dtype)
+
+    offset = 0
+    for x in groups:
+        endpos = offset + x.size(0)
+        output[offset:endpos, :x.size(1)] = x
+        offset = endpos
+
+    return output

colbert/modeling/tokenization/__init__.py

@@ -0,0 +1,3 @@
+from colbert.modeling.tokenization.query_tokenization import *
+from colbert.modeling.tokenization.doc_tokenization import *
+from colbert.modeling.tokenization.utils import tensorize_triples

colbert/modeling/tokenization/doc_tokenization.py

@@ -0,0 +1,63 @@
+import torch
+
+from transformers import BertTokenizerFast
+from colbert.modeling.tokenization.utils import _split_into_batches, _sort_by_length
+
+
+class DocTokenizer():
+    def __init__(self, doc_maxlen):
+        self.tok = BertTokenizerFast.from_pretrained('bert-base-multilingual-uncased')
+        self.doc_maxlen = doc_maxlen
+
+        self.D_marker_token, self.D_marker_token_id = '[D]', self.tok.convert_tokens_to_ids('[unused1]')
+        self.cls_token, self.cls_token_id = self.tok.cls_token, self.tok.cls_token_id
+        self.sep_token, self.sep_token_id = self.tok.sep_token, self.tok.sep_token_id
+
+        assert self.D_marker_token_id == 1
+
+    def tokenize(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        tokens = [self.tok.tokenize(x, add_special_tokens=False) for x in batch_text]
+
+        if not add_special_tokens:
+            return tokens
+
+        prefix, suffix = [self.cls_token, self.D_marker_token], [self.sep_token]
+        tokens = [prefix + lst + suffix for lst in tokens]
+
+        return tokens
+
+    def encode(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        ids = self.tok(batch_text, add_special_tokens=False)['input_ids']
+
+        if not add_special_tokens:
+            return ids
+
+        prefix, suffix = [self.cls_token_id, self.D_marker_token_id], [self.sep_token_id]
+        ids = [prefix + lst + suffix for lst in ids]
+
+        return ids
+
+    def tensorize(self, batch_text, bsize=None):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        # add placehold for the [D] marker
+        batch_text = ['. ' + x for x in batch_text]
+
+        obj = self.tok(batch_text, padding='longest', truncation='longest_first',
+                       return_tensors='pt', max_length=self.doc_maxlen)
+
+        ids, mask = obj['input_ids'], obj['attention_mask']
+
+        # postprocess for the [D] marker
+        ids[:, 1] = self.D_marker_token_id
+
+        if bsize:
+            ids, mask, reverse_indices = _sort_by_length(ids, mask, bsize)
+            batches = _split_into_batches(ids, mask, bsize)
+            return batches, reverse_indices
+
+        return ids, mask

colbert/modeling/tokenization/query_tokenization.py

@@ -0,0 +1,64 @@
+import torch
+
+from transformers import BertTokenizerFast
+from colbert.modeling.tokenization.utils import _split_into_batches
+
+
+class QueryTokenizer():
+    def __init__(self, query_maxlen):
+        self.tok = BertTokenizerFast.from_pretrained('bert-base-multilingual-uncased')
+        self.query_maxlen = query_maxlen
+
+        self.Q_marker_token, self.Q_marker_token_id = '[Q]', self.tok.convert_tokens_to_ids('[unused0]')
+        self.cls_token, self.cls_token_id = self.tok.cls_token, self.tok.cls_token_id
+        self.sep_token, self.sep_token_id = self.tok.sep_token, self.tok.sep_token_id
+        self.mask_token, self.mask_token_id = self.tok.mask_token, self.tok.mask_token_id
+
+        assert self.Q_marker_token_id == 100 and self.mask_token_id == 103
+
+    def tokenize(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        tokens = [self.tok.tokenize(x, add_special_tokens=False) for x in batch_text]
+
+        if not add_special_tokens:
+            return tokens
+
+        prefix, suffix = [self.cls_token, self.Q_marker_token], [self.sep_token]
+        tokens = [prefix + lst + suffix + [self.mask_token] * (self.query_maxlen - (len(lst)+3)) for lst in tokens]
+
+        return tokens
+
+    def encode(self, batch_text, add_special_tokens=False):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        ids = self.tok(batch_text, add_special_tokens=False)['input_ids']
+
+        if not add_special_tokens:
+            return ids
+
+        prefix, suffix = [self.cls_token_id, self.Q_marker_token_id], [self.sep_token_id]
+        ids = [prefix + lst + suffix + [self.mask_token_id] * (self.query_maxlen - (len(lst)+3)) for lst in ids]
+
+        return ids
+
+    def tensorize(self, batch_text, bsize=None):
+        assert type(batch_text) in [list, tuple], (type(batch_text))
+
+        # add placehold for the [Q] marker
+        batch_text = ['. ' + x for x in batch_text]
+
+        obj = self.tok(batch_text, padding='max_length', truncation=True,
+                       return_tensors='pt', max_length=self.query_maxlen)
+
+        ids, mask = obj['input_ids'], obj['attention_mask']
+
+        # postprocess for the [Q] marker and the [MASK] augmentation
+        ids[:, 1] = self.Q_marker_token_id
+        ids[ids == 0] = self.mask_token_id
+
+        if bsize:
+            batches = _split_into_batches(ids, mask, bsize)
+            return batches
+
+        return ids, mask

colbert/modeling/tokenization/utils.py

@@ -0,0 +1,51 @@
+import torch
+
+
+def tensorize_triples(query_tokenizer, doc_tokenizer, queries, positives, negatives, bsize):
+    assert len(queries) == len(positives) == len(negatives)
+    assert bsize is None or len(queries) % bsize == 0
+
+    N = len(queries)
+    Q_ids, Q_mask = query_tokenizer.tensorize(queries)
+    D_ids, D_mask = doc_tokenizer.tensorize(positives + negatives)
+    D_ids, D_mask = D_ids.view(2, N, -1), D_mask.view(2, N, -1)
+
+    # Compute max among {length of i^th positive, length of i^th negative} for i \in N
+    maxlens = D_mask.sum(-1).max(0).values
+
+    # Sort by maxlens
+    indices = maxlens.sort().indices
+    Q_ids, Q_mask = Q_ids[indices], Q_mask[indices]
+    D_ids, D_mask = D_ids[:, indices], D_mask[:, indices]
+
+    (positive_ids, negative_ids), (positive_mask, negative_mask) = D_ids, D_mask
+
+    query_batches = _split_into_batches(Q_ids, Q_mask, bsize)
+    positive_batches = _split_into_batches(positive_ids, positive_mask, bsize)
+    negative_batches = _split_into_batches(negative_ids, negative_mask, bsize)
+
+    batches = []
+    for (q_ids, q_mask), (p_ids, p_mask), (n_ids, n_mask) in zip(query_batches, positive_batches, negative_batches):
+        Q = (torch.cat((q_ids, q_ids)), torch.cat((q_mask, q_mask)))
+        D = (torch.cat((p_ids, n_ids)), torch.cat((p_mask, n_mask)))
+        batches.append((Q, D))
+
+    return batches
+
+
+def _sort_by_length(ids, mask, bsize):
+    if ids.size(0) <= bsize:
+        return ids, mask, torch.arange(ids.size(0))
+
+    indices = mask.sum(-1).sort().indices
+    reverse_indices = indices.sort().indices
+
+    return ids[indices], mask[indices], reverse_indices
+
+
+def _split_into_batches(ids, mask, bsize):
+    batches = []
+    for offset in range(0, ids.size(0), bsize):
+        batches.append((ids[offset:offset+bsize], mask[offset:offset+bsize]))
+
+    return batches