import math import torch import torch.nn as nn class Transformer(nn.Module): “””完全なTransformerモデル（Encoder-Decoder）””” def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, n_heads=8, n_layers=6, d_ff=2048, dropout=0.1, max_len=5000): “”” Args: src_vocab_size: 入力側の語彙サイズ（例: 10000） tgt_vocab_size: 出力側の語彙サイズ（例: 10000） d_model: モデル次元（例: 512） n_heads: ヘッド数（例: 8） n_layers: 層数（例: 6） d_ff: FFN中間層次元（例: 2048） dropout: ドロップアウト率 max_len: 最大系列長 “”” super(Transformer, self).__init__() # ========== 埋め込み層 ========== # 入力側: 単語ID → d_model次元のベクトル self.src_embedding = nn.Embedding(src_vocab_size, d_model) # 出力側: 単語ID → d_model次元のベクトル self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model) # ========== 位置エンコーディング ========== self.pos_encoding = PositionalEncoding(d_model, max_len, dropout) # ========== Encoder ========== self.encoder = Encoder(d_model, n_heads, d_ff, n_layers, dropout) # ========== Decoder ========== self.decoder = Decoder(d_model, n_heads, d_ff, n_layers, dropout) # ========== 出力層 ========== # d_model次元 → 語彙サイズ（各単語の確率） self.fc_out = nn.Linear(d_model, tgt_vocab_size) self.d_model = d_model self.dropout = nn.Dropout(dropout) # パラメータ初期化 self._init_parameters() def _init_parameters(self): “””Xavier初期化””” for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def forward(self, src, tgt, src_mask=None, tgt_mask=None): “”” Args: src: 入力系列 (batch, src_len) – 単語ID tgt: 出力系列 (batch, tgt_len) – 単語ID src_mask: 入力側マスク（Padding用） tgt_mask: 出力側マスク（Look-ahead用） Returns: output: (batch, tgt_len, tgt_vocab_size) – 各位置の単語確率 “”” # ========== Encoder ========== # 埋め込み（√d_modelでスケーリング） src_emb = self.src_embedding(src) * math.sqrt(self.d_model) # 位置エンコーディングを加算 src_emb = self.pos_encoding(src_emb) # Encoderを通す encoder_output = self.encoder(src_emb, src_mask) # ========== Decoder ========== # 埋め込み tgt_emb = self.tgt_embedding(tgt) * math.sqrt(self.d_model) # 位置エンコーディング tgt_emb = self.pos_encoding(tgt_emb) # Decoderを通す decoder_output = self.decoder(tgt_emb, encoder_output, src_mask, tgt_mask) # ========== 出力層 ========== # 語彙への射影 output = self.fc_out(decoder_output) return output def encode(self, src, src_mask=None): “””Encoderのみ実行（推論時に使用）””” src_emb = self.src_embedding(src) * math.sqrt(self.d_model) src_emb = self.pos_encoding(src_emb) return self.encoder(src_emb, src_mask) def decode(self, tgt, encoder_output, src_mask=None, tgt_mask=None): “””Decoderのみ実行（推論時に使用）””” tgt_emb = self.tgt_embedding(tgt) * math.sqrt(self.d_model) tgt_emb = self.pos_encoding(tgt_emb) decoder_output = self.decoder(tgt_emb, encoder_output, src_mask, tgt_mask) return self.fc_out(decoder_output)

import torch.optim as optim # デバイス設定 device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’) model = model.to(device) # Paddingトークンのインデックス PAD_IDX = 0 # 損失関数（Paddingを無視） criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX) # 最適化器（Adam） optimizer = optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9) def train_step(model, src, tgt, optimizer, criterion, device): “”” 1バッチの訓練 Args: src: 入力系列 (batch, src_len) tgt: 出力系列 (batch, tgt_len) – <START>〜<END>を含む “”” model.train() src = src.to(device) tgt = tgt.to(device) # Decoderの入力と教師ラベルを分ける # 入力: <START>〜最後の1つ前 # ラベル: 最初の1つ〜<END> tgt_input = tgt[:, :-1] tgt_output = tgt[:, 1:] # Look-aheadマスク生成 tgt_mask = generate_square_subsequent_mask(tgt_input.size(1)).to(device) # Forward optimizer.zero_grad() output = model(src, tgt_input, tgt_mask=tgt_mask) # 損失計算 # output: (batch, tgt_len-1, vocab_size) # tgt_output: (batch, tgt_len-1) # reshapeして2次元に loss = criterion( output.reshape(-1, output.size(-1)), # (batch * tgt_len-1, vocab_size) tgt_output.reshape(-1) # (batch * tgt_len-1) ) # Backward loss.backward() # 勾配クリッピング（勾配爆発を防ぐ） torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # パラメータ更新 optimizer.step() return loss.item()

def greedy_decode(model, src, start_token, end_token, max_len=50, device=’cpu’): “”” Greedy Searchによる生成 Args: model: Transformerモデル src: 入力系列 (1, src_len) start_token: 開始トークンのID end_token: 終了トークンのID max_len: 最大生成長 device: デバイス Returns: generated: 生成された系列（単語IDのリスト） “”” model.eval() src = src.to(device) # Encoderを1回だけ実行 with torch.no_grad(): encoder_output = model.encode(src) # Decoderの初期入力 generated = [start_token] for _ in range(max_len): # 現在の系列をテンソルに tgt = torch.LongTensor([generated]).to(device) # Look-aheadマスク tgt_mask = generate_square_subsequent_mask(tgt.size(1)).to(device) # Decoder実行 with torch.no_grad(): output = model.decode(tgt, encoder_output, tgt_mask=tgt_mask) # 最後の位置の予測 # output: (1, current_len, vocab_size) next_word_logits = output[0, -1, :] # (vocab_size) # 最も確率が高い単語を選択 next_word = next_word_logits.argmax().item() # 終了トークンなら終了 if next_word == end_token: break # 系列に追加 generated.append(next_word) return generated # 使用例（実際にはデータローダーから取得） START_TOKEN = 1 # <START>のID END_TOKEN = 2 # <END>のID # ダミーの入力 src = torch.randint(3, src_vocab_size, (1, 10)) # 生成 generated = greedy_decode(model, src, START_TOKEN, END_TOKEN, device=device) print(f”生成された系列: {generated}”)