1
import torch
2
import torch.nn as nn
3
from torch.utils.data import Dataset, DataLoader, random_split
4

5
from datasets import load_dataset
6
from tokenizers import Tokenizer
7

8
from tokenizers.models import WordLevel
9
from tokenizers.trainers import WordLevelTrainer # tokenizer相对应的trainer，通过给定的句子创建词表
10
from tokenizers.pre_tokenizers import Whitespace # 按空格分词
11

12
from pathlib import Path
13

14

15
def get_or_build_tokenizer(config, ds, lang):
16
    """
17
    创建分词器
18

19
    参数:
20
        `config`: 模型的配置
21
        `ds`: 数据集
22
        `lang`: 分词器的语言
23
    返回:
24
        `Tokenizer`分词器实例
25
    """
26
    # config['tokenizer_file'] = '../tokenizer/tokenizer_{0}.json'
27
    tokenizer_path = Path(config['tokenizer_file'].format(lang))
28
    if not Path.exists(tokenizer_path):
29
        tokenizer = Tokenizer(WordLevel(unk_token='[UNK]')) # 分词器在遇到不在词表中的词时将其替换为`[UNK]`
30
        tokenizer.pre_tokenizer = Whitespace()
31
        trainer = WordLevelTrainer(special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"], min_frequency=2) # 一个词要出现在词表中，需要至少在句子中出现两次
32
        tokenizer.train_from_iterator(get_all_sentences(ds, lang), trainer=trainer)
33
    else:
34
        tokenizer = Tokenizer.from_file(str(tokenizer_path))
35

36
    return tokenizer

1
def get_all_sentences(ds, lang):
2
    """
3
    从数据集中取指定语言的句子
4

5
    参数:
6
        `ds`: 数据集
7
        `lang`: 指定语言 `"en"`或`"it"`
8
    """
9
    for item in ds:
10
        yield item['translation'][lang]

1
def get_ds(config):
2
    """
3
    加载数据集并创建分词器
4
    """
5
    # 在Hugging Face中下载opus_books数据集中的`en-it`子集，选择`train`部分，我们之后将自行再分出validation部分
6
    ds_raw = load_dataset('opus_books', f'{config["lang_src"]}-{config["lang_tgt"]}', split='train')
7

8
    # 创建分词器
9
    tokenizer_src = get_or_build_tokenizer(config, ds_raw, config['lang_src'])
10
    tokenizer_tgt = get_or_build_tokenizer(config, ds_raw, config['lang_tgt'])
11

12
    # 将数据集分为90%的训练集和10%的验证集
13
    train_ds_size = int(0.9 * len(ds_raw))
14
    val_ds_size = len(ds_raw) - train_ds_size
15
    train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size])

1
class BilingualDataset(Dataset):
2

3
    def __init__(self, ds, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len) -> None:
4
        super().__init__()
5

6
        self.ds = ds
7
        self.tokenizer_src = tokenizer_src
8
        self.tokenizer_tgt = tokenizer_tgt
9
        self.src_lang = src_lang
10
        self.tgt_lang = tgt_lang
11
        self.seq_len = seq_len
12

13
        self.sos_token = torch.tensor([tokenizer_tgt.token_to_id("[SOS]")], dtype=torch.int64)
14
        self.eos_token = torch.tensor([tokenizer_tgt.token_to_id("[EOS]")], dtype=torch.int64)
15
        self.pad_token = torch.tensor([tokenizer_tgt.token_to_id("[PAD]")], dtype=torch.int64)
16

17
    def __len__(self): # 数据集大小
18
        return len(self.ds)

1
    def __getitem__(self, index):
2
        src_target_pair = self.ds[index]
3
        src_text = src_target_pair['translation'][self.src_lang]
4
        tgt_text = src_target_pair['translation'][self.tgt_lang]

1
        # 将文本转换成token，再转换成id，即tokenizer先将句子拆分成词，再将词转换成词表中的id
2
        enc_input_tokens = self.tokenizer_src.encode(src_text).ids
3
        dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids

1
        # 填充句子，使其长度达到seq_len
2
        enc_num_padding_tokens = self.seq_len - len(enc_input_tokens) - 2
3
        # 因为decoder的输入只有sos没有eos，所以padding要多一个token
4
        dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1

1
        # 确保选择的seq_len长度满足所有样本，即padding应该不为负数
