'add_app_files'

app.py

@@ -0,0 +1,54 @@
+import random
+import gradio as gr
+from colbert.data import Queries
+from colbert.infra import Run, RunConfig, ColBERTConfig
+from colbert import Searcher
+
+
+# def init():
+searcher = None
+with Run().context(RunConfig(nranks=1, experiment="medqa")):
+    config = ColBERTConfig(
+        root="./experiments",
+    )
+    searcher = Searcher(index="medqa_idx", config=config)
+    
+
+        
+def search(query):
+    results = searcher.search(query, k=5)
+    responses=[]
+    # idx = 0 
+    for passage_id, _, _ in zip(*results):
+          responses.append(searcher.collection[passage_id])
+        #   idx = idx+1
+    return responses
+
+
+
+
+
+def chat(question):
+    # history = history or []
+    # message = message.lower()
+
+    # if message.startswith("how many"):
+    #     response = random.randint(1, 10)
+    # elif message.startswith("how"):
+    #     response = random.choice(["Great", "Good", "Okay", "Bad"])
+    # elif message.startswith("where"):
+    #     response = random.choice(["Here", "There", "Somewhere"])
+    # else:
+    #     response = "I don't know"
+    responses = search(question)
+    # history.append((message, response))
+    return responses
+
+chatbot = gr.Chatbot().style(color_map=("green", "pink"))
+demo = gr.Interface(
+    chat,
+    inputs=gr.Textbox(lines=2, placeholder="输入你的问题"),
+    outputs =["text", "text","text","text","text"]
+)
+if __name__ == "__main__":
+    demo.launch(share=True)

baleen/condenser/condense.py

@@ -0,0 +1,141 @@
+import torch
+
+from colbert.utils.utils import load_checkpoint
+from colbert.utils.amp import MixedPrecisionManager
+from colbert.utils.utils import flatten
+
+from baleen.utils.loaders import *
+from baleen.condenser.model import ElectraReader
+from baleen.condenser.tokenization import AnswerAwareTokenizer
+
+
+
+class Condenser:
+    def __init__(self, collectionX_path, checkpointL1, checkpointL2, deviceL1='cuda', deviceL2='cuda'):
+        self.modelL1, self.maxlenL1 = self._load_model(checkpointL1, deviceL1)
+        self.modelL2, self.maxlenL2 = self._load_model(checkpointL2, deviceL2)
+
+        assert self.maxlenL1 == self.maxlenL2, "Add support for different maxlens: use two tokenizers."
+
+        self.amp, self.tokenizer = self._setup_inference(self.maxlenL2)
+        self.CollectionX, self.CollectionY = self._load_collection(collectionX_path)
+
+    def condense(self, query, backs, ranking):
+        stage1_preds = self._stage1(query, backs, ranking)
+        stage2_preds, stage2_preds_L3x = self._stage2(query, stage1_preds)
+
+        return stage1_preds, stage2_preds, stage2_preds_L3x
+
+    def _load_model(self, path, device):
+        model = torch.load(path, map_location='cpu')
+        ElectraModels = ['google/electra-base-discriminator', 'google/electra-large-discriminator']
+        assert model['arguments']['model'] in ElectraModels, model['arguments']
+
+        model = ElectraReader.from_pretrained(model['arguments']['model'])
+        checkpoint = load_checkpoint(path, model)
+
+        model = model.to(device)
+        model.eval()
+
+        maxlen = checkpoint['arguments']['maxlen']
+
+        return model, maxlen
+    
+    def _setup_inference(self, maxlen):
+        amp = MixedPrecisionManager(activated=True)
+        tokenizer = AnswerAwareTokenizer(total_maxlen=maxlen)
+
+        return amp, tokenizer
+    
+    def _load_collection(self, collectionX_path):
+        CollectionX = {}
+        CollectionY = {}
+
+        with open(collectionX_path) as f:
+            for line_idx, line in enumerate(f):
+                line = ujson.loads(line)
+
+                assert type(line['text']) is list
+                assert line['pid'] == line_idx, (line_idx, line)
+
+                passage = [line['title']] + line['text']
+                CollectionX[line_idx] = passage
+
+                passage = [line['title'] + ' | ' + sentence for sentence in line['text']]
+
+                for idx, sentence in enumerate(passage):
+                    CollectionY[(line_idx, idx)] = sentence
+        
+        return CollectionX, CollectionY
+    
+    def _stage1(self, query, BACKS, ranking, TOPK=9):
+        model = self.modelL1
+
+        with torch.inference_mode():
+            backs = [self.CollectionY[(pid, sid)] for pid, sid in BACKS if (pid, sid) in self.CollectionY]
+            backs = [query] + backs
+            query = ' # '.join(backs)
+
+            # print(query)
+            # print(backs)
+            passages = []
+            actual_ranking = []
+
+            for pid in ranking:
+                actual_ranking.append(pid)
+                psg = self.CollectionX[pid]
+                psg = ' [MASK] '.join(psg)
+
+                passages.append(psg)
+
+            obj = self.tokenizer.process([query], passages, None)
+
+            with self.amp.context():
+                scores = model(obj.encoding.to(model.device)).float()
+
+            pids = [[pid] * scores.size(1) for pid in actual_ranking]
+            pids = flatten(pids)
+
+            sids = [list(range(scores.size(1))) for pid in actual_ranking]
+            sids = flatten(sids)
+
+            scores = scores.view(-1)
+
+            topk = scores.topk(min(TOPK, len(scores))).indices.tolist()
+            topk_pids = [pids[idx] for idx in topk]
+            topk_sids = [sids[idx] for idx in topk]
+
+            preds = [(pid, sid) for pid, sid in zip(topk_pids, topk_sids)]
+
+            pred_plus = BACKS + preds
+            pred_plus = f7(list(map(tuple, pred_plus)))[:TOPK]
+
+        return pred_plus
+    
+    def _stage2(self, query, preds):
+        model = self.modelL2
+
+        psgX = [self.CollectionY[(pid, sid)] for pid, sid in preds if (pid, sid) in self.CollectionY]
+        psg = ' [MASK] '.join([''] + psgX)
+        passages = [psg]
+        # print(passages)
+
+        obj = self.tokenizer.process([query], passages, None)
+
+        with self.amp.context():
+            scores = model(obj.encoding.to(model.device)).float()
+            scores = scores.view(-1).tolist()
+
+            preds = [(score, (pid, sid)) for (pid, sid), score in zip(preds, scores)]
+            preds = sorted(preds, reverse=True)[:5]
+
+            preds_L3x = [x for score, x in preds if score > min(0, preds[1][0] - 1e-10)] # Take at least 2!
+            preds = [x for score, x in preds if score > 0]
+
+            earliest_pids = f7([pid for pid, _ in preds_L3x])[:4]  # Take at most 4 docs.
+            preds_L3x = [(pid, sid) for pid, sid in preds_L3x if pid in earliest_pids]
+
+            assert len(preds_L3x) >= 2
+            assert len(f7([pid for pid, _ in preds_L3x])) <= 4
+
+        return preds, preds_L3x

