import torch
import torch.nn as nn
 
class LSTMModel(nn.Module):
    """LSTM による系列分類"""
 
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.3,
            bidirectional=True,  # 双方向 LSTM
        )
        self.classifier = nn.Linear(hidden_dim * 2, num_classes)  # *2 for bidirectional
        self.dropout = nn.Dropout(0.3)
 
    def forward(self, x):
        embedded = self.dropout(self.embedding(x))  # [batch, seq_len, embed_dim]
        output, (hidden, cell) = self.lstm(embedded)
        # hidden: [num_layers*2, batch, hidden_dim]
        # 最後の層の forward と backward を結合
        hidden_cat = torch.cat((hidden[-2], hidden[-1]), dim=1)
        logits = self.classifier(self.dropout(hidden_cat))
        return logits
 
# 使用例
model = LSTMModel(vocab_size=30000, embed_dim=256, hidden_dim=512, num_classes=5)

import torch
import torch.nn as nn
import time
 
class GRUModel(nn.Module):
    """GRU による系列分類（LSTM より軽量）"""
 
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.gru = nn.GRU(
            embed_dim, hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=0.3,
            bidirectional=True,
        )
        self.classifier = nn.Linear(hidden_dim * 2, num_classes)
        self.dropout = nn.Dropout(0.3)
 
    def forward(self, x):
        embedded = self.dropout(self.embedding(x))
        output, hidden = self.gru(embedded)  # GRU: cell state なし
        hidden_cat = torch.cat((hidden[-2], hidden[-1]), dim=1)
        logits = self.classifier(self.dropout(hidden_cat))
        return logits
 
 
def compare_rnn_variants():
    """LSTM と GRU のパラメータ数・速度を比較"""
    vocab_size = 30000
    embed_dim = 256
    hidden_dim = 512
    num_classes = 5
    seq_len = 100
    batch_size = 32
 
    models = {
        "LSTM": LSTMModel(vocab_size, embed_dim, hidden_dim, num_classes),
        "GRU": GRUModel(vocab_size, embed_dim, hidden_dim, num_classes),
    }
 
    dummy_input = torch.randint(0, vocab_size, (batch_size, seq_len))
 
    for name, model in models.items():
        # パラメータ数
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
 
        # 推論速度
        model.eval()
        with torch.no_grad():
            start = time.time()
            for _ in range(100):
                _ = model(dummy_input)
            elapsed = time.time() - start
 
        print(f"{name}:")
        print(f"  パラメータ数: {total_params:,}")
        print(f"  学習可能: {trainable_params:,}")
        print(f"  推論時間 (100回): {elapsed:.3f}秒")
        print()
 
compare_rnn_variants()

import torch
import torch.nn as nn
import numpy as np
from torch.utils.data import Dataset, DataLoader
 
class TimeSeriesDataset(Dataset):
    """スライディングウィンドウで時系列データセットを作成"""
 
    def __init__(self, data, window_size, forecast_horizon=1):
        self.data = torch.FloatTensor(data)
        self.window_size = window_size
        self.forecast_horizon = forecast_horizon
 
    def __len__(self):
        return len(self.data) - self.window_size - self.forecast_horizon + 1
 
    def __getitem__(self, idx):
        x = self.data[idx:idx + self.window_size].unsqueeze(-1)  # [window, 1]
        y = self.data[idx + self.window_size:idx + self.window_size + self.forecast_horizon]
        return x, y
 
 
class LSTMForecaster(nn.Module):
    """LSTM による時系列予測モデル"""
 
    def __init__(self, input_dim=1, hidden_dim=64, num_layers=2,
                 forecast_horizon=1, dropout=0.2):
        super().__init__()
        self.lstm = nn.LSTM(
            input_dim, hidden_dim, num_layers=num_layers,
            batch_first=True, dropout=dropout
        )
        self.fc = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, forecast_horizon),
        )
 
    def forward(self, x):
        # x: [batch, window_size, input_dim]
        lstm_out, (h_n, c_n) = self.lstm(x)
        # 最後のタイムステップの隠れ状態を使用
        last_hidden = lstm_out[:, -1, :]  # [batch, hidden_dim]
        prediction = self.fc(last_hidden)  # [batch, forecast_horizon]
        return prediction
 
 