2
        if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:
3
            raise ValueError('Sentence is too long')

1
        # 为encoder的输入填充tensor，添加'[SOS]'、'[EOS]'、'[PAD]'
2
        encoder_input = torch.cat([
3
            self.sos_token,
4
            torch.tensor(enc_input_tokens, dtype=torch.int64),
5
            self.eos_token,
6
            torch.tensor([self.pad_token] * enc_num_padding_tokens, dtype=torch.int64)
7
        ])

1
        # 为decoder的输入填充tensor，只添加'[SOS]'和'[PAD]'
2
        decoder_input = torch.cat([
3
            self.sos_token,
4
            torch.tensor(dec_input_tokens, dtype=torch.int64),
5
            torch.tensor([self.pad_token] * dec_num_padding_tokens, dtype=torch.int64)
6
        ])

1
        # 为decoder的输出（也叫做label，即期望的翻译结果）填充tensor，只添加'[EOS]'
2
        label = torch.cat([
3
            torch.tensor(dec_input_tokens, dtype=torch.int64),
4
            self.eos_token,
5
            torch.tensor([self.pad_token] * dec_num_padding_tokens, dtype=torch.int64)
6
        ])

1
        # 再次检测填充后的tensor长度是否满足seq_len
2
        assert encoder_input.size(0) == self.seq_len
3
        assert decoder_input.size(0) == self.seq_len
4
        assert label.size(0) == self.seq_len
5

6
        return {
7
            'encoder_input': encoder_input, # (seq_len)
8
            'decoder_input': decoder_input, # (seq_len)
9
            'encoder_mask': (encoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int(), # (1, 1, seq_len)
10
            'decoder_mask': (decoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int() & causal_mask(decoder_input.size(0)), # (1, seq_len, seq_len))
11
            'label': label, # (seq_len)
12
            'src_text': src_text, # 原文，用于可视化
13
            'tgt_text': tgt_text, # 译文，用于可视化
14
        }

1
def causal_mask(size):
2
    # triu函数中的diagonal=1表示主对角线+1，即生成上三角矩阵后主对角线再置为0
3
    mask = torch.triu(torch.ones(1, size, size), diagonal=1).type(torch.int)
4
    # 再将矩阵取反，即生成下三角矩阵
5
    return mask == 0

1
from dataset import BilingualDataset, causal_mask
2

3

4
def get_ds(config):
5
    """
6
    加载数据集并创建分词器
7
    """
8
    # 在Hugging Face中下载opus_books数据集中的`en-it`子集，选择`train`部分，我们之后将自行再分出validation部分
9
    ds_raw = load_dataset('opus_books', f'{config["lang_src"]}-{config["lang_tgt"]}', split='train')
10

11
    # 创建分词器
12
    tokenizer_src = get_or_build_tokenizer(config, ds_raw, config['lang_src'])
13
    tokenizer_tgt = get_or_build_tokenizer(config, ds_raw, config['lang_tgt'])
14

15
    # 将数据集分为90%的训练集和10%的验证集
16
    train_ds_size = int(0.9 * len(ds_raw))
17
    val_ds_size = len(ds_raw) - train_ds_size
18
    train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size])
19

