geonmin-kim's picture
Upload folder using huggingface_hub
d6585f5
#
# 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.
#
import os
import subprocess
import shutil
def clean_files(files):
for file in files:
if os.path.exists(file):
if os.path.isdir(file):
shutil.rmtree(file)
else:
os.remove(file)
def run_command(cmd, echo=False):
process = subprocess.Popen(cmd.split(),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
stdout, stderr = process.communicate()
stdout = stdout.decode('utf-8')
stderr = stderr.decode('utf-8')
if stderr and echo:
print(stderr)
if echo:
print(stdout)
return stdout, stderr
def parse_score(output, metric, digits=4):
"""Function for parsing the output from `pyserini.eval.trec_eval`."""
lines = output.split('\n')
# The output begins with a bunch of debug information, get rid of lines until we get to 'Results'
while 'Results' not in lines[0]:
lines.pop(0)
for line in lines:
if metric in line:
score = float(line.split()[-1])
return round(score, digits)
return None
def parse_score_qa(output, metric, digits=4):
"""Function for parsing the output from `pyserini.eval.evaluate_dpr_retrieval`. Currently, the implementation is
the same as `parse_score_msmacro`, but we're keeping separate in case they diverge in the future."""
for line in output.split('\n'):
if metric in line:
score = float(line.split()[-1])
return round(score, digits)
return None
def parse_score_msmarco(output, metric, digits=4):
"""Function for parsing the output from MS MARCO eval scripts. Currently, the implementation is the same as
`parse_score_qa`, but we're keeping separate in case they diverge in the future."""
for line in output.split('\n'):
if metric in line:
score = float(line.split()[-1])
return round(score, digits)
return None
def parse_score_msmarco_as_string(output, metric):
"""Function for parsing the output from MS MARCO eval scripts, but returning result as a string. This is used for
checking results to the entire degree of precision that the script generates."""
for line in output.split('\n'):
if metric in line:
return line.split()[-1]
return None
def run_retrieval_and_return_scores(output_file, retrieval_cmd, qrels, eval_type, metrics):
temp_files = [output_file]
# Take the base retrieval command and append the output file name to it.
os.system(retrieval_cmd + f' --output {output_file}')
scores = {}
# How we compute eval metrics depends on the `eval_type`.
if eval_type == 'trec_eval':
for metric in metrics:
cmd = f'python -m pyserini.eval.trec_eval -m {metric[0]} {qrels} {output_file}'
stdout, stderr = run_command(cmd)
scores[metric[0]] = parse_score(stdout, metric[1])
elif eval_type == 'msmarco_passage':
cmd = f'python -m pyserini.eval.msmarco_passage_eval {qrels} {output_file}'
stdout, stderr = run_command(cmd)
scores['MRR@10'] = parse_score_msmarco(stdout, 'MRR @10')
elif eval_type == 'msmarco_passage_string':
cmd = f'python -m pyserini.eval.msmarco_passage_eval {qrels} {output_file}'
stdout, stderr = run_command(cmd)
scores['MRR@10'] = parse_score_msmarco_as_string(stdout, 'MRR @10')
elif eval_type == 'msmarco_doc':
cmd = f'python -m pyserini.eval.msmarco_doc_eval --judgments {qrels} --run {output_file}'
stdout, stderr = run_command(cmd)
scores['MRR@100'] = parse_score_msmarco(stdout, 'MRR @100')
elif eval_type == 'msmarco_doc_string':
cmd = f'python -m pyserini.eval.msmarco_doc_eval --judgments {qrels} --run {output_file}'
stdout, stderr = run_command(cmd)
scores['MRR@100'] = parse_score_msmarco_as_string(stdout, 'MRR @100')
else:
clean_files(temp_files)
raise ValueError('Unknown eval_type!')
clean_files(temp_files)
return scores