Spaces:
Running
on
A10G
Running
on
A10G
File size: 980 Bytes
5244231 |
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 |
import os
import json
import argparse
import pandas as pd
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--annotation-file", type=str, required=True)
parser.add_argument("--result-dir", type=str, required=True)
parser.add_argument("--upload-dir", type=str, required=True)
parser.add_argument("--experiment", type=str, required=True)
return parser.parse_args()
if __name__ == "__main__":
args = get_args()
df = pd.read_table(args.annotation_file)
cur_df = df.copy()
cur_df = cur_df.drop(columns=['hint', 'category', 'source', 'image', 'comment', 'l2-category'])
cur_df.insert(6, 'prediction', None)
for pred in open(os.path.join(args.result_dir, f"{args.experiment}.jsonl")):
pred = json.loads(pred)
cur_df.loc[df['index'] == pred['question_id'], 'prediction'] = pred['text']
cur_df.to_excel(os.path.join(args.upload_dir, f"{args.experiment}.xlsx"), index=False, engine='openpyxl')
|