20
    train_ds = BilingualDataset(train_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
21
    val_ds = BilingualDataset(val_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
22

23
    # 找出训练集中源语言和目标语言里最长的句子长度分别是多少
24
    max_len_src = 0
25
    max_len_tgt = 0
26

27
    for item in ds_raw:
28
        src_ids = tokenizer_src.encode(item['translation'][config['lang_src']]).ids
29
        tgt_ids = tokenizer_tgt.encode(item['translation'][config['lang_tgt']]).ids
30
        max_len_src = max(max_len_src, len(src_ids))
31
        max_len_tgt = max(max_len_tgt, len(tgt_ids))
32

33
    print(f"Max length of source language: {max_len_src}")
34
    print(f"Max length of target language: {max_len_tgt}")
35

36
    train_dataloader = DataLoader(train_ds, batch_size=config['batch_size'], shuffle=True)
37
    val_dataloader = DataLoader(val_ds, batch_size=1, shuffle=True)
38

39
    return train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt

1
from model import build_transformer
2

3

4
def get_model(config, vocab_src_len, vocab_tgt_len):
5
    """
6
    创建模型
7
    """
8
    model = build_transformer(vocab_src_len, vocab_tgt_len, config['seq_len'], config['seq_len'], config['d_model'])
9
    return model

1
from pathlib import Path
2

3
def get_config():
4
    return {
5
        "batch_size": 8,
6
        "num_epochs": 20,
7
        # 在实际的训练中，lr应该根据实际情况调整，比如一开始设置较大的lr，然后在训练过程中逐渐减小lr
8
        # 但本文的目的在于展示Transformer的原理，使用可变的lr会增加代码复杂度，因此这里使用固定的lr
9
        "lr": 10**-4,
10
        "seq_len": 350,
11
        "d_model": 512,
12
        "lang_src": "en",
13
        "lang_tgt": "it",
14
        "model_folder": "weights",
15
        "model_basename": "tmodel_",
16
        # 预加载模型，如果不为空，则从预加载模型中加载参数（如在训练过程中程序意外中断），否则从头开始训练
17
        "preload": None,
18
        # 分词文件，如`_en.json`和`_it.json`
19
        "tokenizer_file": "tokenizer_{0}.json",
20
        # 为tensorboard提供的log路径
21
        "experiment_name": "runs/tmodel"
22
    }
23

24

25
def get_weights_file_path(config, epoch: str):
26
    model_folder = config["model_folder"]
27
    model_basename = config["model_basename"]
28
    model_filename = f"{model_basename}{epoch}.pt"
29
    # Path重载了`/`操作符，可以直接使用`/`创建子目录
30
    return str(Path('.') / model_folder / model_filename)

1
pip install tensorboard torch_tb_profiler

1
tensorboard --logdir runs

1
from torch.utils.tensorboard import SummaryWriter
2
from config import get_weights_file_path, get_config
3
from tqdm import tqdm
4

5

6
def train_model(config):
7
    """
8
    训练模型
9
    """
10

11
    # 选择设备
12
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
13
    print(f"Using device: {device}")
14

15
    # 确定权重文件目录存在
16
    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)
17

18
    # 加载数据集和分词器
19
    train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_ds(config)
20
    model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size()).to(device)
21

22
    # 启动Tensorboard以可视化loss
23
    writer = SummaryWriter(config['experiment_name'])
24

25
    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], eps=1e-9)
26

27
    # 若指定了预加载权重，则加载预加载权重
28
    initial_epoch = 0
29
    global_step = 0
30
    if config['preload']:
31
        model_filename = get_weights_file_path(config, config['preload'])
32
        print(f'Preloading model {model_filename}')
33
        state = torch.load(model_filename)
34
        initial_epoch = state['epoch'] + 1
35
        optimizer.load_state_dict(state['optimizer_state_dict'])
36
        global_step = state['global_step']
37

38
    # 定义loss函数，声明padding不参与loss计算。同时，使用label smoothing，让模型降低对计算结果的确定性，
39
    # 即减少本次推理结果的概率，并把减少的部分分配到其他可能的推理结果上。实测可以提升模型的泛化能力，降低过拟合。
40
    # 这里的label smoothing的系数设置为0.1，即将最高概率标签的概率降低0.1，再并分配给其他标签的概率。
41
    loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer_src.token_to_id('[PAD]'), label_smoothing=0.1).to(device)

1
    for epoch in range(initial_epoch, config['num_epochs']):
2
        model.train()
3
        # 画进度条
4
        batch_iterator = tqdm(train_dataloader, desc=f"Processing epoch {epoch:02d}")
5
        for batch in batch_iterator:
6

7
            encoder_input = batch['encoder_input'].to(device) # (batch, seq_len)
8
            decoder_input = batch['decoder_input'].to(device) # (batch, seq_len)
9
            encoder_mask = batch['encoder_mask'].to(device) # (batch, 1, 1, seq_len) 只隐藏padding tokens
10
            decoder_mask = batch['decoder_mask'].to(device) # (batch, 1, seq_len, seq_len) 隐藏padding tokens和未来的tokens
11

12
            # 将tensor输入Transformer模型并计算
13
            encoder_output = model.encode(encoder_input, encoder_mask) # (batch, seq_len, d_model)
14
            decoder_output = model.decode(encoder_output, encoder_mask, decoder_input, decoder_mask) # (batch, seq_len, d_model)
15
            proj_output = model.project(decoder_output) # (batch, seq_len, tgt_vocab_size)
16