def train_forecaster(data, window_size=30, forecast_horizon=7,
                     epochs=50, lr=0.001, batch_size=32):
    """時系列予測モデルの学習"""
    # データ分割
    train_size = int(len(data) * 0.8)
    train_data = data[:train_size]
    val_data = data[train_size:]
 
    # 正規化
    mean, std = train_data.mean(), train_data.std()
    train_normalized = (train_data - mean) / std
    val_normalized = (val_data - mean) / std
 
    # データセット
    train_ds = TimeSeriesDataset(train_normalized, window_size, forecast_horizon)
    val_ds = TimeSeriesDataset(val_normalized, window_size, forecast_horizon)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size)
 
    # モデル
    model = LSTMForecaster(
        input_dim=1, hidden_dim=64, num_layers=2,
        forecast_horizon=forecast_horizon
    )
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, patience=5, factor=0.5
    )
    criterion = nn.MSELoss()
 
    best_val_loss = float('inf')
    patience_counter = 0
 
    for epoch in range(epochs):
        # 学習
        model.train()
        train_loss = 0
        for x_batch, y_batch in train_loader:
            optimizer.zero_grad()
            pred = model(x_batch)
            loss = criterion(pred, y_batch)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            train_loss += loss.item()
 
        # 検証
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for x_batch, y_batch in val_loader:
                pred = model(x_batch)
                val_loss += criterion(pred, y_batch).item()
 
        train_loss /= len(train_loader)
        val_loss /= len(val_loader)
        scheduler.step(val_loss)
 
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), "best_forecaster.pt")
            patience_counter = 0
        else:
            patience_counter += 1
 
        if patience_counter >= 10:
            print(f"Early stopping at epoch {epoch}")
            break
 
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}: train_loss={train_loss:.6f}, val_loss={val_loss:.6f}")
 
    return model, mean, std
 
# 使用例
# data = np.sin(np.linspace(0, 100, 1000)) + np.random.randn(1000) * 0.1
# model, mean, std = train_forecaster(data)

import torch
import torch.nn as nn
import random
 
class Encoder(nn.Module):
    """Seq2Seq エンコーダ"""
 
    def __init__(self, input_dim, embed_dim, hidden_dim, num_layers, dropout):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, embed_dim)
        self.rnn = nn.LSTM(embed_dim, hidden_dim, num_layers,
                           batch_first=True, dropout=dropout)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, src):
        embedded = self.dropout(self.embedding(src))  # [batch, src_len, embed_dim]
        outputs, (hidden, cell) = self.rnn(embedded)
        return hidden, cell
 
 
class Decoder(nn.Module):
    """Seq2Seq デコーダ"""
 
    def __init__(self, output_dim, embed_dim, hidden_dim, num_layers, dropout):
        super().__init__()
        self.embedding = nn.Embedding(output_dim, embed_dim)
        self.rnn = nn.LSTM(embed_dim, hidden_dim, num_layers,
                           batch_first=True, dropout=dropout)
        self.fc_out = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, input_token, hidden, cell):
        # input_token: [batch, 1]
        embedded = self.dropout(self.embedding(input_token))
        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))
        prediction = self.fc_out(output.squeeze(1))  # [batch, output_dim]
        return prediction, hidden, cell
 
 
class Seq2Seq(nn.Module):
    """Encoder-Decoder Seq2Seq モデル"""
 
    def __init__(self, encoder, decoder, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device
 
    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        batch_size = src.size(0)
        trg_len = trg.size(1)
        trg_vocab_size = self.decoder.fc_out.out_features
 
        outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)
        hidden, cell = self.encoder(src)
 
        # 最初の入力は <SOS> トークン
        input_token = trg[:, 0:1]  # [batch, 1]
 
        for t in range(1, trg_len):
            prediction, hidden, cell = self.decoder(input_token, hidden, cell)
            outputs[:, t] = prediction
 
            # Teacher Forcing: 一定確率で正解を次の入力にする
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = prediction.argmax(dim=1, keepdim=True)
            input_token = trg[:, t:t+1] if teacher_force else top1
 
        return outputs
 
# 使用例
INPUT_DIM = 10000   # ソース語彙サイズ
OUTPUT_DIM = 8000   # ターゲット語彙サイズ
EMBED_DIM = 256
HIDDEN_DIM = 512
NUM_LAYERS = 2
DROPOUT = 0.3
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
encoder = Encoder(INPUT_DIM, EMBED_DIM, HIDDEN_DIM, NUM_LAYERS, DROPOUT)
decoder = Decoder(OUTPUT_DIM, EMBED_DIM, HIDDEN_DIM, NUM_LAYERS, DROPOUT)
model = Seq2Seq(encoder, decoder, device).to(device)
 
print(f"Seq2Seq パラメータ数: {sum(p.numel() for p in model.parameters()):,}")

import torch
import torch.nn.functional as F
import math
 
def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: [batch, heads, seq_len, d_k]
    K: [batch, heads, seq_len, d_k]
    V: [batch, heads, seq_len, d_v]
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
 
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
 
    attention_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, V)
    return output, attention_weights
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
 
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
 
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
 
        # 線形変換 + マルチヘッド分割
        Q = self.W_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
 
        # Attention 計算
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
 
        # ヘッド結合 + 出力変換
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
        output = self.W_o(attn_output)
        return output

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
 