baleen/condenser/model.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn as nn
+
+from transformers import ElectraPreTrainedModel, ElectraModel
+
+
+class ElectraReader(ElectraPreTrainedModel):
+    def __init__(self, config, learn_labels=False):
+        super(ElectraReader, self).__init__(config)
+
+        self.electra = ElectraModel(config)
+
+        self.relevance = nn.Linear(config.hidden_size, 1)
+
+        if learn_labels:
+            self.linear = nn.Linear(config.hidden_size, 2)
+        else:
+            self.linear = nn.Linear(config.hidden_size, 1)
+
+        self.init_weights()
+
+        self.learn_labels = learn_labels
+
+    def forward(self, encoding):
+        outputs = self.electra(encoding.input_ids,
+                               attention_mask=encoding.attention_mask,
+                               token_type_ids=encoding.token_type_ids)[0]
+
+        scores = self.linear(outputs)
+
+        if self.learn_labels:
+            scores = scores[:, 0].squeeze(1)
+        else:
+            scores = scores.squeeze(-1)
+            candidates = (encoding.input_ids == 103)
+            scores = self._mask_2d_index(scores, candidates)
+
+        return scores
+
+    def _mask_2d_index(self, scores, mask):
+        bsize, maxlen = scores.size()
+        bsize_, maxlen_ = mask.size()
+
+        assert bsize == bsize_, (scores.size(), mask.size())
+        assert maxlen == maxlen_, (scores.size(), mask.size())
+
+        # Get flat scores corresponding to the True mask positions, with -inf at the end
+        flat_scores = scores[mask]
+        flat_scores = torch.cat((flat_scores, torch.ones(1, device=self.device) * float('-inf')))
+
+        # Get 2D indexes
+        rowidxs, nnzs = torch.unique(torch.nonzero(mask, as_tuple=False)[:, 0], return_counts=True)
+        max_nnzs = nnzs.max().item()
+
+        rows = [[-1] * max_nnzs for _ in range(bsize)]
+        offset = 0
+        for rowidx, nnz in zip(rowidxs.tolist(), nnzs.tolist()):
+            rows[rowidx] = [offset + i for i in range(nnz)]
+            rows[rowidx] += [-1] * (max_nnzs - len(rows[rowidx]))
+            offset += nnz
+
+        indexes = torch.tensor(rows).to(self.device)
+
+        # Index with the 2D indexes
+        scores_2d = flat_scores[indexes]
+
+        return scores_2d
+
+    def _2d_index(self, embeddings, positions):
+        bsize, maxlen, hdim = embeddings.size()
+        bsize_, max_out = positions.size()
+
+        assert bsize == bsize_
+        assert positions.max() < maxlen
+
+        embeddings = embeddings.view(bsize * maxlen, hdim)
+        positions = positions + torch.arange(bsize, device=positions.device).unsqueeze(-1) * maxlen
+
+        return embeddings[positions]

baleen/condenser/tokenization.py