17
            label = batch['label'].to(device) # (batch, seq_len)
18

19
            # (batch, seq_len, tgt_vocab_size) -[view]-> (batch * seq_len, tgt_vocab_size)
20
            # (batch, seq_len) -[view]-> (batch * seq_len)
21
            loss = loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))
22
            batch_iterator.set_postfix({f"loss": f"{loss.item():6.3f}"})
23

24
            # 为TensorBoard记录loss
25
            writer.add_scalar('train loss', loss.item(), global_step)
26
            writer.flush()
27

28
            # loss反向传播
29
            loss.backward()
30

31
            # 更新参数
32
            optimizer.step()
33
            optimizer.zero_grad()
34

35
            global_step += 1
36

37
        # 在每个epoch结束后保存模型
38
        # 保存模型是个好习惯，否则每次训练都需要从零开始
39
        model_filename = get_weights_file_path(config, f"{epoch:02d}")
40
        torch.save({
41
            'epoch': epoch,
42
            'model_state_dict': model.state_dict(),
43
            'optimizer_state_dict': optimizer.state_dict(),
44
            'global_step': global_step,
45
        }, model_filename)

1
import warnings
2

3
if __name__ == '__main__':
4
    warnings.filterwarnings('ignore')
5
    config = get_config()
6
    train_model(config)

1
def run_validation(model, validation_ds, tokenizer_src, tokenizer_tgt, max_len, device, print_msg, global_state, writer, num_examples=2):
2

3
    # 告诉pytorch将模型置于eval模式，关闭dropout、batchnorm中的随机因素
4
    model.eval()
5
    # 测试2个句子看看模型输出的翻译结果是什么
6
    count = 0
7

8
    srouce_texts = []
9
    expected = []
10
    predicted = []
11

12
    # 终端宽度
13
    console_width = 80
14

15
    # 关闭pytorch的梯度计算，在这个`with`中我们只需要推理结果，不需要训练，因此不需要计算tensor的梯度
16
    with torch.no_grad():
17
        for batch in validation_ds:
18
            count += 1
19
            encoder_input = batch['encoder_input'].to(device) # (batch, seq_len)
20
            encoder_mask = batch['encoder_mask'].to(device) # (batch, 1, 1, seq_len) 只隐藏padding tokens
21

22
            # 确定validation的batch应为1
23
            assert encoder_input.size(0) == 1, "Batch size must be 1 for validation"
24

25
            # 当模型进行推理时，只需要计算一次`encoder_output`，然后重复使用它来为每个token计算`decoder_output`

1
def greedy_decode(model, source, source_mask, tokenizer_src, tokenizer_tgt, max_len, device):
2
    """
3
    使用贪婪策略，缓存encoder输出，计算decoder输出
4

5
    参数:
6
        `model`: 模型
7
        `source`: 源语言句子
8
        `source_mask`: 源语言句子的padding mask
9
        `tokenizer_src`: 源语言的分词器
10
        `tokenizer_tgt`: 目标语言的分词器
11
        `max_len`: 句子最大长度
12
        `device`: 运行模型的设备
13
    """
14
    # 告诉pytorch将模型置于eval模式，关闭dropout、batchnorm中的随机因素
15
    sos_idx = tokenizer_tgt.token_to_id('[SOS]')
16
    eos_idx = tokenizer_tgt.token_to_id('[EOS]')
17

18
    # 计算encoder输出，并将结果重用于计算每个token的decoder输出
19
    encoder_output = model.encode(source, source_mask)
20

21
    # 模型推理计算过程如下：首先输入句子的起始符[SOS]作为decoder的第一个输入，然后得到模型输出的翻译句子的第一个词；
22
    # 接下在每一步迭代中将上一步的输出附加在输入中，使模型继续输出再下一个词，直到遇到[EOS]或达到句子的最大长度。
23

24
    # 使用[SOS]初始化decoder输入
25
    decoder_input = torch.empty(1, 1).fill_(sos_idx).type_as(source).to(device)
26
    # 接下来让decoder循环输出下一个词，知道遇到[EOS]或达到句子最大长度
27
    while True:
28

29
        # 判断是否达到句子最大长度
30
        if decoder_input.size(1) == max_len:
31
            break
32

33
        # 为target（即decoder input）创建mask，使模型只能看到当前时刻之前的词
