SNLIやMNLIデータのjsonl形式のラベルを取り出す

preprocess_nli.py

import json
import argparse
from collections import defaultdict, Counter


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, required=True,
                        help="jsonl形式のデータ")
    args = parser.parse_args()

    return args


def main(args):

    type_num_distributions = defaultdict(int)
    with open(args.input) as f:
        for l in f:
            d = json.loads(l)
            type_num = max(Counter(d["annotator_labels"]).values())
            type_num_distributions[type_num] += 1
    
    print(type_num_distributions)


if __name__ == "__main__":
    args = parse_args()
    main(args)