@@ -0,0 +1,118 @@
+import torch
+
+from transformers import ElectraTokenizerFast
+
+class AnswerAwareTokenizer():
+    def __init__(self, total_maxlen, bert_model='google/electra-base-discriminator'):
+        self.total_maxlen = total_maxlen
+
+        self.tok = ElectraTokenizerFast.from_pretrained(bert_model)
+
+    def process(self, questions, passages, all_answers=None, mask=None):
+        return TokenizationObject(self, questions, passages, all_answers, mask)
+
+    def tensorize(self, questions, passages):
+        query_lengths = self.tok(questions, padding='longest', return_tensors='pt').attention_mask.sum(-1)
+
+        encoding = self.tok(questions, passages, padding='longest', truncation='longest_first',
+                            return_tensors='pt', max_length=self.total_maxlen, add_special_tokens=True)
+
+        return encoding, query_lengths
+
+    def get_all_candidates(self, encoding, index):
+        offsets, endpositions = self.all_word_positions(encoding, index)
+
+        candidates = [(offset, endpos)
+                      for idx, offset in enumerate(offsets)
+                      for endpos in endpositions[idx:idx+10]]
+
+        return candidates
+
+    def all_word_positions(self, encoding, index):
+        words = encoding.word_ids(index)
+        offsets = [position
+                   for position, (last_word_number, current_word_number) in enumerate(zip([-1] + words, words))
+                   if last_word_number != current_word_number]
+
+        endpositions = offsets[1:] + [len(words)]
+
+        return offsets, endpositions
+
+    def characters_to_tokens(self, text, answers, encoding, index, offset, endpos):
+        # print(text, answers, encoding, index, offset, endpos)
+        # endpos = endpos - 1
+
+        for offset_ in range(offset, len(text)+1):
+            tokens_offset = encoding.char_to_token(index, offset_)
+            # print(f'tokens_offset = {tokens_offset}')
+            if tokens_offset is not None:
+                break
+
+        for endpos_ in range(endpos, len(text)+1):
+            tokens_endpos = encoding.char_to_token(index, endpos_)
+            # print(f'tokens_endpos = {tokens_endpos}')
+            if tokens_endpos is not None:
+                break
+
+        # None on whitespace!
+        assert tokens_offset is not None, (text, answers, offset)
+        # assert tokens_endpos is not None, (text, answers, endpos)
+        tokens_endpos = tokens_endpos if tokens_endpos is not None else len(encoding.tokens(index))
+
+        return tokens_offset, tokens_endpos
+
+    def tokens_to_answer(self, encoding, index, text, tokens_offset, tokens_endpos):
+        # print(encoding, index, text, tokens_offset, tokens_endpos, len(encoding.tokens(index)))
+
+        char_offset = encoding.word_to_chars(index, encoding.token_to_word(index, tokens_offset)).start
+
+        try:
+            char_next_offset = encoding.word_to_chars(index, encoding.token_to_word(index, tokens_endpos)).start
+            char_endpos = char_next_offset
+        except:
+            char_endpos = encoding.word_to_chars(index, encoding.token_to_word(index, tokens_endpos-1)).end
+
+        assert char_offset is not None
+        assert char_endpos is not None
+
+        return text[char_offset:char_endpos].strip()
+
+
+class TokenizationObject():
+    def __init__(self, tokenizer: AnswerAwareTokenizer, questions, passages, answers=None, mask=None):
+        assert type(questions) is list and type(passages) is list
+        assert len(questions) in [1, len(passages)]
+
+        if mask is None:
+            mask = [True for _ in passages]
+
+        self.mask = mask
+
+        self.tok = tokenizer
+        self.questions = questions if len(questions) == len(passages) else questions * len(passages)
+        self.passages = passages
+        self.answers = answers
+
+        self.encoding, self.query_lengths = self._encode()
+        self.passages_only_encoding, self.candidates, self.candidates_list = self._candidize()
+
+        if answers is not None:
+            self.gold_candidates = self.answers  # self._answerize()
+
+    def _encode(self):
+        return self.tok.tensorize(self.questions, self.passages)
+
+    def _candidize(self):
+        encoding = self.tok.tok(self.passages, add_special_tokens=False)
+
+        all_candidates = [self.tok.get_all_candidates(encoding, index) for index in range(len(self.passages))]
+
+        bsize, maxcands = len(self.passages), max(map(len, all_candidates))
+        all_candidates = [cands + [(-1, -1)] * (maxcands - len(cands)) for cands in all_candidates]
+
+        candidates = torch.tensor(all_candidates)
+        assert candidates.size() == (bsize, maxcands, 2), (candidates.size(), (bsize, maxcands, 2), (self.questions, self.passages))
+
+        candidates = candidates + self.query_lengths.unsqueeze(-1).unsqueeze(-1)
+
+        return encoding, candidates, all_candidates

baleen/engine.py

@@ -0,0 +1,58 @@
+from baleen.utils.loaders import *
+from baleen.condenser.condense import Condenser
+
+
+class Baleen:
+    def __init__(self, collectionX_path: str, searcher, condenser: Condenser):
+        self.collectionX = load_collectionX(collectionX_path)
+        self.searcher = searcher
+        self.condenser = condenser
+
+    def search(self, query, num_hops, depth=100, verbose=False):
+        assert depth % num_hops == 0, f"depth={depth} must be divisible by num_hops={num_hops}."
+        k = depth // num_hops
+
+        searcher = self.searcher
+        condenser = self.condenser
+        collectionX = self.collectionX
+
+        facts = []
+        stage1_preds = None
+        context = None
+
+        pids_bag = set()
+
+        for hop_idx in range(0, num_hops):
+            ranking = list(zip(*searcher.search(query, context=context, k=depth)))
+            ranking_ = []
+
+            facts_pids = set([pid for pid, _ in facts])
+
+            for pid, rank, score in ranking:
+                # print(f'[{score}] \t\t {searcher.collection[pid]}')
+                if len(ranking_) < k and pid not in facts_pids:
+                    ranking_.append(pid)
+                
+                if len(pids_bag) < k * (hop_idx+1):
+                    pids_bag.add(pid)
+            
+            stage1_preds, facts, stage2_L3x = condenser.condense(query, backs=facts, ranking=ranking_)
+            context = ' [SEP] '.join([collectionX.get((pid, sid), '') for pid, sid in facts])
+
+        assert len(pids_bag) == depth
+
+        return stage2_L3x, pids_bag, stage1_preds
+
+            
+
+
+            
+                
+
+        
+            
+
+
+
+
+

baleen/hop_searcher.py

@@ -0,0 +1,40 @@
+from typing import Union
+
+from colbert import Searcher
+from colbert.data import Queries
+from colbert.infra.config import ColBERTConfig
+
+
+TextQueries = Union[str, 'list[str]', 'dict[int, str]', Queries]
+
+
+class HopSearcher(Searcher):
+    def __init__(self, *args, config=None, interaction='flipr', **kw_args):
+        defaults = ColBERTConfig(query_maxlen=64, interaction=interaction)
+        config = ColBERTConfig.from_existing(defaults, config)
+
+        super().__init__(*args, config=config, **kw_args)
+
+    def encode(self, text: TextQueries, context: TextQueries):
+        queries = text if type(text) is list else [text]
+        context = context if context is None or type(context) is list else [context]
+        bsize = 128 if len(queries) > 128 else None
+
+        self.checkpoint.query_tokenizer.query_maxlen = self.config.query_maxlen
+        Q = self.checkpoint.queryFromText(queries, context=context, bsize=bsize, to_cpu=True)
+
+        return Q
+
+    def search(self, text: str, context: str, k=10):
+        return self.dense_search(self.encode(text, context), k)
+
+    def search_all(self, queries: TextQueries, context: TextQueries, k=10):
+        queries = Queries.cast(queries)
+        context = Queries.cast(context) if context is not None else context
+
+        queries_ = list(queries.values())
+        context_ = list(context.values()) if context is not None else context
+
+        Q = self.encode(queries_, context_)
+
+        return self._search_all_Q(queries, Q, k)