def visualize_attention(tokens, attention_weights, layer=0, head=0, save_path=None):
    """
    Attention の重みをヒートマップで可視化
 
    Args:
        tokens: トークンのリスト
        attention_weights: [layers, heads, seq_len, seq_len] のテンソル
        layer: 可視化する層
        head: 可視化するヘッド
    """
    attn = attention_weights[layer][head].detach().cpu().numpy()
 
    fig, ax = plt.subplots(figsize=(10, 10))
    im = ax.imshow(attn, cmap="viridis", vmin=0, vmax=1)
 
    ax.set_xticks(range(len(tokens)))
    ax.set_yticks(range(len(tokens)))
    ax.set_xticklabels(tokens, rotation=45, ha="right", fontsize=10)
    ax.set_yticklabels(tokens, fontsize=10)
 
    # 各セルに値を表示
    for i in range(len(tokens)):
        for j in range(len(tokens)):
            ax.text(j, i, f"{attn[i, j]:.2f}",
                    ha="center", va="center", fontsize=8,
                    color="white" if attn[i, j] > 0.5 else "black")
 
    ax.set_xlabel("Key")
    ax.set_ylabel("Query")
    ax.set_title(f"Attention Weights (Layer {layer}, Head {head})")
    plt.colorbar(im, ax=ax)
    plt.tight_layout()
 
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()
 
 
def extract_bert_attention(model, tokenizer, text):
    """BERT から Attention weights を抽出"""
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs, output_attentions=True)
 
    # attentions: tuple of (batch, heads, seq_len, seq_len) per layer
    attentions = outputs.attentions
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
 
    # 各層のAttentionを可視化
    all_attentions = torch.stack([a.squeeze(0) for a in attentions])
    return tokens, all_attentions
 
 
def attention_rollout(attentions, head_fusion="mean"):
    """
    Attention Rollout: 複数層のAttentionを統合して
    入力トークンへの最終的な寄与度を計算
 
    Args:
        attentions: [num_layers, num_heads, seq_len, seq_len]
        head_fusion: ヘッドの統合方法 ("mean", "max", "min")
    """
    num_layers = attentions.shape[0]
    seq_len = attentions.shape[-1]
 
    # 残差接続を考慮してAttentionに単位行列を加算
    result = torch.eye(seq_len)
 
    for layer in range(num_layers):
        if head_fusion == "mean":
            attn = attentions[layer].mean(dim=0)
        elif head_fusion == "max":
            attn = attentions[layer].max(dim=0).values
        elif head_fusion == "min":
            attn = attentions[layer].min(dim=0).values
 
        # 残差接続の影響
        attn = 0.5 * attn + 0.5 * torch.eye(seq_len)
        # 行方向に正規化
        attn = attn / attn.sum(dim=-1, keepdim=True)
        result = torch.matmul(attn, result)
 
    return result

import torch
import torch.nn as nn
 
class BahdanauAttention(nn.Module):
    """Additive Attention (Bahdanau et al., 2014)"""
 
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        super().__init__()
        self.W_encoder = nn.Linear(encoder_dim, attention_dim, bias=False)
        self.W_decoder = nn.Linear(decoder_dim, attention_dim, bias=False)
        self.v = nn.Linear(attention_dim, 1, bias=False)
 
    def forward(self, encoder_outputs, decoder_hidden):
        """
        encoder_outputs: [batch, src_len, encoder_dim]
        decoder_hidden: [batch, decoder_dim]
        """
        # decoder_hidden を src_len 分繰り返す
        decoder_hidden = decoder_hidden.unsqueeze(1)  # [batch, 1, decoder_dim]
 
        # スコア計算
        energy = torch.tanh(
            self.W_encoder(encoder_outputs) + self.W_decoder(decoder_hidden)
        )  # [batch, src_len, attention_dim]
 
        scores = self.v(energy).squeeze(-1)  # [batch, src_len]
        attention_weights = torch.softmax(scores, dim=-1)  # [batch, src_len]
 
        # コンテキストベクトル
        context = torch.bmm(
            attention_weights.unsqueeze(1), encoder_outputs
        ).squeeze(1)  # [batch, encoder_dim]
 
        return context, attention_weights
 
 
class AttentionDecoder(nn.Module):
    """Attention 付き Decoder"""
 
    def __init__(self, output_dim, embed_dim, encoder_dim,
                 decoder_dim, attention_dim, dropout=0.3):
        super().__init__()
        self.attention = BahdanauAttention(encoder_dim, decoder_dim, attention_dim)
        self.embedding = nn.Embedding(output_dim, embed_dim)
        self.rnn = nn.GRU(embed_dim + encoder_dim, decoder_dim, batch_first=True)
        self.fc_out = nn.Linear(decoder_dim + encoder_dim + embed_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, input_token, decoder_hidden, encoder_outputs):
        """
        input_token: [batch]
        decoder_hidden: [1, batch, decoder_dim]
        encoder_outputs: [batch, src_len, encoder_dim]
        """
        embedded = self.dropout(self.embedding(input_token.unsqueeze(1)))  # [batch, 1, embed_dim]
 
        context, attn_weights = self.attention(
            encoder_outputs, decoder_hidden.squeeze(0)
        )
        context = context.unsqueeze(1)  # [batch, 1, encoder_dim]
 
        rnn_input = torch.cat([embedded, context], dim=2)  # [batch, 1, embed_dim + encoder_dim]
        output, hidden = self.rnn(rnn_input, decoder_hidden)
 
        prediction = self.fc_out(
            torch.cat([output.squeeze(1), context.squeeze(1), embedded.squeeze(1)], dim=1)
        )  # [batch, output_dim]
 
        return prediction, hidden, attn_weights

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, mask=None):
        # Self-Attention + Residual + LayerNorm
        attn_out = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))
 
        # Feed-Forward + Residual + LayerNorm
        ff_out = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_out))
        return x
 
