exploration.py (4499B)
1 import numpy as np 2 import pandas as pd 3 import os 4 import json 5 from scipy.stats import entropy 6 import collections 7 from transformers import AutoTokenizer 8 9 import mmsrl.utils 10 11 12 tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased') 13 14 dataset = {} 15 16 #datasetnames = ['uspolitics','covid19'] 17 datasetnames = ['all'] 18 for name in datasetnames: 19 dataset[name] = {} 20 for split in ['train', 'val', 'test']: 21 annot_file = os.path.join('../mmsrl', name, 'annotations', split + '.jsonl') 22 with open(annot_file, 'r') as json_file: 23 dataset[name][split] = list(map(json.loads, json_file)) 24 25 labels_list = mmsrl.utils.LABELS 26 stats_table = pd.DataFrame() 27 metrics_names = ['size', 'avg_nb_tokens', 'max_nb_tokens', 'nb_token_quantile95','class_distrib', 'avg_nb_ents', 'max_nb_ents', 'avg_nb_token_ents', 'max_nb_token_ents', 'total_nb_ents', 'percent_ents_in_text', 'total_nb_unique_ents', 'nb_duplicated_ent', 'samples_with_duplicate', 'entropy'] 28 stats_table['metric_name'] = metrics_names 29 30 def fill_dict(sample, ent_label_dict): 31 for i, label in enumerate(labels_list): 32 for ent in sample[label]: 33 if ent not in ent_label_dict: 34 ent_label_dict[ent] = [0]*len(labels_list) 35 ent_label_dict[ent][i] +=1 36 37 for name in datasetnames: 38 for split in ['train', 'val', 'test']: 39 colname = name + '_' + split 40 metrics = [] 41 data = dataset[name][split] 42 metrics.append(len(data)) 43 text = [data[i]['OCR'] for i in range(len(data))] 44 tokenized = [tokenizer.tokenize(t) for t in text] 45 nbtok = [len(t) for t in tokenized] 46 metrics.append(np.mean(nbtok)) 47 metrics.append(np.max(nbtok)) 48 metrics.append(round(np.percentile(nbtok, 95), 2)) 49 #print(pd.Series(nbtok).describe()) 50 class_distrib = [0]*len(labels_list) 51 nbents = [] 52 entities = [] 53 duplicated_ents = 0 54 sample_with_duplicates = 0 55 ent_label_dict = {} 56 nb_ent_in_text = [] 57 for sample in data: 58 sample_ents = [] 59 # add entities and labels to the global dictionary 60 fill_dict(sample, ent_label_dict) 61 for i, label in enumerate(labels_list): 62 # compute class distribution 63 class_distrib[i] += len(sample[label]) 64 sample_ents.extend(sample[label]) 65 # check if some entities are in several classes 66 if len(sample_ents) != len(set(sample_ents)): 67 #print("Duplicated entity!") 68 #print(sample) 69 duplicated_ents += (len(sample_ents)- len(set(sample_ents))) 70 sample_with_duplicates +=1 71 entities.extend(sample_ents) 72 nbents.append(len(sample_ents)) 73 # check if the entities are in the text 74 ents_in_text = [(ent in sample['OCR']) for ent in sample_ents] 75 nb_ent_in_text.append(np.sum(ents_in_text)) 76 77 class_distrib /= np.sum(class_distrib) 78 class_distrib = [round(num, 2) for num in class_distrib] 79 metrics.append(class_distrib) 80 metrics.append(np.mean(nbents)) 81 metrics.append(np.max(nbents)) 82 83 tokenized = [tokenizer.tokenize(t) for t in entities] 84 nbtokent = [len(t) for t in tokenized] 85 metrics.append(np.mean(nbtokent)) 86 metrics.append(np.max(nbtokent)) 87 88 metrics.append(len(entities)) 89 90 percent_ents_in_text = np.sum(nb_ent_in_text) / len(entities) * 100 91 metrics.append(round(percent_ents_in_text, 2)) 92 metrics.append(len(set(entities))) 93 94 metrics.append(duplicated_ents) 95 metrics.append(sample_with_duplicates) 96 97 # entropy: how are the classes distributed fow each entity 98 99 dataset_avg_var = 0 100 dataset_avg_entropy = 0 101 for ent, distrib in ent_label_dict.items(): 102 mean_abs_dev_ent_ditrib = np.mean(distrib - np.mean(distrib)**2) 103 dataset_avg_var += mean_abs_dev_ent_ditrib 104 dataset_avg_entropy += entropy(distrib) 105 #mean_abs_dev_ent_ditrib = np.mean(np.absolute(distrib - np.mean(distrib))) 106 dataset_avg_var /= len(ent_label_dict) 107 dataset_avg_entropy /= len(ent_label_dict) 108 109 metrics.append((round(dataset_avg_var,2), round(dataset_avg_entropy, 2))) 110 111 stats_table[colname] = metrics 112 113 114 stats_table[colname] = metrics 115 116 117 print(stats_table) 118 119 stats_table.to_csv('stats_table.csv', sep='\t')