34
        # 对比`dataset.py`中使用的`causal_mask`，这里使用的是`decoder_input.size(1)`作为`tgt_mask`的长度，并不需要使用额外的padding
35
        decoder_mask = causal_mask(decoder_input.size(1)).type_as(source_mask).to(device)
36

37
        # 计算decoder输出，这里复用了`encoder_output`
38
        out = model.decode(encoder_output, source_mask, decoder_input, decoder_mask) # (1, seq_len, d_model)
39

40
        # 计算下一个词：使用projection层计算下一个token的概率分布，只取最后一个词的projection，out[:, -1]为(1, 1, d_model)，prob为(1, 1, vocab_size)
41
        prob = model.project(out[:, -1])
42
        # 取概率最大的词作为下一个token（这里使用贪婪策略，即只取概率最大的词，而不是所有词的概率之和）
43
        _, next_word = torch.max(prob, dim=1)
44
        # 再将next_word添加到decoder_input的末尾
45
        decoder_input = torch.cat([decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(device)], dim=1)
46

47
        # 如果遇到[EOS]，则停止循环
48
        if next_word == eos_idx:
49
            break
50

51
    return decoder_input.squeeze(0) # 输出翻译句子

1
            # 确定validation的batch应为1
2
            assert encoder_input.size(0) == 1, "Batch size must be 1 for validation"
3

4
            # 当我们想要模型进行推理时，我们只需要计算一次`encoder_output`，然后重复使用它来为每个token计算`decoder_output`
5
            model_out = greedy_decode(model, encoder_input, encoder_mask, tokenizer_src, tokenizer_tgt, max_len, device)
6

7
            source_text = batch['src_text'][0]
8
            target_text = batch['tgt_text'][0]
9
            # 用目标语言的分词器将模型输出的token序列转换为文本
10
            model_out_text = tokenizer_tgt.decode(model_out.detach().cpu().numpy())
11

12
            source_texts.append(source_text)
13
            expected.append(target_text)
14
            predicted.append(model_out_text)
15

16
            # 打印输出，为了不影响`tqdm`进度条正常刷新，我们使用`print_msg`而不是`print`来打印输出。
17
            print_msg('-'*console_width)
18
            print_msg(f"Source text: {source_text}")
19
            print_msg(f"TARGET text: {target_text}")
20
            print_msg(f"PREDICTED text: {model_out_text}")
21
            print_msg('-'*console_width)
22

23
            # 如果预测了足够的句子，则停止循环
24
            if count >= num_examples:
25
                break

1
            global_step += 1
2

3
        run_validation(model, val_dataloader, tokenizer_src, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step, writer)
4

5
        # 在每个epoch结束后保存模型
6
        # 保存模型是个好习惯，否则每次训练都需要从零开始
7
        model_filename = get_weights_file_path(config, f"{epoch:02d}")

1
import torch
2
import torch.nn as nn
3

4
from model import Transformer
5
from config import get_config, get_weights_file_path
6
from train import get_model, get_ds, greedy_decode
7

8
import altair as alt
9
import pandas as pd
10
import numpy as np
11
import warnings
12

13

14
warnings.filterwarnings('ignore')

1
# 从`train_model`中拷贝选择设备的代码即可
2
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
3
print(f'Using device: {device}')

1
config = get_config()
2
train_dataloader, val_dataloader, vocab_src, vocab_tgt = get_ds(config)
3
model = get_model(config, vocab_src.get_vocab_size(), vocab_tgt.get_vocab_size()).to(device)
4

5
# 加载预训练权重
6
model_filename = get_weights_file_path(config, f"03")
7
state = torch.load(model_filename)
8
model.load_state_dict(state['model_state_dict'])

1
def load_next_batch():
2
    # Load a sample batch from the validation set
3
    batch = next(iter(val_dataloader))
4
    encoder_input = batch["encoder_input"].to(device)
5
    encoder_mask = batch["encoder_mask"].to(device)
6
    decoder_input = batch["decoder_input"].to(device)
7
    decoder_mask = batch["decoder_mask"].to(device)
8

9
    encoder_input_tokens = [vocab_src.id_to_token(idx) for idx in encoder_input[0].cpu().numpy()]
10
    decoder_input_tokens = [vocab_tgt.id_to_token(idx) for idx in decoder_input[0].cpu().numpy()]
11

12
    # check that the batch size is 1
13
    assert encoder_input.size(
14
        0) == 1, "Batch size must be 1 for validation"
