Spaces:
Runtime error
Runtime error
| 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. | |