| import os |
| import torch |
|
|
| from tqdm import tqdm |
| from typing import Union |
|
|
| from colbert.data import Collection, Queries, Ranking |
|
|
| from colbert.modeling.checkpoint import Checkpoint |
| from colbert.search.index_storage import IndexScorer |
|
|
| from colbert.infra.provenance import Provenance |
| from colbert.infra.run import Run |
| from colbert.infra.config import ColBERTConfig, RunConfig |
| from colbert.infra.launcher import print_memory_stats |
|
|
| import time |
| from transformers import BertTokenizer, BertModel |
| import torch |
| import os |
| os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com' |
| os.environ['https_proxy'] = 'http://127.0.0.1:7890/' |
| os.environ['http_proxy'] = 'http://127.0.0.1:7890/' |
| myberttokenizer = BertTokenizer.from_pretrained('bert-base-uncased') |
|
|
|
|
| TextQueries = Union[str, 'list[str]', 'dict[int, str]', Queries] |
|
|
|
|
| class Searcher: |
| def __init__(self, index, checkpoint=None, collection=None, config=None): |
| print_memory_stats() |
|
|
| initial_config = ColBERTConfig.from_existing(config, Run().config) |
|
|
| default_index_root = initial_config.index_root_ |
| self.index = os.path.join(default_index_root, index) |
| self.index_config = ColBERTConfig.load_from_index(self.index) |
|
|
| self.checkpoint = checkpoint or self.index_config.checkpoint |
| self.checkpoint_config = ColBERTConfig.load_from_checkpoint(self.checkpoint) |
| self.config = ColBERTConfig.from_existing(self.checkpoint_config, self.index_config, initial_config) |
|
|
| self.collection = Collection.cast(collection or self.config.collection) |
| self.configure(checkpoint=self.checkpoint, collection=self.collection) |
|
|
| self.checkpoint = Checkpoint(self.checkpoint, colbert_config=self.config) |
| use_gpu = self.config.total_visible_gpus > 0 |
| if use_gpu: |
| self.checkpoint = self.checkpoint.cuda() |
| self.ranker = IndexScorer(self.index, use_gpu) |
|
|
| print_memory_stats() |
|
|
| def configure(self, **kw_args): |
| self.config.configure(**kw_args) |
|
|
| def encode(self, text: TextQueries): |
| queries = text if type(text) is list else [text] |
| bsize = 128 if len(queries) > 128 else None |
|
|
| self.checkpoint.query_tokenizer.query_maxlen = self.config.query_maxlen |
| |
| Q = self.checkpoint.queryFromText(queries, bsize=bsize, to_cpu=True) |
| return Q |
|
|
| def search(self, text: str, k=10, filter_fn=None): |
| Q = self.encode(text) |
| print(Q.shape) |
| return self.dense_search(Q, k, filter_fn=filter_fn) |
| |
| def search_with_mask(self, text: str, mask, k=100, filter_fn=None): |
| queries = text if type(text) is list else [text] |
| self.checkpoint.query_tokenizer.query_maxlen = self.config.query_maxlen |
| Q = self.checkpoint.queryFromText_withmask(queries, mask, bsize=None, to_cpu=True) |
| |
| return self.dense_search(Q, k, filter_fn=filter_fn) |
|
|
| def search_with_sentence(self, query, text_ls, k=100, filter_fn=None): |
| Q = [] |
| token_ls = myberttokenizer.tokenize(query)[:29] |
| token_emb_dic = {} |
| ori_Q = self.encode(query).squeeze(0) |
| token_emb_dic['101'] = ori_Q[0, :] |
| token_emb_dic['1'] = ori_Q[1, :] |
| Q.append(token_emb_dic['101']) |
| Q.append(token_emb_dic['1']) |
|
|
| for i in range(len(token_ls)): |
| token_emb_dic[token_ls[i]] = ori_Q[2+i, :] |
| token_emb_dic['102'] = ori_Q[len(token_ls)+2 , :] |
|
|
| for text in text_ls: |
| tmp = [] |
| tmp_token_ls = myberttokenizer.tokenize(text) |
| for tmp_token in tmp_token_ls: |
| if tmp_token in token_ls: |
| if tmp_token in token_emb_dic.keys(): |
| tmp.append(token_emb_dic[tmp_token]) |
| if len(tmp) > 0: |
| stacked_embeddings = torch.mean(torch.stack(tmp), dim = 0) |
| Q.append(stacked_embeddings) |
| Q.append(token_emb_dic['102']) |
| if len(Q) > 3: |
| Q = torch.stack(Q, dim=0) |
| Q = torch.nn.functional.pad(Q, (0, 0, 0, 32 - min(Q.shape[0], 32))) |
| Q = Q.unsqueeze(0) |
| return self.dense_search(Q, k, filter_fn=filter_fn) |
| else: |
| print("still keep") |
| return self.dense_search(ori_Q, k, filter_fn=filter_fn) |
| |
|
|
| def search_all(self, queries: TextQueries, k=10, filter_fn=None): |
| queries = Queries.cast(queries) |
| queries_ = list(queries.values()) |
|
|
| Q = self.encode(queries_) |
|
|
| return self._search_all_Q(queries, Q, k, filter_fn=filter_fn) |
|
|
| def _search_all_Q(self, queries, Q, k, filter_fn=None): |
| all_scored_pids = [list(zip(*self.dense_search(Q[query_idx:query_idx+1], k, filter_fn=filter_fn))) |
| for query_idx in tqdm(range(Q.size(0)))] |
|
|
| data = {qid: val for qid, val in zip(queries.keys(), all_scored_pids)} |
|
|
| provenance = Provenance() |
| provenance.source = 'Searcher::search_all' |
| provenance.queries = queries.provenance() |
| provenance.config = self.config.export() |
| provenance.k = k |
|
|
| return Ranking(data=data, provenance=provenance) |
|
|
| def dense_search(self, Q: torch.Tensor, k=10, filter_fn=None): |
| if k <= 10: |
| if self.config.ncells is None: |
| self.configure(ncells=1) |
| if self.config.centroid_score_threshold is None: |
| self.configure(centroid_score_threshold=0.5) |
| if self.config.ndocs is None: |
| self.configure(ndocs=256) |
| elif k <= 100: |
| if self.config.ncells is None: |
| self.configure(ncells=2) |
| if self.config.centroid_score_threshold is None: |
| self.configure(centroid_score_threshold=0.45) |
| if self.config.ndocs is None: |
| self.configure(ndocs=1024) |
| else: |
| if self.config.ncells is None: |
| self.configure(ncells=4) |
| if self.config.centroid_score_threshold is None: |
| self.configure(centroid_score_threshold=0.4) |
| if self.config.ndocs is None: |
| self.configure(ndocs=max(k * 4, 4096)) |
|
|
| pids, scores = self.ranker.rank(self.config, Q, filter_fn=filter_fn) |
|
|
| return pids[:k], list(range(1, k+1)), scores[:k] |
|
|
