Spaces:
Runtime error
Runtime error
import os | |
import ujson | |
import random | |
from argparse import ArgumentParser | |
from colbert.utils.utils import print_message, create_directory | |
from utility.utils.save_metadata import save_metadata | |
def main(args): | |
AllMetrics = {} | |
Scores = {} | |
for path in args.paths: | |
with open(path) as f: | |
metric = ujson.load(f) | |
AllMetrics[path] = metric | |
for k in args.metric: | |
metric = metric[k] | |
assert type(metric) is float | |
Scores[path] = metric | |
MaxKey = max(Scores, key=Scores.get) | |
MaxCKPT = int(MaxKey.split('/')[-2].split('.')[-1]) | |
MaxARGS = os.path.join(os.path.dirname(MaxKey), 'logs', 'args.json') | |
with open(MaxARGS) as f: | |
logs = ujson.load(f) | |
MaxCHECKPOINT = logs['checkpoint'] | |
assert MaxCHECKPOINT.endswith(f'colbert-{MaxCKPT}.dnn'), (MaxCHECKPOINT, MaxCKPT) | |
with open(args.output, 'w') as f: | |
f.write(MaxCHECKPOINT) | |
args.Scores = Scores | |
args.AllMetrics = AllMetrics | |
save_metadata(f'{args.output}.meta', args) | |
print('\n\n', args, '\n\n') | |
print(args.output) | |
print_message("#> Done.") | |
if __name__ == "__main__": | |
random.seed(12345) | |
parser = ArgumentParser(description='.') | |
# Input / Output Arguments | |
parser.add_argument('--metric', dest='metric', required=True, type=str) # e.g., success.20 | |
parser.add_argument('--paths', dest='paths', required=True, type=str, nargs='+') | |
parser.add_argument('--output', dest='output', required=True, type=str) | |
args = parser.parse_args() | |
args.metric = args.metric.split('.') | |
assert not os.path.exists(args.output), args.output | |
create_directory(os.path.dirname(args.output)) | |
main(args) | |