class TextClassifier(nn.Module):
    """Transformer ベースのテキスト分類器"""
 
    def __init__(self, vocab_size, d_model=512, num_heads=8,
                 num_layers=6, d_ff=2048, num_classes=5, max_len=512):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = nn.Embedding(max_len, d_model)
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ])
        self.classifier = nn.Linear(d_model, num_classes)
        self.dropout = nn.Dropout(0.1)
 
    def forward(self, x):
        seq_len = x.size(1)
        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)
        x = self.dropout(self.embedding(x) + self.pos_encoding(positions))
 
        for layer in self.layers:
            x = layer(x)
 
        # [CLS] トークンの表現を分類に使用
        cls_output = x[:, 0]
        return self.classifier(cls_output)

import torch
import torch.nn as nn
 
class TransformerDecoderLayer(nn.Module):
    """Transformer Decoder Layer（Causal Attention + Cross-Attention）"""
 
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        # Masked Self-Attention
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        # Cross-Attention (Encoder-Decoder)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        # Feed-Forward
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, encoder_output, src_mask=None, tgt_mask=None):
        # Masked Self-Attention
        attn_out = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_out))
 
        # Cross-Attention
        cross_out = self.cross_attn(x, encoder_output, encoder_output, src_mask)
        x = self.norm2(x + self.dropout(cross_out))
 
        # Feed-Forward
        ff_out = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_out))
        return x
 
 
class TransformerModel(nn.Module):
    """完全な Transformer (Encoder + Decoder)"""
 
    def __init__(self, src_vocab, tgt_vocab, d_model=512, num_heads=8,
                 num_encoder_layers=6, num_decoder_layers=6,
                 d_ff=2048, max_len=512, dropout=0.1):
        super().__init__()
        self.d_model = d_model
 
        # Embeddings
        self.src_embedding = nn.Embedding(src_vocab, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab, d_model)
        self.pos_encoding = SinusoidalPositionalEncoding(d_model, max_len)
 
        # Encoder
        self.encoder_layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_encoder_layers)
        ])
 
        # Decoder
        self.decoder_layers = nn.ModuleList([
            TransformerDecoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_decoder_layers)
        ])
 
        self.output_proj = nn.Linear(d_model, tgt_vocab)
        self.dropout = nn.Dropout(dropout)
 
    def generate_causal_mask(self, seq_len, device):
        """未来のトークンをマスクする Causal Mask を生成"""
        mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1)
        mask = mask.bool()
        return ~mask  # True = attend, False = mask
 
    def encode(self, src):
        x = self.dropout(self.pos_encoding(
            self.src_embedding(src) * (self.d_model ** 0.5)
        ))
        for layer in self.encoder_layers:
            x = layer(x)
        return x
 
    def decode(self, tgt, encoder_output, tgt_mask=None):
        x = self.dropout(self.pos_encoding(
            self.tgt_embedding(tgt) * (self.d_model ** 0.5)
        ))
        for layer in self.decoder_layers:
            x = layer(x, encoder_output, tgt_mask=tgt_mask)
        return self.output_proj(x)
 
    def forward(self, src, tgt):
        tgt_mask = self.generate_causal_mask(tgt.size(1), tgt.device)
        encoder_output = self.encode(src)
        output = self.decode(tgt, encoder_output, tgt_mask)
        return output

from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
 
# 事前学習済みモデルのロード
tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
model = BertForSequenceClassification.from_pretrained(
    'bert-base-multilingual-cased',
    num_labels=3,
)
 
# テキストのトークナイズ
texts = ["この映画は素晴らしい", "つまらない作品だった"]
inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
 
# 推論
outputs = model(**inputs)
predictions = outputs.logits.argmax(dim=-1)
 
# ファインチューニング
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_steps=500,
)
 
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    get_linear_schedule_with_warmup
)
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import numpy as np
 
class TextDataset(Dataset):
    """テキスト分類用データセット"""
 
    def __init__(self, texts, labels, tokenizer, max_length=256):
        self.encodings = tokenizer(
            texts, truncation=True, padding=True,
            max_length=max_length, return_tensors="pt"
        )
        self.labels = torch.tensor(labels, dtype=torch.long)
 
    def __len__(self):
        return len(self.labels)
 
    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item["labels"] = self.labels[idx]
        return item
 
 