baleen/utils/annotate.py

@@ -0,0 +1,37 @@
+import os
+import ujson
+
+from colbert.utils.utils import print_message, file_tqdm
+
+
+def annotate_to_file(qas_path, ranking_path):
+    output_path = f'{ranking_path}.annotated'
+    assert not os.path.exists(output_path), output_path
+
+    QID2pids = {}
+
+    with open(qas_path) as f:
+        print_message(f"#> Reading QAs from {f.name} ..")
+
+        for line in file_tqdm(f):
+            example = ujson.loads(line)
+            QID2pids[example['qid']] = example['support_pids']
+
+    with open(ranking_path) as f:
+        print_message(f"#> Reading ranked lists from {f.name} ..")
+
+        with open(output_path, 'w') as g:
+            for line in file_tqdm(f):
+                qid, pid, *other = line.strip().split('\t')
+                qid, pid = map(int, [qid, pid])
+
+                label = int(pid in QID2pids[qid])
+
+                line_ = [qid, pid, *other, label]
+                line_ = '\t'.join(map(str, line_)) + '\n'
+                g.write(line_)
+
+    print_message(g.name)
+    print_message("#> Done!")
+
+    return g.name

baleen/utils/loaders.py

@@ -0,0 +1,50 @@
+import os
+import time
+import torch
+import ujson
+
+from colbert.utils.utils import f7, print_message, timestamp
+
+
+def load_contexts(first_hop_topk_path):
+    qid2backgrounds = {}
+
+    with open(first_hop_topk_path) as f:
+        print_message(f"#> Loading backgrounds from {f.name} ..")
+
+        last = None
+        for line in f:
+            qid, facts = ujson.loads(line)
+            facts = [(tuple(f) if type(f) is list else f) for f in facts]
+            qid2backgrounds[qid] = facts
+            last = (qid, facts)
+
+    # assert len(qid2backgrounds) in [0, len(queries)], (len(qid2backgrounds), len(queries))
+    print_message(f"#> {first_hop_topk_path} has {len(qid2backgrounds)} qids. Last = {last}")
+
+    return qid2backgrounds
+
+def load_collectionX(collection_path, dict_in_dict=False):
+    print_message("#> Loading collection...")
+
+    collectionX = {}
+
+    with open(collection_path) as f:
+        for line_idx, line in enumerate(f):
+            line = ujson.loads(line)
+
+            assert type(line['text']) is list
+            assert line['pid'] == line_idx, (line_idx, line)
+
+            passage = [line['title'] + ' | ' + sentence for sentence in line['text']]
+
+            if dict_in_dict:
+                collectionX[line_idx] = {}
+
+            for idx, sentence in enumerate(passage):
+                if dict_in_dict:
+                    collectionX[line_idx][idx] = sentence
+                else:
+                    collectionX[(line_idx, idx)] = sentence
+
+    return collectionX

colbert/__init__.py

@@ -0,0 +1,5 @@
+from .trainer import Trainer
+from .indexer import Indexer
+from .searcher import Searcher
+
+from .modeling.checkpoint import Checkpoint

colbert/data/__init__.py

@@ -0,0 +1,5 @@
+from .collection import *
+from .queries import *
+
+from .ranking import *
+from .examples import *

colbert/data/collection.py