15

16
    model_out = greedy_decode(
17
        model, encoder_input, encoder_mask, vocab_src, vocab_tgt, config['seq_len'], device)
18

19
    return batch, encoder_input_tokens, decoder_input_tokens

1
def mtx2df(m, max_row, max_col, row_tokens, col_tokens):
2
    """
3
    生成一个可视化矩阵
4
    """
5
    return pd.DataFrame(
6
        [
7
            (
8
                r,
9
                c,
10
                float(m[r, c]),
11
                "%.3d %s" % (r, row_tokens[r] if len(row_tokens) > r else "<blank>"),
12
                "%.3d %s" % (c, col_tokens[c] if len(col_tokens) > c else "<blank>"),
13
            )
14
            for r in range(m.shape[0])
15
            for c in range(m.shape[1])
16
            if r < max_row and c < max_col
17
        ],
18
        columns=["row", "column", "value", "row_token", "col_token"],
19
    )
20

21
def get_attn_map(attn_type: str, layer: int, head: int):
22
    """
23
    从指定的层和头获取attention map
24
    """
25

26
    # 我们在`model.py`的`MultiHeadAttentionBlock`中，定义`attention`函数的返回值时，同时返回了`attention_scores`，
27
    # 之后又在`forward`中将其保存在`self.attention_scores`中。
28
    # 因此，可以通过`model.encoder.layers[layer].self_attention_block.attention_scores`
29
    # 来获取encoder的self-attention的attention_scores。
30
    if attn_type == "encoder":
31
        attn = model.encoder.layers[layer].self_attention_block.attention_scores
32
    elif attn_type == "decoder":
33
        attn = model.decoder.layers[layer].self_attention_block.attention_scores
34
    elif attn_type == "encoder-decoder":
35
        attn = model.decoder.layers[layer].cross_attention_block.attention_scores
36
    return attn[0, head].data
37

38
def attn_map(attn_type, layer, head, row_tokens, col_tokens, max_sentence_len):
39
    """
40
    使用altair绘制attention map
41
    """
42
    df = mtx2df(
43
        get_attn_map(attn_type, layer, head),
44
        max_sentence_len,
45
        max_sentence_len,
46
        row_tokens,
47
        col_tokens,
48
    )
49
    return (
50
        alt.Chart(data=df)
51
        .mark_rect()
52
        .encode(
53
            x=alt.X("col_token", axis=alt.Axis(title="")),
54
            y=alt.Y("row_token", axis=alt.Axis(title="")),
55
            color="value",
56
            tooltip=["row", "column", "value", "row_token", "col_token"],
57
        )
58
        #.title(f"Layer {layer} Head {head}")
59
        .properties(height=400, width=400, title=f"Layer {layer} Head {head}")
60
        .interactive()
61
    )
62

63
def get_all_attention_maps(attn_type: str, layers: list[int], heads: list[int], row_tokens: list, col_tokens, max_sentence_len: int):
64
    charts = []
65
    for layer in layers:
66
        rowCharts = []
67
        for head in heads:
68
            rowCharts.append(attn_map(attn_type, layer, head, row_tokens, col_tokens, max_sentence_len))
69
        charts.append(alt.hconcat(*rowCharts))
70
    return alt.vconcat(*charts)

1
batch, encoder_input_tokens, decoder_input_tokens = load_next_batch()
2
print(f'Source: {batch["src_text"][0]}')
3
print(f'Target: {batch["tgt_text"][0]}')
4
sentence_len = encoder_input_tokens.index("[PAD]")

1
layers = [0, 1, 2]
2
heads = [0, 1, 2, 3, 4, 5, 6, 7]
3

4
# 可视化Encoder的Self-Attention
5
get_all_attention_maps("encoder", layers, heads, encoder_input_tokens, encoder_input_tokens, min(20, sentence_len))

1
# 可视化Decoder的Self-Attention
2
get_all_attention_maps("decoder", layers, heads, decoder_input_tokens, decoder_input_tokens, min(20, sentence_len))

1
# 可视化Decoder的Cross-Attention
2
# 这是实际上计算en-it翻译任务的地方，cross attention使用encoder的key/value，以及decoder的query计算attention。
3
get_all_attention_maps("encoder-decoder", layers, heads, encoder_input_tokens, decoder_input_tokens, min(20, sentence_len))