transformer 翻译可视化代码

以下是一个简单的Transformer翻译可视化代码示例：

import torch
import torch.nn.functional as F

class TransformerTranslator:
    def __init__(self, model, src_vocab, tgt_vocab):
        self.model = model
        self.src_vocab = src_vocab
        self.tgt_vocab = tgt_vocab

    def translate_sentence(self, sentence):
        # 将输入句子转换为张量
        src_indexes = [self.src_vocab.stoi[word] for word in sentence.split()]
        src_tensor = torch.LongTensor(src_indexes).unsqueeze(0)

        # 使用模型进行翻译
        with torch.no_grad():
            enc_src = self.model.encoder(src_tensor)
            output = [self.tgt_vocab.stoi["<sos>"]]
            for i in range(10):  # 限制输出到最大长度为10个单词
                trg_tensor = torch.LongTensor(output).unsqueeze(0)
                with torch.no_grad():
                    output_logits, _ = self.model.decoder(trg_tensor, enc_src)
                pred_token = output_logits.argmax(2)[:, -1].item()
                output.append(pred_token)
                if pred_token == self.tgt_vocab.stoi["<eos>"]:
                    break

        # 将输出单词转换为句子
        translation = [self.tgt_vocab.itos[idx] for idx in output]
        translation = ' '.join(translation[1:-1])  # 去掉<sos>和<eos>
        return translation

在上面的代码中，我们首先将输入句子转换为源语言的单词索引序列，并将其传递给Transformer模型进行编码。然后，我们使用解码器逐个生成目标语言的单词，并将其添加到输出序列中，直到生成结束标记或达到输出序列的最大长度为止。最后，我们将输出序列转换回目标语言的单词，并将它们连接起来以形成翻译后的句子。