thu-coai/esconv

1import json2import tqdm3import numpy as np4import multiprocessing as mp5import random6from collections import Counter7random.seed(13)8 9 10def _norm(x):11    return ' '.join(x.strip().split())12 13 14strategies = json.load(open('./strategy.json'))15strategies = [e[1:-1] for e in strategies]16strat2id = {strat: i for i, strat in enumerate(strategies)}17original = json.load(open('./ESConv.json'))18 19def process_data(d):20    dial = []21    for uttr in d['dialog']:22        text = _norm(uttr['content'])23        role = uttr['speaker']24        if role == 'seeker':25            dial.append({26                'text': text,27                'speaker': 'usr',28            })29        else:30            dial.append({31                'text': text,32                'speaker': 'sys',33                'strategy': uttr['annotation']['strategy'],34            })35    d['dialog'] = dial36    return d37 38data = []39 40for e in map(process_data, tqdm.tqdm(original, total=len(original))):41    data.append(e)42 43emotions = Counter([e['emotion_type'] for e in data])44problems = Counter([e['problem_type'] for e in data])45print('emotion', emotions)46print('problem', problems)47 48 49random.shuffle(data)50dev_size = int(0.15 * len(data))51test_size = int(0.15 * len(data))52valid = data[:dev_size]53test = data[dev_size: dev_size + test_size]54train = data[dev_size + test_size:]55 56print('train', len(train))57with open('./train.txt', 'w') as f:58    for e in train:59        f.write(json.dumps(e) + '\n')60with open('./sample.json', 'w') as f:61    json.dump(train[:10], f, ensure_ascii=False, indent=2)62 63print('valid', len(valid))64with open('./valid.txt', 'w') as f:65    for e in valid:66        f.write(json.dumps(e) + '\n')67 68print('test', len(test))69with open('./test.txt', 'w') as f:70    for e in test:71        f.write(json.dumps(e) + '\n')72