@@ -0,0 +1,100 @@
+
+# Could be .tsv or .json. The latter always allows more customization via optional parameters.
+# I think it could be worth doing some kind of parallel reads too, if the file exceeds 1 GiBs.
+# Just need to use a datastructure that shares things across processes without too much pickling.
+# I think multiprocessing.Manager can do that!
+
+import os
+import itertools
+
+from colbert.evaluation.loaders import load_collection
+from colbert.infra.run import Run
+
+
+class Collection:
+    def __init__(self, path=None, data=None):
+        self.path = path
+        self.data = data or self._load_file(path)
+
+    def __iter__(self):
+        # TODO: If __data isn't there, stream from disk!
+        return self.data.__iter__()
+
+    def __getitem__(self, item):
+        # TODO: Load from disk the first time this is called. Unless self.data is already not None.
+        return self.data[item]
+
+    def __len__(self):
+        # TODO: Load here too. Basically, let's make data a property function and, on first call, either load or get __data.
+        return len(self.data)
+
+    def _load_file(self, path):
+        self.path = path
+        return self._load_tsv(path) if path.endswith('.tsv') else self._load_jsonl(path)
+
+    def _load_tsv(self, path):
+        return load_collection(path)
+
+    def _load_jsonl(self, path):
+        raise NotImplementedError()
+
+    def provenance(self):
+        return self.path
+    
+    def toDict(self):
+        return {'provenance': self.provenance()}
+
+    def save(self, new_path):
+        assert new_path.endswith('.tsv'), "TODO: Support .json[l] too."
+        assert not os.path.exists(new_path), new_path
+
+        with Run().open(new_path, 'w') as f:
+            # TODO: expects content to always be a string here; no separate title!
+            for pid, content in enumerate(self.data):
+                content = f'{pid}\t{content}\n'
+                f.write(content)
+            
+            return f.name
+
+    def enumerate(self, rank):
+        for _, offset, passages in self.enumerate_batches(rank=rank):
+            for idx, passage in enumerate(passages):
+                yield (offset + idx, passage)
+
+    def enumerate_batches(self, rank, chunksize=None):
+        assert rank is not None, "TODO: Add support for the rank=None case."
+
+        chunksize = chunksize or self.get_chunksize()
+
+        offset = 0
+        iterator = iter(self)
+
+        for chunk_idx, owner in enumerate(itertools.cycle(range(Run().nranks))):
+            L = [line for _, line in zip(range(chunksize), iterator)]
+
+            if len(L) > 0 and owner == rank:
+                yield (chunk_idx, offset, L)
+
+            offset += len(L)
+
+            if len(L) < chunksize:
+                return
+    
+    def get_chunksize(self):
+        return min(25_000, 1 + len(self) // Run().nranks)  # 25k is great, 10k allows things to reside on GPU??
+
+    @classmethod
+    def cast(cls, obj):
+        if type(obj) is str:
+            return cls(path=obj)
+
+        if type(obj) is list:
+            return cls(data=obj)
+
+        if type(obj) is cls:
+            return obj
+
+        assert False, f"obj has type {type(obj)} which is not compatible with cast()"
+
+
+# TODO: Look up path in some global [per-thread or thread-safe] list.

colbert/data/dataset.py

@@ -0,0 +1,14 @@
+
+
+# Not just the corpus, but also an arbitrary number of query sets, indexed by name in a dictionary/dotdict.
+# And also query sets with top-k PIDs.
+# QAs too? TripleSets too?
+
+
+class Dataset:
+    def __init__(self):
+        pass
+
+    def select(self, key):
+        # Select the {corpus, queryset, tripleset, rankingset} determined by uniqueness or by key and return a "unique" dataset (e.g., for key=train)
+        pass

colbert/data/examples.py

@@ -0,0 +1,82 @@
+from colbert.infra.run import Run
+import os
+import ujson
+
+from colbert.utils.utils import print_message
+from colbert.infra.provenance import Provenance
+from utility.utils.save_metadata import get_metadata_only
+
+
+class Examples:
+    def __init__(self, path=None, data=None, nway=None, provenance=None):
+        self.__provenance = provenance or path or Provenance()
+        self.nway = nway
+        self.path = path
+        self.data = data or self._load_file(path)
+
+    def provenance(self):
+        return self.__provenance
+    
+    def toDict(self):
+        return self.provenance()
+
+    def _load_file(self, path):
+        nway = self.nway + 1 if self.nway else self.nway
+        examples = []
+
+        with open(path) as f:
+            for line in f:
+                example = ujson.loads(line)[:nway]
+                examples.append(example)
+
+        return examples
+
+    def tolist(self, rank=None, nranks=None):
+        """
+        NOTE: For distributed sampling, this isn't equivalent to perfectly uniform sampling.
+        In particular, each subset is perfectly represented in every batch! However, since we never
+        repeat passes over the data, we never repeat any particular triple, and the split across
+        nodes is random (since the underlying file is pre-shuffled), there's no concern here.
+        """
+
+        if rank or nranks:
+            assert rank in range(nranks), (rank, nranks)
+            return [self.data[idx] for idx in range(0, len(self.data), nranks)]  # if line_idx % nranks == rank
+
+        return list(self.data)
+
+    def save(self, new_path):
+        assert 'json' in new_path.strip('/').split('/')[-1].split('.'), "TODO: Support .json[l] too."
+
+        print_message(f"#> Writing {len(self.data) / 1000_000.0}M examples to {new_path}")
+
+        with Run().open(new_path, 'w') as f:
+            for example in self.data:
+                ujson.dump(example, f)
+                f.write('\n')
+
+            output_path = f.name
+            print_message(f"#> Saved examples with {len(self.data)} lines to {f.name}")
+        
+        with Run().open(f'{new_path}.meta', 'w') as f:
+            d = {}
+            d['metadata'] = get_metadata_only()
+            d['provenance'] = self.provenance()
+            line = ujson.dumps(d, indent=4)
+            f.write(line)
+
+        return output_path
+
+    @classmethod
+    def cast(cls, obj, nway=None):
+        if type(obj) is str:
+            return cls(path=obj, nway=nway)
+
+        if isinstance(obj, list):
+            return cls(data=obj, nway=nway)
+
+        if type(obj) is cls:
+            assert nway is None, nway
+            return obj
+
+        assert False, f"obj has type {type(obj)} which is not compatible with cast()"

colbert/data/queries.py

@@ -0,0 +1,163 @@
+from colbert.infra.run import Run
+import os
+import ujson
+
+from colbert.evaluation.loaders import load_queries
+
+# TODO: Look up path in some global [per-thread or thread-safe] list.
+# TODO: path could be a list of paths...? But then how can we tell it's not a list of queries..
+
+
+class Queries:
+    def __init__(self, path=None, data=None):
+        self.path = path
+
+        if data:
+            assert isinstance(data, dict), type(data)
+        self._load_data(data) or self._load_file(path)
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __iter__(self):
+        return iter(self.data.items())
+
+    def provenance(self):
+        return self.path
+    
+    def toDict(self):
+        return {'provenance': self.provenance()}
+
+    def _load_data(self, data):
+        if data is None:
+            return None
+
+        self.data = {}
+        self._qas = {}
+
+        for qid, content in data.items():
+            if isinstance(content, dict):
+                self.data[qid] = content['question']
+                self._qas[qid] = content
+            else:
+                self.data[qid] = content
+
+        if len(self._qas) == 0:
+            del self._qas
+
+        return True
+
+    def _load_file(self, path):
+        if not path.endswith('.json'):
+            self.data = load_queries(path)
+            return True
+        
+        # Load QAs
+        self.data = {}
+        self._qas = {}
+
+        with open(path) as f:
+            for line in f:
+                qa = ujson.loads(line)
+
+                assert qa['qid'] not in self.data
+                self.data[qa['qid']] = qa['question']
+                self._qas[qa['qid']] = qa
+
+        return self.data
+
+    def qas(self):
+        return dict(self._qas)
+
+    def __getitem__(self, key):
+        return self.data[key]
+
+    def keys(self):
+        return self.data.keys()
+
+    def values(self):
+        return self.data.values()
+
+    def items(self):
+        return self.data.items()
+
+    def save(self, new_path):
+        assert new_path.endswith('.tsv')
+        assert not os.path.exists(new_path), new_path
+
+        with Run().open(new_path, 'w') as f:
+            for qid, content in self.data.items():
+                content = f'{qid}\t{content}\n'
+                f.write(content)
+            
+            return f.name
+
+    def save_qas(self, new_path):
+        assert new_path.endswith('.json')
+        assert not os.path.exists(new_path), new_path
+
+        with open(new_path, 'w') as f:
+            for qid, qa in self._qas.items():
+                qa['qid'] = qid
+                f.write(ujson.dumps(qa) + '\n')
+
+    def _load_tsv(self, path):
+        raise NotImplementedError
+
+    def _load_jsonl(self, path):
+        raise NotImplementedError
+
+    @classmethod
+    def cast(cls, obj):
+        if type(obj) is str:
+            return cls(path=obj)
+
+        if isinstance(obj, dict) or isinstance(obj, list):
+            return cls(data=obj)
+
+        if type(obj) is cls:
+            return obj
+
+        assert False, f"obj has type {type(obj)} which is not compatible with cast()"
+
+
+# class QuerySet:
+#     def __init__(self, *paths, renumber=False):
+#         self.paths = paths
+#         self.original_queries = [load_queries(path) for path in paths]
+
+#         if renumber:
+#             self.queries = flatten([q.values() for q in self.original_queries])
+#             self.queries = {idx: text for idx, text in enumerate(self.queries)}
+
+#         else:
+#             self.queries = {}
+
+#             for queries in self.original_queries:
+#                 assert len(set.intersection(set(queries.keys()), set(self.queries.keys()))) == 0, \
+#                     "renumber=False requires non-overlapping query IDs"
+
+#                 self.queries.update(queries)
+
+#         assert len(self.queries) == sum(map(len, self.original_queries))
+
+#     def todict(self):
+#         return dict(self.queries)
+
+#     def tolist(self):
+#         return list(self.queries.values())
+
+#     def query_sets(self):
+#         return self.original_queries
+
+#     def split_rankings(self, rankings):
+#         assert type(rankings) is list
+#         assert len(rankings) == len(self.queries)
+
+#         sub_rankings = []
+#         offset = 0
+#         for source in self.original_queries:
+#             sub_rankings.append(rankings[offset:offset+len(source)])
+#             offset += len(source)
+
+#         return sub_rankings

colbert/data/ranking.py

@@ -0,0 +1,94 @@
+import os
+import tqdm
+import ujson
+from colbert.infra.provenance import Provenance
+
+from colbert.infra.run import Run
+from colbert.utils.utils import print_message, groupby_first_item
+from utility.utils.save_metadata import get_metadata_only
+
+
+def numericize(v):
+    if '.' in v:
+        return float(v)
+
+    return int(v)
+
+
+def load_ranking(path):  # works with annotated and un-annotated ranked lists
+    print_message("#> Loading the ranked lists from", path)
+
+    with open(path) as f:
+        return [list(map(numericize, line.strip().split('\t'))) for line in f]
+
+
+class Ranking:
+    def __init__(self, path=None, data=None, metrics=None, provenance=None):
+        self.__provenance = provenance or path or Provenance()
+        self.data = self._prepare_data(data or self._load_file(path))
+
+    def provenance(self):
+        return self.__provenance
+    
+    def toDict(self):
+        return {'provenance': self.provenance()}
+
+    def _prepare_data(self, data):
+        # TODO: Handle list of lists???
+        if isinstance(data, dict):
+            self.flat_ranking = [(qid, *rest) for qid, subranking in data.items() for rest in subranking]
+            return data
+
+        self.flat_ranking = data
+        return groupby_first_item(tqdm.tqdm(self.flat_ranking))
+
+    def _load_file(self, path):
+        return load_ranking(path)
+
+    def todict(self):
+        return dict(self.data)
+
+    def tolist(self):
+        return list(self.flat_ranking)
+
+    def items(self):
+        return self.data.items()
+
+    def _load_tsv(self, path):
+        raise NotImplementedError
+
+    def _load_jsonl(self, path):
+        raise NotImplementedError
+
+    def save(self, new_path):
+        assert 'tsv' in new_path.strip('/').split('/')[-1].split('.'), "TODO: Support .json[l] too."
+
+        with Run().open(new_path, 'w') as f:
+            for items in self.flat_ranking:
+                line = '\t'.join(map(lambda x: str(int(x) if type(x) is bool else x), items)) + '\n'
+                f.write(line)
+
+            output_path = f.name
+            print_message(f"#> Saved ranking of {len(self.data)} queries and {len(self.flat_ranking)} lines to {f.name}")
+        
+        with Run().open(f'{new_path}.meta', 'w') as f:
+            d = {}
+            d['metadata'] = get_metadata_only()
+            d['provenance'] = self.provenance()
+            line = ujson.dumps(d, indent=4)
+            f.write(line)
+        
+        return output_path
+
+    @classmethod
+    def cast(cls, obj):
+        if type(obj) is str:
+            return cls(path=obj)
+
+        if isinstance(obj, dict) or isinstance(obj, list):
+            return cls(data=obj)
+
+        if type(obj) is cls:
+            return obj
+
+        assert False, f"obj has type {type(obj)} which is not compatible with cast()"

colbert/distillation/ranking_scorer.py

@@ -0,0 +1,52 @@
+import tqdm
+import ujson
+
+from collections import defaultdict
+
+from colbert.utils.utils import print_message, zipstar
+from utility.utils.save_metadata import get_metadata_only
+
+from colbert.infra import Run
+from colbert.data import Ranking
+from colbert.infra.provenance import Provenance
+from colbert.distillation.scorer import Scorer
+
+
+class RankingScorer:
+    def __init__(self, scorer: Scorer, ranking: Ranking):
+        self.scorer = scorer
+        self.ranking = ranking.tolist()
+        self.__provenance = Provenance()
+
+        print_message(f"#> Loaded ranking with {len(self.ranking)} qid--pid pairs!")
+    
+    def provenance(self):
+        return self.__provenance
+
+    def run(self):
+        print_message(f"#> Starting..")
+
+        qids, pids, *_ = zipstar(self.ranking)
+        distillation_scores = self.scorer.launch(qids, pids)
+
+        scores_by_qid = defaultdict(list)
+
+        for qid, pid, score in tqdm.tqdm(zip(qids, pids, distillation_scores)):
+            scores_by_qid[qid].append((score, pid))
+
+        with Run().open('distillation_scores.json', 'w') as f:
+            for qid in tqdm.tqdm(scores_by_qid):
+                obj = (qid, scores_by_qid[qid])
+                f.write(ujson.dumps(obj) + '\n')
+
+            output_path = f.name
+            print_message(f'#> Saved the distillation_scores to {output_path}')
+
+        with Run().open(f'{output_path}.meta', 'w') as f:
+            d = {}
+            d['metadata'] = get_metadata_only()
+            d['provenance'] = self.provenance()
+            line = ujson.dumps(d, indent=4)
+            f.write(line)
+
+        return output_path

colbert/distillation/scorer.py

@@ -0,0 +1,68 @@
+import torch
+import tqdm
+
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+from colbert.infra.launcher import Launcher
+from colbert.infra import Run, RunConfig
+from colbert.modeling.reranker.electra import ElectraReranker
+from colbert.utils.utils import flatten
+
+
+DEFAULT_MODEL = 'cross-encoder/ms-marco-MiniLM-L-6-v2'
+
+
+class Scorer:
+    def __init__(self, queries, collection, model=DEFAULT_MODEL, maxlen=180, bsize=256):
+        self.queries = queries
+        self.collection = collection
+        self.model = model
+
+        self.maxlen = maxlen
+        self.bsize = bsize
+
+    def launch(self, qids, pids):
+        launcher = Launcher(self._score_pairs_process, return_all=True)
+        outputs = launcher.launch(Run().config, qids, pids)
+
+        return flatten(outputs)
+
+    def _score_pairs_process(self, config, qids, pids):
+        assert len(qids) == len(pids), (len(qids), len(pids))
+        share = 1 + len(qids) // config.nranks
+        offset = config.rank * share
+        endpos = (1 + config.rank) * share
+
+        return self._score_pairs(qids[offset:endpos], pids[offset:endpos], show_progress=(config.rank < 1))
+
+    def _score_pairs(self, qids, pids, show_progress=False):
+        tokenizer = AutoTokenizer.from_pretrained(self.model)
+        model = AutoModelForSequenceClassification.from_pretrained(self.model).cuda()
+
+        assert len(qids) == len(pids), (len(qids), len(pids))
+
+        scores = []
+
+        model.eval()
+        with torch.inference_mode():
+            with torch.cuda.amp.autocast():
+                for offset in tqdm.tqdm(range(0, len(qids), self.bsize), disable=(not show_progress)):
+                    endpos = offset + self.bsize
+
+                    queries_ = [self.queries[qid] for qid in qids[offset:endpos]]
+                    passages_ = [self.collection[pid] for pid in pids[offset:endpos]]
+
+                    features = tokenizer(queries_, passages_, padding='longest', truncation=True,
+                                            return_tensors='pt', max_length=self.maxlen).to(model.device)
+
+                    scores.append(model(**features).logits.flatten())
+
+        scores = torch.cat(scores)
+        scores = scores.tolist()
+
+        Run().print(f'Returning with {len(scores)} scores')
+
+        return scores
+
+
+# LONG-TERM TODO: This can be sped up by sorting by length in advance.

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-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,198 @@
+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)
+    for qid in qrels:
+        qrels[qid] = list(set(qrels[qid]))
+
+    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('\n\r ').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/index.py