class BERTFineTuner:
    """BERT ファインチューニングの完全パイプライン"""
 
    def __init__(self, model_name="bert-base-multilingual-cased",
                 num_labels=3, max_length=256, device=None):
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=num_labels
        ).to(self.device)
        self.max_length = max_length
 
    def prepare_data(self, texts, labels, test_size=0.2, batch_size=16):
        """データの分割とDataLoader作成"""
        train_texts, val_texts, train_labels, val_labels = train_test_split(
            texts, labels, test_size=test_size, stratify=labels, random_state=42
        )
 
        self.train_dataset = TextDataset(
            train_texts, train_labels, self.tokenizer, self.max_length
        )
        self.val_dataset = TextDataset(
            val_texts, val_labels, self.tokenizer, self.max_length
        )
 
        self.train_loader = DataLoader(
            self.train_dataset, batch_size=batch_size, shuffle=True
        )
        self.val_loader = DataLoader(
            self.val_dataset, batch_size=batch_size
        )
 
    def train(self, epochs=3, lr=2e-5, warmup_ratio=0.1,
              weight_decay=0.01, max_grad_norm=1.0):
        """学習ループ"""
        # オプティマイザ（bias と LayerNorm に weight_decay を適用しない）
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in self.model.named_parameters()
                          if not any(nd in n for nd in no_decay)],
                "weight_decay": weight_decay,
            },
            {
                "params": [p for n, p in self.model.named_parameters()
                          if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
            },
        ]
        optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=lr)
 
        # スケジューラ
        total_steps = len(self.train_loader) * epochs
        warmup_steps = int(total_steps * warmup_ratio)
        scheduler = get_linear_schedule_with_warmup(
            optimizer, warmup_steps, total_steps
        )
 
        best_val_acc = 0.0
 
        for epoch in range(epochs):
            # 学習フェーズ
            self.model.train()
            total_loss = 0
            correct = 0
            total = 0
 
            for batch in self.train_loader:
                batch = {k: v.to(self.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                loss = outputs.loss
 
                loss.backward()
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), max_grad_norm
                )
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
 
                total_loss += loss.item()
                preds = outputs.logits.argmax(dim=-1)
                correct += (preds == batch["labels"]).sum().item()
                total += len(batch["labels"])
 
            train_acc = correct / total
            avg_loss = total_loss / len(self.train_loader)
 
            # 検証フェーズ
            val_acc, val_report = self.evaluate()
 
            print(f"Epoch {epoch+1}/{epochs}")
            print(f"  Train Loss: {avg_loss:.4f}, Train Acc: {train_acc:.4f}")
            print(f"  Val Acc: {val_acc:.4f}")
 
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                torch.save(self.model.state_dict(), "best_bert_model.pt")
                print(f"  -> Best model saved (acc={val_acc:.4f})")
 
        return best_val_acc
 
    def evaluate(self):
        """検証データで評価"""
        self.model.eval()
        all_preds = []
        all_labels = []
 
        with torch.no_grad():
            for batch in self.val_loader:
                batch = {k: v.to(self.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                preds = outputs.logits.argmax(dim=-1)
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(batch["labels"].cpu().numpy())
 
        acc = np.mean(np.array(all_preds) == np.array(all_labels))
        report = classification_report(all_labels, all_preds)
        return acc, report
 
    def predict(self, texts):
        """新しいテキストの推論"""
        self.model.eval()
        inputs = self.tokenizer(
            texts, truncation=True, padding=True,
            max_length=self.max_length, return_tensors="pt"
        ).to(self.device)
 
        with torch.no_grad():
            outputs = self.model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)
            preds = probs.argmax(dim=-1)
 
        return preds.cpu().numpy(), probs.cpu().numpy()
 
 
# 使用例
# fine_tuner = BERTFineTuner(num_labels=3)
# fine_tuner.prepare_data(texts, labels, batch_size=16)
# best_acc = fine_tuner.train(epochs=3, lr=2e-5)
# preds, probs = fine_tuner.predict(["この映画は最高だった"])

class SinusoidalPositionalEncoding(nn.Module):
    """Transformer 原論文の正弦波位置エンコーディング"""
 
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))
 
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

import torch
import torch.nn as nn
 
class RotaryPositionalEmbedding(nn.Module):
    """
    Rotary Positional Embedding (RoPE)
    - LLaMA、GPT-NeoX 等で採用
    - 相対位置情報を回転行列で埋め込む
    - 長さの外挿（extrapolation）に強い
    """
 
    def __init__(self, dim, max_seq_len=4096, base=10000):
        super().__init__()
        self.dim = dim
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
 
        # 事前計算
        t = torch.arange(max_seq_len).float()
        freqs = torch.outer(t, inv_freq)  # [max_seq_len, dim/2]
        emb = torch.cat([freqs, freqs], dim=-1)  # [max_seq_len, dim]
        self.register_buffer("cos_cached", emb.cos())
        self.register_buffer("sin_cached", emb.sin())
 
    def forward(self, x, seq_len=None):
        # x: [batch, heads, seq_len, dim]
        if seq_len is None:
            seq_len = x.shape[2]
        return (
            self.cos_cached[:seq_len].unsqueeze(0).unsqueeze(0),
            self.sin_cached[:seq_len].unsqueeze(0).unsqueeze(0),
        )
 
 
