CSS-LM

1
import json
2
import sys
3
import random
4

5
file = sys.argv[1]
6
fewshot_n = int(sys.argv[2])
7
with open(file) as f:
8
    data = json.load(f)
9

10
label_list = list()
11
for line in data:
12
    label_list.append(line["sentiment"])
13

14
total_len = len(label_list)
15
label_list = list(set(label_list))
16

17
label_dict = dict()
18
for line in data:
19
    try:
20
        label_dict[line["sentiment"]].append(line)
21
    except:
22
        label_dict[line["sentiment"]] = []
23
        label_dict[line["sentiment"]].append(line)
24

25
print("=========")
26
print("Number:",fewshot_n*len(label_list))
27
#print("Sample N:",fewshot_n)
28
#print("Total len",total_len)
29
print("=========")
30

31
train_n = list()
32
for label in label_list:
33
    print("+++",len(label_dict[label]))
34
    #train_n += random.choices(label_dict[label],k=fewshot_n)
35
    #samples = random.sample(label_dict[label],min(len(label_dict[label]),fewshot_n))
36

37

38
    #ratio_num =round( fewshot_n*(len(label_dict[label])/total_len))+1
39
    #print("==",ratio_num)
40
    samples = random.sample(label_dict[label],min(len(label_dict[label]),fewshot_n))
41
    if len(samples) < fewshot_n:
42
        samples += random.choices(label_dict[label],k=fewshot_n-len(samples))
43
    train_n += samples
44
    #print(train_n)
45
    print(label,len(samples))
46
    print("--------")
47

48
print("=========")
49
print("Final Sample",len(train_n))
50

51
#with open("train_"+str(fewshot_n)+".json", 'w') as f:
52
with open("train.json_"+str(fewshot_n), 'w') as f:
53
    json.dump(train_n, f)
54