@@ -0,0 +1,17 @@
+
+
+# TODO: This is the loaded index, underneath a searcher.
+
+
+"""
+## Operations:
+
+index = Index(index='/path/to/index')
+index.load_to_memory()
+
+batch_of_pids = [2324,32432,98743,23432]
+index.lookup(batch_of_pids, device='cuda:0') -> (N, doc_maxlen, dim)
+
+index.iterate_over_parts()
+
+"""

colbert/indexer.py

@@ -0,0 +1,84 @@
+import os
+import time
+
+import torch.multiprocessing as mp
+
+from colbert.infra.run import Run
+from colbert.infra.config import ColBERTConfig, RunConfig
+from colbert.infra.launcher import Launcher
+
+from colbert.utils.utils import create_directory, print_message
+
+from colbert.indexing.collection_indexer import encode
+
+
+class Indexer:
+    def __init__(self, checkpoint, config=None):
+        """
+           Use Run().context() to choose the run's configuration. They are NOT extracted from `config`.
+        """
+
+        self.index_path = None
+        self.checkpoint = checkpoint
+        self.checkpoint_config = ColBERTConfig.load_from_checkpoint(checkpoint)
+
+        self.config = ColBERTConfig.from_existing(self.checkpoint_config, config, Run().config)
+        self.configure(checkpoint=checkpoint)
+
+    def configure(self, **kw_args):
+        self.config.configure(**kw_args)
+
+    def get_index(self):
+        return self.index_path
+
+    def erase(self):
+        assert self.index_path is not None
+        directory = self.index_path
+        deleted = []
+
+        for filename in sorted(os.listdir(directory)):
+            filename = os.path.join(directory, filename)
+
+            delete = filename.endswith(".json")
+            delete = delete and ('metadata' in filename or 'doclen' in filename or 'plan' in filename)
+            delete = delete or filename.endswith(".pt")
+            
+            if delete:
+                deleted.append(filename)
+        
+        if len(deleted):
+            print_message(f"#> Will delete {len(deleted)} files already at {directory} in 20 seconds...")
+            time.sleep(20)
+
+            for filename in deleted:
+                os.remove(filename)
+
+        return deleted
+
+    def index(self, name, collection, overwrite=False):
+        assert overwrite in [True, False, 'reuse', 'resume']
+
+        self.configure(collection=collection, index_name=name, resume=overwrite=='resume')
+        self.configure(bsize=64, partitions=None)
+
+        self.index_path = self.config.index_path_
+        index_does_not_exist = (not os.path.exists(self.config.index_path_))
+
+        assert (overwrite in [True, 'reuse', 'resume']) or index_does_not_exist, self.config.index_path_
+        create_directory(self.config.index_path_)
+
+        if overwrite is True:
+            self.erase()
+
+        if index_does_not_exist or overwrite != 'reuse':
+            self.__launch(collection)
+
+        return self.index_path
+
+    def __launch(self, collection):
+        manager = mp.Manager()
+        shared_lists = [manager.list() for _ in range(self.config.nranks)]
+        shared_queues = [manager.Queue(maxsize=1) for _ in range(self.config.nranks)]
+
+        launcher = Launcher(encode)
+        launcher.launch(self.config, collection, shared_lists, shared_queues)