def rotate_half(x):
    """次元の前半と後半を入れ替えて符号反転"""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)
 
 
def apply_rotary_pos_emb(q, k, cos, sin):
    """Q, K に RoPE を適用"""
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed
 
 
class RoPEMultiHeadAttention(nn.Module):
    """RoPE を使用した Multi-Head Attention"""
 
    def __init__(self, d_model, num_heads, max_seq_len=4096):
        super().__init__()
        self.d_k = d_model // num_heads
        self.num_heads = num_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.rope = RotaryPositionalEmbedding(self.d_k, max_seq_len)
 
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape
 
        Q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
 
        # RoPE 適用
        cos, sin = self.rope(Q, seq_len)
        Q, K = apply_rotary_pos_emb(Q, K, cos, sin)
 
        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn, V)
 
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        return self.W_o(output)

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
class ModelOptimizer:
    """Transformer モデルの最適化ツール"""
 
    @staticmethod
    def quantize_dynamic(model):
        """動的量子化（CPU 推論向け）"""
        quantized = torch.quantization.quantize_dynamic(
            model, {torch.nn.Linear}, dtype=torch.qint8
        )
        return quantized
 
    @staticmethod
    def compare_model_sizes(original, quantized):
        """モデルサイズの比較"""
        import tempfile
        import os
 
        with tempfile.NamedTemporaryFile(suffix=".pt") as f:
            torch.save(original.state_dict(), f.name)
            original_size = os.path.getsize(f.name)
 
        with tempfile.NamedTemporaryFile(suffix=".pt") as f:
            torch.save(quantized.state_dict(), f.name)
            quantized_size = os.path.getsize(f.name)
 
        print(f"元のモデル: {original_size / 1e6:.1f} MB")
        print(f"量子化後: {quantized_size / 1e6:.1f} MB")
        print(f"圧縮率: {quantized_size / original_size:.1%}")
 
    @staticmethod
    def benchmark_inference(model, tokenizer, texts, num_runs=100):
        """推論速度のベンチマーク"""
        import time
 
        model.eval()
        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
 
        # ウォームアップ
        with torch.no_grad():
            for _ in range(10):
                _ = model(**inputs)
 
        # 計測
        start = time.time()
        with torch.no_grad():
            for _ in range(num_runs):
                _ = model(**inputs)
        elapsed = time.time() - start
 
        print(f"平均推論時間: {elapsed / num_runs * 1000:.2f} ms")
        print(f"スループット: {num_runs / elapsed:.1f} 推論/秒")
 
    @staticmethod
    def export_onnx(model, tokenizer, output_path="model.onnx", max_length=128):
        """ONNX 形式でエクスポート"""
        model.eval()
        dummy_input = tokenizer(
            "サンプルテキスト", return_tensors="pt",
            padding="max_length", max_length=max_length, truncation=True
        )
 
        torch.onnx.export(
            model,
            (dummy_input["input_ids"], dummy_input["attention_mask"]),
            output_path,
            input_names=["input_ids", "attention_mask"],
            output_names=["logits"],
            dynamic_axes={
                "input_ids": {0: "batch_size"},
                "attention_mask": {0: "batch_size"},
                "logits": {0: "batch_size"},
            },
            opset_version=14,
        )
        print(f"ONNX モデル出力: {output_path}")
 
 
# KV Cache の実装
class KVCache:
    """
    Key-Value Cache: 自己回帰生成の高速化
    - 既に計算した K, V を再利用
    - 生成ステップごとの計算量を O(n) に削減
    """
 
    def __init__(self, max_batch_size=1, max_seq_len=2048,
                 num_heads=8, head_dim=64, num_layers=12):
        self.max_seq_len = max_seq_len
        self.cache = {}
        for layer in range(num_layers):
            self.cache[layer] = {
                "key": torch.zeros(max_batch_size, num_heads, max_seq_len, head_dim),
                "value": torch.zeros(max_batch_size, num_heads, max_seq_len, head_dim),
            }
        self.current_len = 0
 
    def update(self, layer, key, value):
        """新しい K, V をキャッシュに追加"""
        new_len = key.shape[2]
        self.cache[layer]["key"][:, :, self.current_len:self.current_len + new_len] = key
        self.cache[layer]["value"][:, :, self.current_len:self.current_len + new_len] = value
 
    def get(self, layer):
        """キャッシュから K, V を取得"""
        return (
            self.cache[layer]["key"][:, :, :self.current_len + 1],
            self.cache[layer]["value"][:, :, :self.current_len + 1],
        )
 
    def advance(self):
        """位置を1つ進める"""
        self.current_len += 1
 
    def reset(self):
        """キャッシュをリセット"""
        self.current_len = 0
        for layer in self.cache:
            self.cache[layer]["key"].zero_()
            self.cache[layer]["value"].zero_()

import torch
import torch.nn as nn
import math
 
