<jupyter_start><jupyter_text>Data<jupyter_code>path_params = "../MEDeA/datasets/file_parameters/default/"
path_data = "../MEDeA/datasets/romance_languages/cog/"
langs = ["es", "it", "la"]
test_data = MedeaDatasetReader(folder_path=path_data, langs=langs, name="Triple", phonetized=True)

#lang_pairs = [["es", "it", "la"]]

all_langs = ["es", "it", "la"]
lang_pairs = [all_langs,
              ["it", "la"], ["es", "it"], ["es", "la"],
             ]
tokens = ["EOW", "es", "it", "la"]
phones = [p for p in set(get_all_phones(path_data, lang_pairs)) if p not in tokens]

all_words = sorted(phones)
vocab_dim = len(all_words) + len(tokens)

embed_1_params = EmbeddingParamsOneLang.from_json_file(f"{path_params}embed_1_parameters.json")
embed_params = {}
for lang in ["es", "it", "la"]:
    embed_params[lang] = embed_1_params
    embed_params[lang].phone_embedding.phone_categories = list(all_words)
    embed_params[lang].pseudotoken_embedding.list = tokens
embed_params = EmbeddingParams(embed_params)

# Used for finetuning - we keep all the pairs
train_set = DatasetBuilderEmbedder.build(data_list=[test_data], parameters=embed_params)
batches = train_set.to_batch(ShuffleType.DESCENDING, batch_size=1, 
                             langs=langs, lang_of_reference_for_shuffling=langs[0])
batch_info = train_set.get_batch_info()
batch = batches[100]
in_lang = langs[0]
out_lang = langs[1]
batch<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>seed = 0
epoch = 20
num_head = 4
num_layers = 2

model_path = os.path.expanduser("/Users/cfourrie/Desktop/Experiments_synthesis_Y1/Exp5/" + \
                                f"Transformer{num_head}head{num_layers}layer_20_36_{seed}" + \
                                f"/results/experiment/models/epoch_{epoch}")
model = CompositeModel.load(model_path)

predictions, attn = model.predict(batch, batch_info)
local_prediction = [
    train_set[out_lang].ix_to_item[int(c)]
    for c in predictions["decoder"][f"{in_lang}_{out_lang}"][0][0]
]
local_target = [
    train_set[out_lang].ix_to_item[int(c)]
    for c in batch[out_lang][0]
]
local_input = [
    train_set[in_lang].ix_to_item[int(c)]
    for c in batch[in_lang][0]
]

local_self_attn = []
local_enc_attn = []
local_attn = []
for l_num in range(num_layers):
    local_self_attn.append(torch.stack(
        [attn[f"{in_lang}_{out_lang}"][0][f"decoder_self_attn/layer{l_num}_head{h}"]
         for h in range(num_head)]
    ).unsqueeze(0))

    local_enc_attn.append(torch.stack(
        [attn[f"{in_lang}_{out_lang}"][0][f"encoder_attn/layer{l_num}_head{h}"]
         for h in range(num_head)], dim=0).transpose(0, 1))
                          
    local_attn.append(torch.stack(
        [attn[f"{in_lang}_{out_lang}"][0][f"decoder_attn/layer{l_num}_head{h}"]
         for h in range(num_head)]
    ).unsqueeze(0))
print(
    "".join([train_set[out_lang].ix_to_item[int(c)] for c in batch[in_lang][0]][1:-1]), 
    "=>", 
    "".join([train_set[out_lang].ix_to_item[int(c)] for c in batch[out_lang][0]][1:-1]))

enc = False
self = False

if enc:
    print("Encoder attention")
    head_view(local_enc_attn, local_input, local_input, prettify_tokens=False) 
else:
    if self:
        print("Decoder self attention")
        head_view(local_self_attn, local_prediction[1:], local_prediction, prettify_tokens=False)
    else:
        print("Decoder attention")
        head_view(local_attn, local_prediction[1:], local_input, prettify_tokens=False)<jupyter_output>tʃɾistjano => kristja:no
Decoder attention
{'fraction_non_alphanumeric': 0.15048268029528677, 'fraction_numerical': 0.011641113003975015, 'mean_word_length': 4.428351309707242, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '241264', 'license': 'Apache-2.0', 'n_tokens_mistral': 1361, 'n_tokens_neox': 1337, 'n_words': 200, 'path': 'MEDeA_headview_realdata.ipynb', 'repo_name': 'clefourrier/bertviz', 'size': '3522'}