wehub-resource-sync
93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
#!/usr/bin/env python3 -u
|
|
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
import argparse
|
|
import glob
|
|
|
|
import numpy as np
|
|
|
|
|
|
DIM = 1024
|
|
|
|
|
|
def compute_dist(source_embs, target_embs, k=5, return_sim_mat=False):
|
|
target_ids = [tid for tid in target_embs]
|
|
source_mat = np.stack(source_embs.values(), axis=0)
|
|
normalized_source_mat = source_mat / np.linalg.norm(
|
|
source_mat, axis=1, keepdims=True
|
|
)
|
|
target_mat = np.stack(target_embs.values(), axis=0)
|
|
normalized_target_mat = target_mat / np.linalg.norm(
|
|
target_mat, axis=1, keepdims=True
|
|
)
|
|
sim_mat = normalized_source_mat.dot(normalized_target_mat.T)
|
|
if return_sim_mat:
|
|
return sim_mat
|
|
neighbors_map = {}
|
|
for i, sentence_id in enumerate(source_embs):
|
|
idx = np.argsort(sim_mat[i, :])[::-1][:k]
|
|
neighbors_map[sentence_id] = [target_ids[tid] for tid in idx]
|
|
return neighbors_map
|
|
|
|
|
|
def load_embeddings(directory, LANGS):
|
|
sentence_embeddings = {}
|
|
sentence_texts = {}
|
|
for lang in LANGS:
|
|
sentence_embeddings[lang] = {}
|
|
sentence_texts[lang] = {}
|
|
lang_dir = f"{directory}/{lang}"
|
|
embedding_files = glob.glob(f"{lang_dir}/all_avg_pool.{lang}.*")
|
|
for embed_file in embedding_files:
|
|
shard_id = embed_file.split(".")[-1]
|
|
embeddings = np.fromfile(embed_file, dtype=np.float32)
|
|
num_rows = embeddings.shape[0] // DIM
|
|
embeddings = embeddings.reshape((num_rows, DIM))
|
|
|
|
with open(f"{lang_dir}/sentences.{lang}.{shard_id}") as sentence_file:
|
|
for idx, line in enumerate(sentence_file):
|
|
sentence_id, sentence = line.strip().split("\t")
|
|
sentence_texts[lang][sentence_id] = sentence
|
|
sentence_embeddings[lang][sentence_id] = embeddings[idx, :]
|
|
|
|
return sentence_embeddings, sentence_texts
|
|
|
|
|
|
def compute_accuracy(directory, LANGS):
|
|
sentence_embeddings, sentence_texts = load_embeddings(directory, LANGS)
|
|
|
|
top_1_accuracy = {}
|
|
|
|
top1_str = " ".join(LANGS) + "\n"
|
|
for source_lang in LANGS:
|
|
top_1_accuracy[source_lang] = {}
|
|
top1_str += f"{source_lang} "
|
|
for target_lang in LANGS:
|
|
top1 = 0
|
|
top5 = 0
|
|
neighbors_map = compute_dist(
|
|
sentence_embeddings[source_lang], sentence_embeddings[target_lang]
|
|
)
|
|
for sentence_id, neighbors in neighbors_map.items():
|
|
if sentence_id == neighbors[0]:
|
|
top1 += 1
|
|
if sentence_id in neighbors[:5]:
|
|
top5 += 1
|
|
n = len(sentence_embeddings[target_lang])
|
|
top1_str += f"{top1/n} "
|
|
top1_str += "\n"
|
|
|
|
print(top1_str)
|
|
print(top1_str, file=open(f"{directory}/accuracy", "w"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Analyze encoder outputs")
|
|
parser.add_argument("directory", help="Source language corpus")
|
|
parser.add_argument("--langs", help="List of langs")
|
|
args = parser.parse_args()
|
|
langs = args.langs.split(",")
|
|
compute_accuracy(args.directory, langs)
|