class LoRALayer(nn.Module):
    """
    LoRA: Low-Rank Adaptation
    - 元の重み行列 W に低ランク行列 A*B を加算
    - W' = W + alpha/r * A * B
    - 学習パラメータ数を大幅に削減
    """
 
    def __init__(self, original_layer, rank=8, alpha=16):
        super().__init__()
        self.original_layer = original_layer
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
 
        in_features = original_layer.in_features
        out_features = original_layer.out_features
 
        # 低ランク行列
        self.lora_A = nn.Parameter(torch.zeros(in_features, rank))
        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))
 
        # 初期化
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)
 
        # 元の重みをフリーズ
        for param in self.original_layer.parameters():
            param.requires_grad = False
 
    def forward(self, x):
        original_output = self.original_layer(x)
        lora_output = (x @ self.lora_A @ self.lora_B) * self.scaling
        return original_output + lora_output
 
 
def apply_lora_to_model(model, rank=8, alpha=16, target_modules=None):
    """
    モデルの特定の Linear 層に LoRA を適用
 
    Args:
        model: 対象モデル
        rank: LoRA のランク
        alpha: スケーリングファクター
        target_modules: LoRA を適用するモジュール名のリスト
    """
    if target_modules is None:
        target_modules = ["query", "key", "value", "dense"]
 
    total_params = 0
    lora_params = 0
 
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            total_params += module.weight.numel()
 
            if any(target in name for target in target_modules):
                # LoRA レイヤーで置換
                parent_name = ".".join(name.split(".")[:-1])
                child_name = name.split(".")[-1]
 
                parent = model
                for part in parent_name.split("."):
                    if part:
                        parent = getattr(parent, part)
 
                lora_layer = LoRALayer(module, rank=rank, alpha=alpha)
                setattr(parent, child_name, lora_layer)
                lora_params += rank * (module.in_features + module.out_features)
 
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad)
 
    print(f"元のパラメータ数: {total_params:,}")
    print(f"LoRA パラメータ数: {lora_params:,}")
    print(f"学習可能パラメータ数: {trainable:,} ({trainable/(trainable+frozen):.2%})")
    print(f"フリーズパラメータ数: {frozen:,}")
 
    return model
 
# 使用例
# from transformers import AutoModelForSequenceClassification
# model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
# model = apply_lora_to_model(model, rank=8, alpha=16)

import torch
import torch.nn as nn
 
class SelectiveSSM(nn.Module):
    """
    Selective State Space Model (Mamba の簡略実装)
    - 入力依存のゲーティングで選択的に情報を保持
    - RNN のような逐次処理が可能（推論時は高速）
    - Transformer のような並列学習も可能
    """
 
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_inner = int(expand * d_model)
 
        # 入力射影
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
 
        # 1D 畳み込み（局所依存性）
        self.conv1d = nn.Conv1d(
            self.d_inner, self.d_inner,
            kernel_size=d_conv, padding=d_conv - 1,
            groups=self.d_inner
        )
 
        # SSM パラメータ
        self.x_proj = nn.Linear(self.d_inner, d_state * 2 + 1, bias=False)
        self.dt_proj = nn.Linear(1, self.d_inner, bias=True)
 
        # A パラメータ（対角行列として初期化）
        A = torch.arange(1, d_state + 1, dtype=torch.float32)
        self.A_log = nn.Parameter(torch.log(A).unsqueeze(0).expand(self.d_inner, -1))
 
        self.D = nn.Parameter(torch.ones(self.d_inner))
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
 
    def forward(self, x):
        """
        x: [batch, seq_len, d_model]
        """
        batch, seq_len, _ = x.shape
 
        # 入力射影 → x, z
        xz = self.in_proj(x)  # [batch, seq_len, d_inner * 2]
        x_branch, z = xz.chunk(2, dim=-1)
 
        # 1D 畳み込み
        x_conv = self.conv1d(x_branch.transpose(1, 2))[:, :, :seq_len].transpose(1, 2)
        x_conv = torch.silu(x_conv)
 
        # SSM パラメータの計算（入力依存）
        x_ssm = self.x_proj(x_conv)  # [batch, seq_len, d_state*2 + 1]
        B = x_ssm[:, :, :self.d_state]
        C = x_ssm[:, :, self.d_state:self.d_state*2]
        dt = torch.softplus(self.dt_proj(x_ssm[:, :, -1:]))
 
        # A の離散化
        A = -torch.exp(self.A_log)  # [d_inner, d_state]
 
        # SSM の逐次計算（簡略版）
        y = self._ssm_scan(x_conv, A, B, C, dt)
 
        # ゲーティング
        y = y * torch.silu(z)
        return self.out_proj(y)
 
    def _ssm_scan(self, x, A, B, C, dt):
        """SSM の逐次スキャン"""
        batch, seq_len, d_inner = x.shape
        d_state = A.shape[1]
 
        h = torch.zeros(batch, d_inner, d_state, device=x.device)
        outputs = []
 
        for t in range(seq_len):
            # 状態更新: h = A_bar * h + B_bar * x
            dt_t = dt[:, t].unsqueeze(-1)  # [batch, d_inner, 1]
            A_bar = torch.exp(A.unsqueeze(0) * dt_t)  # [batch, d_inner, d_state]
            B_bar = dt_t * B[:, t].unsqueeze(1)  # [batch, d_inner, d_state]
 
            h = A_bar * h + B_bar * x[:, t].unsqueeze(-1)
 
            # 出力: y = C * h
            y_t = (h * C[:, t].unsqueeze(1)).sum(dim=-1)  # [batch, d_inner]
            outputs.append(y_t)
 
        return torch.stack(outputs, dim=1)  # [batch, seq_len, d_inner]
 
 
