Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add colbert version of IBM reranker #918

Merged
merged 20 commits into from
Jan 3, 2022
Merged

Add colbert version of IBM reranker #918

Changes from 4 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
e4a8cda
add non-parameter ibm baseline reranker
Nov 12, 2021
79beba1
merge with master
Nov 12, 2021
575cd16
Merge branch 'ibm_model1' of https://github.com/yuki617/pyserini into…
stephaniewhoo Nov 13, 2021
073123e
update doc_id loading method
Nov 13, 2021
29b1e44
add multithread
Nov 21, 2021
1e0982f
update ibmmodel equation
Nov 22, 2021
ddd3c37
Merge branch 'ibm_model1' of https://github.com/yuki617/pyserini into…
stephaniewhoo Nov 22, 2021
6b5fad3
add more typing hint and normalize function
Nov 29, 2021
2cfadc0
merge with master
Nov 29, 2021
f4baff6
remove unnecessary lines
Nov 29, 2021
217c2b2
Merge branch 'ibm_model1' of https://github.com/yuki617/pyserini into…
stephaniewhoo Nov 29, 2021
17e2b51
update with python name convention
Nov 29, 2021
480e7de
add rank
stephaniewhoo Dec 20, 2021
ad557c1
merge with master
Dec 21, 2021
8f45938
add ibm version of Colbert script
Dec 21, 2021
7d24e34
add both formula
stephaniewhoo Dec 21, 2021
dbccf20
add rbf kernel function
Dec 23, 2021
a4df378
fix a bug on selftranslation ratio adjustment
Dec 30, 2021
fed8a5d
rename the file and fix conflict
Dec 30, 2021
252a134
update parser comment
Dec 30, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 15 additions & 8 deletions scripts/rank_ibm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
JIndexReader = autoclass('io.anserini.index.IndexReaderUtils')
JTerm = autoclass('org.apache.lucene.index.Term')

SELF_TRAN = 0.35
SELF_TRAN = 0.1
MIN_PROB=0.0025
LAMBDA_VALUE = 0.3
MIN_COLLECT_PROB=1e-9
Expand Down Expand Up @@ -109,6 +109,7 @@ def get_ibm_score(arguments):
target_lookup = arguments['target_lookup']
tran = arguments['tran']
collect_probs = arguments['collect_probs']
max_sim = arguments['max_sim']

if searcher.documentRaw(test_doc) ==None:
print(f'{test_doc} is not found in searcher')
Expand All @@ -121,6 +122,7 @@ def get_ibm_score(arguments):
target_map = {}
total_tran_prob = 0
collect_prob = collect_probs[querytoken]
max_sim_score = 0
if querytoken in target_lookup.keys():
query_word_id = target_lookup[querytoken]
if query_word_id in tran.keys():
Expand All @@ -133,10 +135,13 @@ def get_ibm_score(arguments):
if doctoken in source_lookup.keys():
doc_word_id = source_lookup[doctoken]
if doc_word_id in target_map.keys():
tran_prob = max(target_map[doc_word_id],tran_prob)
total_tran_prob += (tran_prob/doc_size)

query_word_prob=math.log((1 - LAMBDA_VALUE) * total_tran_prob + LAMBDA_VALUE * collect_prob)
tran_prob = max(target_map[doc_word_id],tran_prob)
max_sim_score = max(tran_prob, max_sim_score)
total_tran_prob += (tran_prob/doc_size)
if (max_sim):
query_word_prob=math.log((1 - LAMBDA_VALUE) * max_sim_score + LAMBDA_VALUE * collect_prob)
else:
query_word_prob=math.log((1 - LAMBDA_VALUE) * total_tran_prob + LAMBDA_VALUE * collect_prob)

total_query_prob += query_word_prob
return total_query_prob /query_size
Expand Down Expand Up @@ -229,7 +234,7 @@ def load_tranprobs_table(dir_path: str):


def rank(qrels: str, base: str,tran_path:str, query_path:str, lucene_index_path: str,output_path:str, \
score_path:str,field_name:str, tag: str,alpha:int,num_threads:int):
score_path:str,field_name:str, tag: str,alpha:int,num_threads:int, max_sim:bool):

pool = ThreadPool(num_threads)
searcher = JSimpleSearcher(JString(lucene_index_path))
Expand Down Expand Up @@ -257,7 +262,7 @@ def rank(qrels: str, base: str,tran_path:str, query_path:str, lucene_index_path:
collect_probs[querytoken] = max(reader.totalTermFreq(JTerm(field_name, querytoken))/total_term_freq, MIN_COLLECT_PROB)
arguments = [{"query_text_lst":query_text_lst,"test_doc":test_doc, "searcher":searcher,\
"field_name":field_name,"source_lookup":source_lookup,"target_lookup":target_lookup,\
"tran":tran,"collect_probs":collect_probs} for test_doc in test_docs]
"tran":tran,"collect_probs":collect_probs, "max_sim":max_sim} for test_doc in test_docs]
rank_scores = pool.map(get_ibm_score, arguments)

ibm_scores = normalize([p for p in rank_scores])
Expand Down Expand Up @@ -302,9 +307,11 @@ def rank(qrels: str, base: str,tran_path:str, query_path:str, lucene_index_path:
metavar="type of field", help='interpolation weight')
parser.add_argument('-num_threads', type=int, default="12",
metavar="num_of_threads", help='number of threads to use')
parser.add_argument('-max_sim', type=bool, default=False,
metavar="bool for max sim operator", help='whether we use max sim operator')
args = parser.parse_args()

print('Using base run:', args.base)

rank(args.qrels, args.base, args.tran_path, args.query_path, args.index, args.output, \
args.score_path,args.field_name, args.tag,args.alpha,args.num_threads)
args.score_path,args.field_name, args.tag,args.alpha,args.num_threads, args.max_sim)