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