Spaces:
Runtime error
Runtime error
| import argparse | |
| from pyserini.search.faiss import AutoQueryEncoder, AnceQueryEncoder, DprQueryEncoder, TctColBertQueryEncoder | |
| def _init_encoder_from_str(encoder, device='cpu'): | |
| encoder_lower = encoder.lower() | |
| if 'dpr' in encoder_lower: | |
| return DprQueryEncoder(encoder_dir=encoder, device=device) | |
| elif 'tct_colbert' in encoder_lower: | |
| return TctColBertQueryEncoder(encoder_dir=encoder, device=device) | |
| elif 'ance' in encoder_lower: | |
| return AnceQueryEncoder(encoder_dir=encoder, device=device) | |
| elif 'sentence' in encoder_lower: | |
| return AutoQueryEncoder(encoder_dir=encoder, pooling='mean', l2_norm=True, device=device) | |
| else: | |
| return AutoQueryEncoder(encoder_dir=encoder, device=device) | |
| def load_index(searcher_class, index_path, query_encoder=None): | |
| if query_encoder is not None: | |
| searcher = searcher_class(index_path, query_encoder) | |
| else: | |
| searcher = searcher_class(index_path) | |
| return searcher | |
| class OnlineSearcher(object): | |
| def __init__(self, args): | |
| self.args = args | |
| # TEMPORARY: deal with missing args | |
| if 'index_type' not in args: | |
| args.index_type = 'hybrid' | |
| if 'index' not in args: | |
| args.index = '/root/indexes/mrtydi-korean/sparse,/root/indexes/mrtydi-korean/dense' | |
| if 'lang_abbr' not in args: | |
| args.lang_abbr = 'ko' | |
| if 'encoder' not in args: | |
| args.encoder = 'castorini/mdpr-question-nq' | |
| if 'device' not in args: | |
| args.device = 'cuda:0' | |
| if 'alpha' not in args: | |
| args.alpha = 0.5 | |
| if 'normalization' not in args: | |
| args.normalization = True | |
| if 'k' not in args: | |
| args.k = 10 | |
| if args.index_type == 'sparse': | |
| query_encoder = None | |
| elif args.index_type == 'dense' or args.index_type == 'hybrid': | |
| query_encoder = _init_encoder_from_str(encoder=args.encoder, device=args.device) | |
| else: | |
| raise ValueError(f"index_type {args.index_type} should be chosen among sparse, dense, or hybrid") | |
| # load index | |
| if args.index_type == 'hybrid': | |
| args.index = args.index.split(',') | |
| assert(len(args.index) == 2), "require both sparse and dense index delimited by comma" | |
| from pyserini.search.lucene import LuceneSearcher | |
| self.ssearcher = load_index(searcher_class=LuceneSearcher, index_path=args.index[0]) | |
| self.ssearcher.set_language(args.lang_abbr) | |
| from pyserini.search.faiss import FaissSearcher | |
| self.dsearcher = load_index(searcher_class=FaissSearcher, index_path=args.index[1], query_encoder=query_encoder) | |
| from pyserini.search.hybrid import HybridSearcher | |
| self.searcher = HybridSearcher(self.dsearcher, self.ssearcher) | |
| print(f"load {self.ssearcher.num_docs} documents from {args.index}") | |
| else: | |
| if args.index_type == 'sparse': | |
| from pyserini.search.lucene import LuceneSearcher as Searcher | |
| elif args.index_type == 'dense': | |
| from pyserini.search.faiss import FaissSearcher as Searcher | |
| self.searcher = load_index(searcher_class=Searcher, index_path=args.index, query_encoder=query_encoder) | |
| if args.index_type == 'sparse': | |
| self.searcher.set_language(args.lang_abbr) | |
| print(f"load {self.searcher.num_docs} documents from {args.index}") | |
| def search(self, query, k=10): | |
| if self.args.index_type == 'hybrid': | |
| hits = self.searcher.search(query, alpha=self.args.alpha, normalization=self.args.normalization, k=k) | |
| else: | |
| hits = self.searcher.search(query) | |
| return hits | |
| def print_result(self, hits, k): | |
| # Print the first 10 hits: | |
| for i in range(0, k): | |
| print(f'{i+1:2} {hits[i].docid:15} {hits[i].score:.5f}') | |
| if args.index_type == 'sparse': # faiss searcher does not store document raw text | |
| doc = self.searcher.doc(hits[i].docid) | |
| elif args.index_type == 'hybrid': | |
| doc = self.searcher.sparse_searcher.doc(hits[i].docid) | |
| else: | |
| doc = None | |
| if doc is not None: | |
| print(doc.raw()) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description='Search interactively') | |
| parser.add_argument('--index_type', type=str, required=True, | |
| help="choose indexing type", choices=['sparse', 'dense', 'hybrid']) | |
| parser.add_argument('--index', type=str, required=True, | |
| help="Path to index or name of prebuilt index.") | |
| parser.add_argument('--query', type=str, required=True, | |
| help="Query text") | |
| parser.add_argument('--lang_abbr', type=str, required=False, default='ko', | |
| help="for language specific algorithms for sparse retrieveal)") | |
| parser.add_argument('--encoder', type=str, required=False, | |
| help="encoder name or checkpoint path") | |
| parser.add_argument('--device', type=str, required=False, default='cpu', | |
| help="device to use for encoding queries (cf. pyserini does not support faiss-gpu)") | |
| # for hybrid search | |
| parser.add_argument('--alpha', type=float, default=0.5, | |
| help="weight for hybrid search: alpha*score(sparse) + score(dense)") | |
| parser.add_argument('--normalization', action='store_true', | |
| help="normalize sparse & dens score before fusion") | |
| # search range | |
| parser.add_argument('--k', type=int, default=10, | |
| help="the number of passages to return (default: 10)") | |
| args = parser.parse_args() | |
| # make searcher | |
| searcher = OnlineSearcher(args) | |
| print(f"given query: {args.query}") | |
| # search | |
| hits = searcher.search(args.query) | |
| # print results | |
| searcher.print_result(hits, args.k) |