interpret_entity_embeddings.py (3839B)
1 import numpy as np 2 import os 3 import json 4 import collections 5 import ipdb 6 import clip 7 import torch 8 from sklearn.manifold import TSNE 9 from sklearn.decomposition import PCA 10 import tqdm 11 import pandas as pd 12 import matplotlib.pyplot as plt 13 import seaborn as sns 14 import pickle 15 16 labels_list = ["hero", "villain", "victim", "other"] 17 label_frequencies = [0.027121160214685396, 0.13857485440219253, 0.05195843325339728, 0.7823455521297248] 18 19 20 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") 21 min_occurences=30 #min nb of occurrences of entities 22 23 entities_dict = collections.defaultdict(list) 24 annot_file = os.path.join('../mmsrl', 'all', 'annotations', 'train' + '.jsonl') 25 with open(annot_file, 'r') as json_file: 26 dataset = list(map(json.loads, json_file)) 27 for sample in dataset: 28 for i, label in enumerate(labels_list): 29 for entity in sample[label]: 30 entities_dict[entity].append(label) 31 32 def normalise_label_freq(counter): 33 counts = np.array([counter.get(label, 0) for label in labels_list]) 34 counts = np.divide(counts, label_frequencies) 35 counts = counts / sum(counts) 36 return counts 37 38 39 entities = [] 40 main_label = [] 41 label_distrib = [] 42 frequency = [] 43 for entity, labels in entities_dict.items(): 44 entities.append(entity) 45 counter = collections.Counter(labels) 46 distrib = normalise_label_freq(counter) 47 main_label.append(distrib.argmax()) 48 label_distrib.append(distrib) 49 frequency.append(len(labels)) 50 51 print(f"Extracted {len(entities_dict)} entities.") 52 entities_dict_reduc = {k:v for k,v in entities_dict.items() if len(v)>min_occurences} 53 print(f"keeping {len(entities_dict_reduc)} entities appearing more than {min_occurences} times in the dataset.") 54 labels_reduc = [l for l, e in zip(main_label, entities) if e in entities_dict_reduc] 55 distrib_reduc = [l for l, e in zip(label_distrib, entities) if e in entities_dict_reduc] 56 57 print(collections.Counter(labels_reduc)) 58 # Get list of entities and encode them 59 entities_list_final = entities_dict_reduc.keys() 60 entities_input_ids = [clip.tokenize(entity, truncate=True).squeeze(0).to(device) for entity in entities_list_final] 61 embeddings = [] 62 63 # load model and get entities embeddings 64 65 model_path = '/data/almanach/user/smontari/scratch/constraint_challenge/Multimodal-semantic-role-labeling/models/model_clip_0.bin' 66 model = torch.load(model_path, map_location=torch.device(device)) 67 68 print("loading model...") 69 clip_model, _ = clip.load('ViT-L/14', device=device) 70 for entity_encoding in tqdm.tqdm(entities_input_ids): 71 entities_features = clip_model.encode_text(entity_encoding.unsqueeze(0)).to(dtype=model['entities_linear.weight'].dtype) 72 entities_embedding = torch.nn.functional.linear(entities_features, weight = model['entities_linear.weight'], bias = model['entities_linear.bias']) 73 embeddings.append(np.array(entities_embedding.squeeze(0).cpu().detach(), dtype='float32')) 74 75 76 # tsne = TSNE(n_components=2, init='random').fit_transform(embeddings) 77 78 79 pca = PCA(n_components=2) 80 pca.fit(embeddings) 81 print(pca.explained_variance_ratio_) #array([0.3405144 , 0.18068498]) 82 83 pca_features = pca.transform(embeddings) 84 85 df = pd.DataFrame() 86 df["comp-1"] = pca_features[:,0] 87 df["comp-2"] = pca_features[:,1] 88 89 final_dict = {} 90 for i, entity in enumerate(entities_list_final): 91 final_dict[entity] = {} 92 counter = collections.Counter(entities_dict[entity]) 93 final_dict[entity]['distrib'] = np.array([counter.get(label, 0) for label in labels_list]) 94 final_dict[entity]['feature'] = pca_features[i] 95 96 pickle.dump(final_dict, open('pca_coordinates.pkl', 'wb')) 97 98 plt.figure(figsize=(8,5)) 99 100 sns.scatterplot(x="comp-1", y="comp-2", hue=labels_reduc, 101 palette=sns.color_palette("hls", 4), 102 data=df, legend="full").set(title="T-SNE projection") 103 plt.savefig('embeddings_pca_correct.png')