class MambaBlock(nn.Module):
    """Mamba Block (SSM + Skip Connection + Norm)"""
 
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.ssm = SelectiveSSM(d_model, d_state, d_conv, expand)
 
    def forward(self, x):
        return x + self.ssm(self.norm(x))

import torch
import torch.nn.functional as F
import math
 
def flash_attention_reference(Q, K, V, block_size=256):
    """
    Flash Attention の概念実装（参考用）
    - 実際の Flash Attention は CUDA カーネルで実装
    - ここではアルゴリズムの理解のための参考実装
 
    Key Idea:
    - Attention を小ブロックに分割して計算
    - SRAM（高速メモリ）に収まるサイズで処理
    - Softmax をオンラインで計算（log-sum-exp トリック）
    """
    batch, heads, seq_len, d_k = Q.shape
 
    # 出力とログサムエクスプの初期化
    O = torch.zeros_like(V)
    L = torch.zeros(batch, heads, seq_len, 1, device=Q.device)  # log-sum-exp
    M = torch.full((batch, heads, seq_len, 1), float('-inf'), device=Q.device)  # max
 
    num_blocks = math.ceil(seq_len / block_size)
 
    for j in range(num_blocks):
        # K, V のブロック
        j_start = j * block_size
        j_end = min((j + 1) * block_size, seq_len)
        K_block = K[:, :, j_start:j_end]
        V_block = V[:, :, j_start:j_end]
 
        for i in range(num_blocks):
            # Q のブロック
            i_start = i * block_size
            i_end = min((i + 1) * block_size, seq_len)
            Q_block = Q[:, :, i_start:i_end]
 
            # ブロック間の Attention スコア
            S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / math.sqrt(d_k)
 
            # オンライン Softmax
            M_block = S_block.max(dim=-1, keepdim=True).values
            M_new = torch.max(M[:, :, i_start:i_end], M_block)
 
            # 指数関数の安定計算
            exp_old = torch.exp(M[:, :, i_start:i_end] - M_new)
            exp_new = torch.exp(S_block - M_new)
 
            L_new = exp_old * L[:, :, i_start:i_end] + exp_new.sum(dim=-1, keepdim=True)
 
            # 出力の更新
            O[:, :, i_start:i_end] = (
                exp_old * O[:, :, i_start:i_end] +
                torch.matmul(exp_new, V_block)
            )
 
            M[:, :, i_start:i_end] = M_new
            L[:, :, i_start:i_end] = L_new
 
    # 正規化
    O = O / L
    return O
 
 
# PyTorch 2.0+ では torch.nn.functional.scaled_dot_product_attention を使用
def efficient_attention_pytorch2(Q, K, V, is_causal=False):
    """PyTorch 2.0 の Flash Attention API"""
    # 自動的に Flash Attention または Memory-Efficient Attention を選択
    return F.scaled_dot_product_attention(
        Q, K, V,
        is_causal=is_causal,
        # dropout_p=0.0,  # 推論時は0
    )

import torch
from torch.cuda.amp import autocast, GradScaler
 
class MixedPrecisionTrainer:
    """Mixed Precision Training でメモリ使用量を削減し高速化"""
 
    def __init__(self, model, optimizer, criterion):
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.scaler = GradScaler()
 
    def train_step(self, batch):
        self.optimizer.zero_grad()
 
        with autocast():
            outputs = self.model(**batch)
            loss = outputs.loss if hasattr(outputs, 'loss') else self.criterion(
                outputs, batch["labels"]
            )
 
        # Scaled backward
        self.scaler.scale(loss).backward()
 
        # Gradient clipping (unscale first)
        self.scaler.unscale_(self.optimizer)
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
 
        # Optimizer step
        self.scaler.step(self.optimizer)
        self.scaler.update()
 
        return loss.item()
 
    def train_epoch(self, dataloader, accumulation_steps=4):
        """Gradient Accumulation 付きの学習エポック"""
        self.model.train()
        total_loss = 0
 
        for step, batch in enumerate(dataloader):
            batch = {k: v.cuda() for k, v in batch.items()}
 
            with autocast():
                outputs = self.model(**batch)
                loss = outputs.loss / accumulation_steps
 
            self.scaler.scale(loss).backward()
 
            if (step + 1) % accumulation_steps == 0:
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()
 
            total_loss += loss.item() * accumulation_steps
 
        return total_loss / len(dataloader)