Spaces:
Runtime error
Runtime error
| # | |
| # Pyserini: Reproducible IR research with sparse and dense representations | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| """ | |
| This module provides Pyserini's hybrid searcher by Dense + Sparse | |
| """ | |
| from typing import List, Dict | |
| from pyserini.search.lucene import LuceneSearcher | |
| from pyserini.search.faiss import FaissSearcher, DenseSearchResult | |
| class HybridSearcher: | |
| """Hybrid Searcher for dense + sparse | |
| Parameters | |
| ---------- | |
| dense_searcher : FaissSearcher | |
| sparse_searcher : LuceneSearcher | |
| """ | |
| def __init__(self, dense_searcher, sparse_searcher): | |
| self.dense_searcher = dense_searcher | |
| self.sparse_searcher = sparse_searcher | |
| def search(self, query: str, k0: int = 10, k: int = 10, alpha: float = 0.1, normalization: bool = False, weight_on_dense: bool = False) -> List[DenseSearchResult]: | |
| dense_hits = self.dense_searcher.search(query, k0) | |
| sparse_hits = self.sparse_searcher.search(query, k0) | |
| return self._hybrid_results(dense_hits, sparse_hits, alpha, k, normalization, weight_on_dense) | |
| def batch_search(self, queries: List[str], q_ids: List[str], k0: int = 10, k: int = 10, threads: int = 1, | |
| alpha: float = 0.1, normalization: bool = False, weight_on_dense: bool = False) \ | |
| -> Dict[str, List[DenseSearchResult]]: | |
| dense_result = self.dense_searcher.batch_search(queries, q_ids, k0, threads) | |
| sparse_result = self.sparse_searcher.batch_search(queries, q_ids, k0, threads) | |
| hybrid_result = { | |
| key: self._hybrid_results(dense_result[key], sparse_result[key], alpha, k, normalization, weight_on_dense) | |
| for key in dense_result | |
| } | |
| return hybrid_result | |
| def _hybrid_results(dense_results, sparse_results, alpha, k, normalization=False, weight_on_dense=False): | |
| dense_hits = {hit.docid: hit.score for hit in dense_results} | |
| sparse_hits = {hit.docid: hit.score for hit in sparse_results} | |
| hybrid_result = [] | |
| min_dense_score = min(dense_hits.values()) if len(dense_hits) > 0 else 0 | |
| max_dense_score = max(dense_hits.values()) if len(dense_hits) > 0 else 1 | |
| min_sparse_score = min(sparse_hits.values()) if len(sparse_hits) > 0 else 0 | |
| max_sparse_score = max(sparse_hits.values()) if len(sparse_hits) > 0 else 1 | |
| for doc in set(dense_hits.keys()) | set(sparse_hits.keys()): | |
| if doc not in dense_hits: | |
| sparse_score = sparse_hits[doc] | |
| dense_score = min_dense_score | |
| elif doc not in sparse_hits: | |
| sparse_score = min_sparse_score | |
| dense_score = dense_hits[doc] | |
| else: | |
| sparse_score = sparse_hits[doc] | |
| dense_score = dense_hits[doc] | |
| if normalization: | |
| sparse_score = (sparse_score - (min_sparse_score + max_sparse_score) / 2) \ | |
| / (max_sparse_score - min_sparse_score) | |
| dense_score = (dense_score - (min_dense_score + max_dense_score) / 2) \ | |
| / (max_dense_score - min_dense_score) | |
| score = alpha * sparse_score + dense_score if not weight_on_dense else sparse_score + alpha * dense_score | |
| hybrid_result.append(DenseSearchResult(doc, score)) | |
| return sorted(hybrid_result, key=lambda x: x.score, reverse=True)[:k] | |