def search(self,
q: Union[str, JQuery],
k: int = 10,
query_generator: JQueryGenerator = None,
strip_segment_id=False,
remove_dups=False) -> List[JLuceneSearcherResult]:
trec_runs, docid_to_search_result = list(), dict()
for searcher in self.searchers:
docid_score_pair = list()
hits = searcher.search(q,
k=k,
query_generator=query_generator,
strip_segment_id=strip_segment_id,
remove_dups=remove_dups)
for hit in hits:
docid_to_search_result[hit.docid] = hit
docid_score_pair.append((hit.docid, hit.score))
run = TrecRun.from_search_results(docid_score_pair)
trec_runs.append(run)
if self.method == FusionMethod.RRF:
fused_run = reciprocal_rank_fusion(trec_runs,
rrf_k=60,
depth=1000,
k=1000)
else:
raise NotImplementedError()
return self.convert_to_search_result(fused_run, docid_to_search_result)
from ..trectools import TrecRun
parser = argparse.ArgumentParser(description='Perform various ways of fusion given a list of trec run files.')
parser.add_argument('--runs', type=str, nargs='+', default=[], required=True,
help='List of run files separated by space.')
parser.add_argument('--output', type=str, required=True, help="Path to resulting fused txt.")
parser.add_argument('--runtag', type=str, default="pyserini.fusion", help="Tag name of fused run.")
parser.add_argument('--method', type=FusionMethod, default=FusionMethod.RRF, help="The fusion method to be used.")
parser.add_argument('--rrf.k', dest='rrf_k', type=int, default=60,
help="Parameter k needed for reciprocal rank fusion.")
parser.add_argument('--alpha', type=float, default=0.5, required=False, help='Alpha value used for interpolation.')
parser.add_argument('--depth', type=int, default=1000, required=False, help='Pool depth per topic.')
parser.add_argument('--k', type=int, default=1000, required=False, help='Number of documents to output per topic.')
parser.add_argument('--resort', action='store_true', help='We resort the Trec run files or not')
args = parser.parse_args()
trec_runs = [TrecRun(filepath=path,resort=args.resort) for path in args.runs]
fused_run = None
if args.method == FusionMethod.RRF:
fused_run = reciprocal_rank_fusion(trec_runs, rrf_k=args.rrf_k, depth=args.depth, k=args.k)
elif args.method == FusionMethod.INTERPOLATION:
fused_run = interpolation(trec_runs, alpha=args.alpha, depth=args.depth, k=args.k)
elif args.method == FusionMethod.AVERAGE:
fused_run = average(trec_runs, depth=args.depth, k=args.k)
else:
raise NotImplementedError(f'Fusion method {args.method} not implemented.')
fused_run.save_to_txt(args.output, tag=args.runtag)