colbert/indexing/codecs/decompress_residuals.cpp

@@ -0,0 +1,23 @@
+#include <torch/extension.h>
+
+torch::Tensor decompress_residuals_cuda(
+    const torch::Tensor binary_residuals, const torch::Tensor bucket_weights,
+    const torch::Tensor reversed_bit_map,
+    const torch::Tensor bucket_weight_combinations, const torch::Tensor codes,
+    const torch::Tensor centroids, const int dim, const int nbits);
+
+torch::Tensor decompress_residuals(
+    const torch::Tensor binary_residuals, const torch::Tensor bucket_weights,
+    const torch::Tensor reversed_bit_map,
+    const torch::Tensor bucket_weight_combinations, const torch::Tensor codes,
+    const torch::Tensor centroids, const int dim, const int nbits) {
+    // Add input verification
+    return decompress_residuals_cuda(
+        binary_residuals, bucket_weights, reversed_bit_map,
+        bucket_weight_combinations, codes, centroids, dim, nbits);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("decompress_residuals_cpp", &decompress_residuals,
+          "Decompress residuals");
+}

colbert/indexing/codecs/decompress_residuals.cu

@@ -0,0 +1,75 @@
+#include <assert.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <stdint.h>
+#include <torch/extension.h>
+
+__global__ void decompress_residuals_kernel(
+    const uint8_t* binary_residuals,
+    const torch::PackedTensorAccessor32<at::Half, 1, torch::RestrictPtrTraits>
+        bucket_weights,
+    const torch::PackedTensorAccessor32<uint8_t, 1, torch::RestrictPtrTraits>
+        reversed_bit_map,
+    const torch::PackedTensorAccessor32<uint8_t, 2, torch::RestrictPtrTraits>
+        bucket_weight_combinations,
+    const torch::PackedTensorAccessor32<int, 1, torch::RestrictPtrTraits> codes,
+    const torch::PackedTensorAccessor32<at::Half, 2, torch::RestrictPtrTraits>
+        centroids,
+    const int n, const int dim, const int nbits, const int packed_size,
+    at::Half* output) {
+    const int packed_dim = (int)(dim * nbits / packed_size);
+    const int i = blockIdx.x;
+    const int j = threadIdx.x;
+
+    if (i >= n) return;
+    if (j >= dim * nbits / packed_size) return;
+
+    const int code = codes[i];
+
+    uint8_t x = binary_residuals[i * packed_dim + j];
+    x = reversed_bit_map[x];
+    int output_idx = (int)(j * packed_size / nbits);
+    for (int k = 0; k < packed_size / nbits; k++) {
+        assert(output_idx < dim);
+        const int bucket_weight_idx = bucket_weight_combinations[x][k];
+        output[i * dim + output_idx] = bucket_weights[bucket_weight_idx];
+        output[i * dim + output_idx] += centroids[code][output_idx];
+        output_idx++;
+    }
+}
+
+torch::Tensor decompress_residuals_cuda(
+    const torch::Tensor binary_residuals, const torch::Tensor bucket_weights,
+    const torch::Tensor reversed_bit_map,
+    const torch::Tensor bucket_weight_combinations, const torch::Tensor codes,
+    const torch::Tensor centroids, const int dim, const int nbits) {
+    auto options = torch::TensorOptions()
+                       .dtype(torch::kFloat16)
+                       .device(torch::kCUDA, 0)
+                       .requires_grad(false);
+    torch::Tensor output =
+        torch::zeros({(int)binary_residuals.size(0), (int)dim}, options);
+
+    // TODO: Set this automatically?
+    const int packed_size = 8;
+
+    const int threads = dim / (packed_size / nbits);
+    const int blocks =
+        (binary_residuals.size(0) * binary_residuals.size(1)) / threads;
+
+    decompress_residuals_kernel<<<blocks, threads>>>(
+        binary_residuals.data<uint8_t>(),
+        bucket_weights
+            .packed_accessor32<at::Half, 1, torch::RestrictPtrTraits>(),
+        reversed_bit_map
+            .packed_accessor32<uint8_t, 1, torch::RestrictPtrTraits>(),
+        bucket_weight_combinations
+            .packed_accessor32<uint8_t, 2, torch::RestrictPtrTraits>(),
+        codes.packed_accessor32<int, 1, torch::RestrictPtrTraits>(),
+        centroids.packed_accessor32<at::Half, 2, torch::RestrictPtrTraits>(),
+        binary_residuals.size(0), dim, nbits, packed_size,
+        output.data<at::Half>());
+
+    return output;
+}