pminervini's picture
update
5ca644e
#!/usr/bin/env python3
import glob
import os
import random
from tqdm import tqdm
from datasets import Dataset, DatasetDict, load_dataset
def convert(list_of_dicts):
res = {}
for d in list_of_dicts:
for k, v in d.items():
res.setdefault(k, []).append(v)
return res
v10 = load_dataset("fever", "v1.0")
name_lst = ['train', 'labelled_dev']
old_to_new_label_map = {
'SUPPORTS': 'supported',
'REFUTES': 'refuted'
}
data_map = {}
for name in name_lst:
instance_lst = []
for entry in tqdm(v10[name]):
id_ = entry['id']
label = entry['label']
claim = entry['claim']
evidence_id = entry['evidence_id']
evidence_wiki_url = entry['evidence_wiki_url']
if evidence_id != -1:
assert label in {'SUPPORTS', 'REFUTES'}
instance = {'id': id_, 'label': old_to_new_label_map[label], 'claim': claim}
instance_lst.append(instance)
key = 'dev' if name in {'labelled_dev'} else name
instance_lst = sorted([dict(t) for t in {tuple(d.items()) for d in instance_lst}], key=lambda d: d['claim'])
label_to_instance_lst = {}
for e in instance_lst:
if e['label'] not in label_to_instance_lst:
label_to_instance_lst[e['label']] = []
label_to_instance_lst[e['label']].append(e)
min_len = min(len(v) for k, v in label_to_instance_lst.items())
new_instance_lst = []
for k in sorted(label_to_instance_lst.keys()):
new_instance_lst += label_to_instance_lst[k][:min_len]
random.Random(42).shuffle(new_instance_lst)
data_map[key] = new_instance_lst
ds_path = 'pminervini/hl-fever'
task_to_ds_map = {k: Dataset.from_dict(convert(v)) for k, v in data_map.items()}
ds_dict = DatasetDict(task_to_ds_map)
ds_dict.push_to_hub(ds_path, "v1.0")
# breakpoint()