Spaces:
Running
on
Zero
Running
on
Zero
import os | |
import pickle | |
import matplotlib.pyplot as plt | |
import networkx as nx | |
import torch | |
from tqdm import tqdm | |
from Modules.ToucanTTS.InferenceToucanTTS import ToucanTTS | |
from Utility.utils import load_json_from_path | |
distance_types = ["tree", "asp", "map", "learned", "l1"] | |
modes = ["plot_all", "plot_neighbors"] | |
neighbor = "Latin" | |
num_neighbors = 12 | |
distance_type = distance_types[0] # switch here | |
mode = modes[1] | |
edge_threshold = 0.01 | |
# TODO histograms to figure out a good threshold | |
cache_root = "." | |
supervised_iso_codes = load_json_from_path(os.path.join(cache_root, "supervised_languages.json")) | |
if distance_type == "l1": | |
iso_codes_to_ids = load_json_from_path(os.path.join(cache_root, "iso_lookup.json"))[-1] | |
model_path = "../../Models/ToucanTTS_Meta/best.pt" | |
checkpoint = torch.load(model_path, map_location='cpu') | |
embedding_provider = ToucanTTS(weights=checkpoint["model"], config=checkpoint["config"]).encoder.language_embedding | |
embedding_provider.requires_grad_(False) | |
l1_dist = dict() | |
seen_langs = set() | |
for lang_1 in supervised_iso_codes: | |
if lang_1 not in seen_langs: | |
seen_langs.add(lang_1) | |
l1_dist[lang_1] = dict() | |
for lang_2 in supervised_iso_codes: | |
if lang_2 not in seen_langs: # it's symmetric | |
l1_dist[lang_1][lang_2] = torch.nn.functional.mse_loss(embedding_provider(torch.LongTensor([iso_codes_to_ids[lang_1]])).squeeze(), embedding_provider(torch.LongTensor([iso_codes_to_ids[lang_2]])).squeeze()) | |
largest_value_l1_dist = 0.0 | |
for _, values in l1_dist.items(): | |
for _, value in values.items(): | |
largest_value_l1_dist = max(largest_value_l1_dist, value) | |
for key1 in l1_dist: | |
for key2 in l1_dist[key1]: | |
l1_dist[key1][key2] = l1_dist[key1][key2] / largest_value_l1_dist | |
distance_measure = l1_dist | |
if distance_type == "tree": | |
tree_lookup_path = os.path.join(cache_root, "lang_1_to_lang_2_to_tree_dist.json") | |
tree_dist = load_json_from_path(tree_lookup_path) | |
distance_measure = tree_dist | |
if distance_type == "map": | |
map_lookup_path = os.path.join(cache_root, "lang_1_to_lang_2_to_map_dist.json") | |
map_dist = load_json_from_path(map_lookup_path) | |
largest_value_map_dist = 0.0 | |
for _, values in map_dist.items(): | |
for _, value in values.items(): | |
largest_value_map_dist = max(largest_value_map_dist, value) | |
for key1 in map_dist: | |
for key2 in map_dist[key1]: | |
map_dist[key1][key2] = map_dist[key1][key2] / largest_value_map_dist | |
distance_measure = map_dist | |
if distance_type == "learned": | |
learned_lookup_path = os.path.join(cache_root, "lang_1_to_lang_2_to_map_dist.json") | |
learned_dist = load_json_from_path(learned_lookup_path) | |
largest_value_learned_dist = 0.0 | |
for _, values in learned_dist.items(): | |
for _, value in values.items(): | |
largest_value_learned_dist = max(largest_value_learned_dist, value) | |
for key1 in learned_dist: | |
for key2 in learned_dist[key1]: | |
learned_dist[key1][key2] = learned_dist[key1][key2] / largest_value_learned_dist | |
distance_measure = learned_dist | |
if distance_type == "asp": | |
asp_dict_path = os.path.join(cache_root, "asp_dict.pkl") | |
with open(asp_dict_path, 'rb') as dictfile: | |
asp_sim = pickle.load(dictfile) | |
lang_list = list(asp_sim.keys()) | |
asp_dist = dict() | |
seen_langs = set() | |
for lang_1 in lang_list: | |
if lang_1 not in seen_langs: | |
seen_langs.add(lang_1) | |
asp_dist[lang_1] = dict() | |
for index, lang_2 in enumerate(lang_list): | |
if lang_2 not in seen_langs: # it's symmetric | |
asp_dist[lang_1][lang_2] = 1 - asp_sim[lang_1][index] | |
distance_measure = asp_dist | |
iso_codes_to_names = load_json_from_path(os.path.join(cache_root, "iso_to_fullname.json")) | |
distances = list() | |
for lang_1 in distance_measure: | |
if lang_1 not in iso_codes_to_names: | |
continue | |
if lang_1 not in supervised_iso_codes and iso_codes_to_names[lang_1] != neighbor: | |
continue | |
for lang_2 in distance_measure[lang_1]: | |
try: | |
if lang_2 not in supervised_iso_codes and iso_codes_to_names[lang_2] != neighbor: | |
continue | |
except KeyError: | |
continue | |
distances.append((iso_codes_to_names[lang_1], iso_codes_to_names[lang_2], distance_measure[lang_1][lang_2])) | |
# Create a graph | |
G = nx.Graph() | |
# Add edges along with distances as weights | |
min_dist = min(d for _, _, d in distances) | |
max_dist = max(d for _, _, d in distances) | |
normalized_distances = [(entity1, entity2, (d - min_dist) / (max_dist - min_dist)) for entity1, entity2, d in distances] | |
if mode == "plot_neighbors": | |
fullnames = list() | |
fullnames.append(neighbor) | |
for code in supervised_iso_codes: | |
fullnames.append(iso_codes_to_names[code]) | |
supervised_iso_codes = fullnames | |
d_dist = list() | |
for entity1, entity2, d in tqdm(normalized_distances): | |
if (neighbor == entity2 or neighbor == entity1) and (entity1 in supervised_iso_codes and entity2 in supervised_iso_codes): | |
if entity1 != entity2: | |
d_dist.append(d) | |
thresh = sorted(d_dist)[num_neighbors] | |
# distance_scores = sorted(d_dist)[:num_neighbors] | |
neighbors = list() | |
for entity1, entity2, d in tqdm(normalized_distances): | |
if (d < thresh and (neighbor == entity2 or neighbor == entity1)) and (entity1 in supervised_iso_codes and entity2 in supervised_iso_codes): | |
neighbors.append(entity1) | |
neighbors.append(entity2) | |
unique_neighbors = list(set(neighbors)) | |
unique_neighbors.remove(neighbor) | |
for entity1, entity2, d in tqdm(normalized_distances): | |
if (neighbor == entity2 or neighbor == entity1) and (entity1 in supervised_iso_codes and entity2 in supervised_iso_codes): | |
if entity1 != entity2 and d < thresh: | |
spring_tension = ((thresh - d) ** 2) * 20000 # for vis purposes | |
print(f"{d}-->{spring_tension}") | |
G.add_edge(entity1, entity2, weight=spring_tension) | |
for entity1, entity2, d in tqdm(normalized_distances): | |
if (entity2 in unique_neighbors and entity1 in unique_neighbors) and (entity1 in supervised_iso_codes and entity2 in supervised_iso_codes): | |
if entity1 != entity2: | |
spring_tension = 1 - d | |
G.add_edge(entity1, entity2, weight=spring_tension) | |
# Draw the graph | |
pos = nx.spring_layout(G, weight="weight") # Positions for all nodes | |
edges = G.edges(data=True) | |
# Draw nodes | |
nx.draw_networkx_nodes(G, pos, node_size=1, alpha=0.01) | |
# Draw edges with labels | |
edges_connected_to_specific_node = [(u, v) for u, v in G.edges() if u == neighbor or v == neighbor] | |
# nx.draw_networkx_edges(G, pos, alpha=0.1) | |
nx.draw_networkx_edges(G, pos, edgelist=edges_connected_to_specific_node, edge_color='red', alpha=0.3, width=3) | |
for u, v, d in edges: | |
if u == neighbor or v == neighbor: | |
nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): round((thresh - (d['weight'] / 20000) ** (1 / 2)) * 10, 2)}, font_color="red", alpha=0.3) # reverse modifications | |
else: | |
pass | |
# nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d['weight']}) | |
# Draw node labels | |
nx.draw_networkx_labels(G, pos, font_size=14, font_family='sans-serif', font_color='green') | |
nx.draw_networkx_labels(G, pos, labels={neighbor: neighbor}, font_size=14, font_family='sans-serif', font_color='red') | |
plt.title(f'Graph of {distance_type} Distances') | |
plt.subplots_adjust(left=0, right=1, top=1, bottom=0) | |
plt.tight_layout(pad=0) | |
plt.savefig("avg.png", dpi=300) | |
plt.show() | |
elif mode == "plot_all": | |
for entity1, entity2, d in tqdm(normalized_distances): | |
if d < edge_threshold and entity1 != entity2: | |
spring_tension = edge_threshold - d | |
G.add_edge(entity1, entity2, weight=spring_tension) | |
# Draw the graph | |
pos = nx.spring_layout(G, weight="weight") # Positions for all nodes | |
edges = G.edges(data=True) | |
# Draw nodes | |
nx.draw_networkx_nodes(G, pos, node_size=1, alpha=0.01) | |
# Draw edges with labels | |
nx.draw_networkx_edges(G, pos, alpha=0.1, edge_color="blue") | |
# nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d['weight'] for u, v, d in edges}) | |
# Draw node labels | |
nx.draw_networkx_labels(G, pos, font_size=10, font_family='sans-serif') | |
plt.title(f'Graph of {distance_type} Distances') | |
plt.subplots_adjust(left=0, right=1, top=1, bottom=0) | |
plt.tight_layout(pad=0) | |
plt.show() | |