# Multi-Head Attentionの実装例（PyTorch）
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        self.depth = d_model // num_heads
        
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.dense = nn.Linear(d_model, d_model)
        
    def scaled_dot_product_attention(self, q, k, v, mask=None):
        matmul_qk = torch.matmul(q, k.transpose(-2, -1))
        dk = torch.tensor(k.shape[-1], dtype=torch.float32)
        scaled_attention_logits = matmul_qk / math.sqrt(dk)
        
        if mask is not None:
            scaled_attention_logits += (mask * -1e9)
            
        attention_weights = torch.softmax(scaled_attention_logits, dim=-1)
        output = torch.matmul(attention_weights, v)
        return output, attention_weights

# Positional Encodingの実装例
import numpy as np
import torch

def positional_encoding(seq_len, d_model):
    pos = np.arange(seq_len)[:, np.newaxis]
    i = np.arange(d_model)[np.newaxis, :]
    
    # 角度計算
    angles = pos / np.power(10000, (2 * (i // 2)) / d_model)
    
    # sin/cosを適用
    angles[:, 0::2] = np.sin(angles[:, 0::2])  # 偶数インデックス
    angles[:, 1::2] = np.cos(angles[:, 1::2])  # 奇数インデックス
    
    pos_encoding = torch.tensor(angles, dtype=torch.float32)
    return pos_encoding

# 使用例
pe = positional_encoding(seq_len=100, d_model=512)
print(f"Position Encoding shape: {pe.shape}")  # [100, 512]

# 1. 必要なライブラリのインストール
pip install transformers torch datasets

# 2. 基本的な使用例
from transformers import AutoTokenizer, AutoModel
import torch

# モデルとトークナイザーの読み込み
model_name = "bert-base-japanese"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# テキストの処理
text = "今日は良い天気です。"
inputs = tokenizer(text, return_tensors="pt")

# 推論実行
with torch.no_grad():
    outputs = model(**inputs)
    
# 結果の取得
last_hidden_states = outputs.last_hidden_state
print(f"Output shape: {last_hidden_states.shape}")  # [1, seq_len, hidden_size]

# シンプルなTransformerブロックの実装
import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(TransformerBlock, self).__init__()
        self.attention = nn.MultiheadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        # Self-attention
        attn_output, _ = self.attention(x, x, x)
        x = self.norm1(x + self.dropout(attn_output))
        
        # Feed forward
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        
        return x

# 使用例
d_model = 512
num_heads = 8
d_ff = 2048
seq_len = 10
batch_size = 2

transformer_block = TransformerBlock(d_model, num_heads, d_ff)
x = torch.randn(seq_len, batch_size, d_model)
output = transformer_block(x)
print(f"Output shape: {output.shape}")  # [10, 2, 512]