-
Notifications
You must be signed in to change notification settings - Fork 340
/
__main__.py
103 lines (90 loc) · 4.29 KB
/
__main__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#
# Pyserini: Python interface to the Anserini IR toolkit built on Lucene
#
# 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.
#
import argparse
from pyserini.search import get_topics, SimpleSearcher
from pyserini.search.reranker import ClassifierType, PseudoRelevanceClassifierReranker
from tqdm import tqdm
parser = argparse.ArgumentParser(description='Search a Lucene index.')
parser.add_argument('--index', type=str, metavar='path to index', required=True, help="Path to Lucene index.")
parser.add_argument('--topics', type=str, metavar='topic_name', required=True,
help="Name of topics. Available: robust04, robust05, core17, core18.")
parser.add_argument('--output', type=str, metavar='path', help="Path to output file.")
parser.add_argument('--bm25', action='store_true', default=True, help="Use BM25 (default).")
parser.add_argument('--rm3', action='store_true', help="Use RM3")
parser.add_argument('--qld', action='store_true', help="Use QLD")
parser.add_argument('--prcl', type=ClassifierType, nargs='+', default=[],
help='Specify the classifier PseudoRelevanceClassifierReranker uses.')
parser.add_argument('--prcl.vectorizer', dest='vectorizer', type=str,
help='Type of vectorizer. Available: TfidfVectorizer, BM25Vectorizer.')
parser.add_argument('--prcl.r', dest='r', type=int, default=10,
help='Number of positive labels in pseudo relevance feedback.')
parser.add_argument('--prcl.n', dest='n', type=int, default=100,
help='Number of negative labels in pseudo relevance feedback.')
parser.add_argument('--prcl.alpha', dest='alpha', type=float, default=0.5,
help='Alpha value for interpolation in pseudo relevance feedback.')
args = parser.parse_args()
topics = get_topics(args.topics)
searcher = SimpleSearcher(args.index)
search_rankers = []
if args.qld:
search_rankers.append('qld')
searcher.set_qld()
else:
search_rankers.append('bm25')
if args.rm3:
search_rankers.append('rm3')
searcher.set_rm3()
# invalid topics name
if topics == {}:
print(f'Topic {args.topics} Not Found')
exit()
# get re-ranker
use_prcl = args.prcl and len(args.prcl) > 0 and args.alpha > 0
if use_prcl is True:
ranker = PseudoRelevanceClassifierReranker(
args.index, args.vectorizer, args.prcl, r=args.r, n=args.n, alpha=args.alpha)
# build output path
output_path = args.output
if output_path is None:
if use_prcl is True:
clf_rankers = []
for t in args.prcl:
if t == ClassifierType.LR:
clf_rankers.append('lr')
elif t == ClassifierType.SVM:
clf_rankers.append('svm')
r_str = f'prcl.r_{args.r}'
n_str = f'prcl.n_{args.n}'
a_str = f'prcl.alpha_{args.alpha}'
clf_str = 'prcl_' + '+'.join(clf_rankers)
tokens1 = ['run', args.topics, '+'.join(search_rankers)]
tokens2 = [args.vectorizer, clf_str, r_str, n_str, a_str]
output_path = '.'.join(tokens1) + '-' + '-'.join(tokens2) + ".txt"
else:
tokens = ['run', args.topics, '+'.join(search_rankers), 'txt']
output_path = '.'.join(tokens)
print(f'Running {args.topics} topics, saving to {output_path}...')
with open(output_path, 'w') as target_file:
for index, topic in enumerate(tqdm(sorted(topics.keys()))):
search = topics[topic].get('title')
hits = searcher.search(search, 1000)
doc_ids = [hit.docid.strip() for hit in hits]
scores = [hit.score for hit in hits]
if use_prcl and len(hits) > (args.r + args.n):
scores, doc_ids = ranker.rerank(doc_ids, scores)
tag = output_path[:-4] if args.output is None else 'Anserini'
for i, (doc_id, score) in enumerate(zip(doc_ids, scores)):
target_file.write(f'{topic} Q0 {doc_id} {i + 1} {score:.6f} {tag}\n')