from transformers import BertTokenizer, BertModel
import torch

# 1. 加載預(yù)訓(xùn)練的 BERT tokenizer 和模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 2. 輸入文本
text = "A Titan RTX has 24GB of VRAM"

# 3. 分詞并映射為 token ID 序列
inputs = tokenizer(text, return_tensors="pt", truncatinotallow=True, padding=True)

# 輸出 token IDs
print("Token IDs:", inputs['input_ids'])

# 4. 傳入模型，獲取輸出
outputs = model(**inputs)

# 5. 獲取最后一層的隱藏狀態(tài)表示
last_hidden_states = outputs.last_hidden_state
print("Last hidden states shape:", last_hidden_states.shape)

[101, 138, 28318, 56898, 12674, 10393, 10233, 32469, 10108, 74727, 36535, 102]

from transformers import BertTokenizer
import torch.nn as nn

## 1, 使用 BERT tokenizer 將批量輸入的字符串文本序列轉(zhuǎn)化為 input_ids
tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased") 
batch_text = ["A Titan RTX has 24GB of VRAM", "I have a dog and cat"]
inputs = tokenizer(batch_text, return_tensors="pt", truncation=True, padding=True)
input_ids = inputs["input_ids"]

# 2. 創(chuàng)建一個(gè) nn.Embedding 層
vocab_size = tokenizer.vocab_size  # 詞表大小取決于你加載的具體 tokenizer 模型
embedding_dim = 512  # 嵌入向量的維度，參考 transformer 論文的大小
embedding_layer = nn.Embedding(vocab_size, embedding_dim)

# 3. 通過 nn.Embedding 層，將輸入的 IDs 映射到嵌入向量
embedded_output = embedding_layer(input_ids)

# 4. 輸出嵌入向量的形狀
print("嵌入向量的形狀：", embedded_output.shape)  # (batch_size, sequence_length, embedding_dim), torch.Size([2, 12, 512])

# 5. 打印嵌入向量
print(embedded_output)

class PositionalEncoding(nn.Module):
    """
    compute sinusoid encoding.
    """
    def __init__(self, d_model, max_len, device):
        """
        constructor of sinusoid encoding class

        :param d_model: dimension of model
        :param max_len: max sequence length
        :param device: hardware device setting
        """
        super(PositionalEncoding, self).__init__()

        # same size with input matrix (for adding with input matrix)
        self.encoding = torch.zeros(max_len, d_model, device=device)
        self.encoding.requires_grad = False  # we don't need to compute gradient

        pos = torch.arange(0, max_len, device=device)
        pos = pos.float().unsqueeze(dim=1)
        # 1D => 2D unsqueeze to represent word's position

        _2i = torch.arange(0, d_model, step=2, device=device).float()
        # 'i' means index of d_model (e.g. embedding size = 50, 'i' = [0,50])
        # "step=2" means 'i' multiplied with two (same with 2 * i)

        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))
        # compute positional encoding to consider positional information of words

    def forward(self, x):
        # self.encoding
        # [max_len = 512, d_model = 512]

        batch_size, seq_len = x.size()
        # [batch_size = 128, seq_len = 30]

        return self.encoding[:seq_len, :]
        # [seq_len = 30, d_model = 512]
        # it will add with tok_emb : [128, 30, 512]         

class TokenEmbedding(nn.Embedding):
    """
    Token Embedding using torch.nn
    they will dense representation of word using weighted matrix
    """

    def __init__(self, vocab_size, d_model):
        """
        class for token embedding that included positional information
        :param vocab_size: size of vocabulary
        :param d_model: dimensions of model
        """
        super(TokenEmbedding, self).__init__(vocab_size, d_model, padding_idx=1)

class TransformerEmbedding(nn.Module):
    """
    token embedding + positional encoding (sinusoid)
    positional encoding can give positional information to network
    """

    def __init__(self, vocab_size, max_len, d_model, drop_prob, device):
        """
        class for word embedding that included positional information
        :param vocab_size: size of vocabulary
        :param d_model: dimensions of model
        """
        super(TransformerEmbedding, self).__init__()
        self.tok_emb = TokenEmbedding(vocab_size, d_model)
        # self.position_embedding = nn.Embedding(max_len, embed_size)
        self.pos_emb = PositionalEncoding(d_model, max_len, device)
        self.drop_out = nn.Dropout(p=drop_prob)

    def forward(self, x):
        tok_emb = self.tok_emb(x)
        pos_emb = self.pos_emb(x)
        return self.drop_out(tok_emb + pos_emb)

import torch
import math
import torch.nn as nn

class ScaleDotProductAttention(nn.Module):
    def __init__(self):
        """
        初始化 Self-Attention 層，僅包含一個(gè) softmax 操作。
        """
        super(ScaleDotProductAttention, self).__init__()
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: torch.Tensor = None):
        """
        Self-Attention 前向傳播函數(shù)

        :param Q: Query 向量，形狀為 [batch_size, seq_len, d_k]
        :param K: Key 向量，形狀為 [batch_size, seq_len, d_k]
        :param V: Value 向量，形狀為 [batch_size, seq_len, d_v]
        :param mask: 掩碼張量，形狀為 [batch_size, seq_len, seq_len]
        :return: 
            output: 加權(quán)后的 Value 向量，形狀為 [batch_size, seq_len, d_v]
            attn_weights: 注意力權(quán)重矩陣，形狀為 [batch_size, seq_len, seq_len]
        """
        # 1. 計(jì)算 QK^T 得到注意力分?jǐn)?shù)
        K_T = K.transpose(-1, -2)  # [batch_size, d_k, seq_len]
        scores = torch.matmul(Q, K_T) / math.sqrt(Q.size(-1))  # [batch_size, seq_len, seq_len]

        # 2. 如果有 mask，應(yīng)用掩碼（例如 Decoder 中防止看到未來詞）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 3. 應(yīng)用 softmax 得到注意力權(quán)重
        attn_weights = self.softmax(scores)  # [batch_size, seq_len, seq_len]

        # 4. 權(quán)重 × Value 得到最終輸出
        output = torch.matmul(attn_weights, V)  # [batch_size, seq_len, d_v]

        return output, attn_weights

# 創(chuàng)建 Q、K、V 張量
Q = torch.randn(5, 10, 64)  # [batch_size=5, seq_len=10, d_k=64]
K = torch.randn(5, 10, 64)
V = torch.randn(5, 10, 64)

# 創(chuàng)建 Self-Attention 層
attention = ScaleDotProductAttention()

# 前向傳播
output, attn_weights = attention(Q, K, V)

# 打印輸出形狀
print(f"ScaleDotProductAttention output shape: {output.shape}")      # [5, 10, 64]
print(f"attn_weights shape: {attn_weights.shape}")                # [5, 10, 10]

import torch
import math
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    """Multi-Head Attention Layer"""
    
    def __init__(self, d_model, n_head):
        """
        Args:
            d_model: 模型嵌入維度（通常為 512 或 768）；
            n_head: 注意力頭的數(shù)量（如 8）；
        """
        super(MultiHeadAttention, self).__init__()
        
        # 初始化參數(shù)
        self.n_head = n_head
        self.attention = ScaleDotProductAttention()  # 使用前面定義的 Self-Attention
        
        # 線性變換層
        self.w_q = nn.Linear(d_model, d_model)       # Query 變換
        self.w_k = nn.Linear(d_model, d_model)       # Key 變換
        self.w_v = nn.Linear(d_model, d_model)       # Value 變換
        self.fc = nn.Linear(d_model, d_model)         # 輸出投影層

    def forward(self, q, k, v, mask=None):
        """
        Args:
            q: Query 張量，[batch_size, seq_len, d_model]
            k: Key 張量，[batch_size, seq_len, d_model]
            v: Value 張量，[batch_size, seq_len, d_model]
            mask: 掩碼張量，[batch_size, seq_len, seq_len]
        """
        # Step 1: 線性變換
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)

        # Step 2: 拆分到多個(gè) head
        q = self.split(q)   # [batch_size, n_head, seq_len, d_tensor]
        k = self.split(k)   # [batch_size, n_head, seq_len, d_tensor]
        v = self.split(v)   # [batch_size, n_head, seq_len, d_tensor]

        # Step 3: 計(jì)算每個(gè) head 的 attention
        sa_output, attn_weights = self.attention(q, k, v, mask)

        # Step 4: 拼接所有 head 的輸出
        mha_output = self.concat(sa_output)  # [batch_size, seq_len, d_model]

        # Step 5: 最終線性變換
        mha_output = self.fc(mha_output)

        return mha_output, attn_weights

    def split(self, tensor):
        """
        拆分輸入張量為多個(gè) head

        Args:
            tensor: [batch_size, seq_len, d_model]
        Returns:
            [batch_size, n_head, seq_len, d_tensor]
        """
        batch_size, seq_len, d_model = tensor.size()
        d_tensor = d_model // self.n_head  # 每個(gè) head 的維度
        
        # reshape + transpose 實(shí)現(xiàn)拆分
        tensor = tensor.view(batch_size, seq_len, self.n_head, d_tensor)
        tensor = tensor.transpose(1, 2)  # [batch_size, n_head, seq_len, d_tensor]
        
        return tensor

    def concat(self, sa_output):
        """
        拼接多個(gè) head 的輸出

        Args:
            sa_output: [batch_size, n_head, seq_len, d_tensor]
        Returns:
            [batch_size, seq_len, d_model]
        """
        batch_size, n_head, seq_len, d_tensor = sa_output.size()
        d_model = n_head * d_tensor
        
        # transpose + reshape 實(shí)現(xiàn)合并
        sa_output = sa_output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        
        return sa_output

# 定義參數(shù)
d_model = 512
n_head = 8
seq_len = 10
batch_size = 32

# 創(chuàng)建 Q、K、V 張量
Q = torch.randn(batch_size, seq_len, d_model)
K = torch.randn(batch_size, seq_len, d_model)
V = torch.randn(batch_size, seq_len, d_model)

# 構(gòu)建 MHA 層
mha_layer = MultiHeadAttention(d_model=d_model, n_head=n_head)

# 前向傳播
output, weights = mha_layer(Q, K, V)

# 打印輸出形狀
print("MHA Output Shape:", output.shape)      # [32, 10, 512]
print("Attn Weights Shape:", weights.shape)   # [32, 8, 10, 10]

sa_output, attn_weights = MultiHeadAttention(q=x, k=x, v=x, mask=src_mask)

x = x + dropout(sa_output)
x = layer_norm(x)

nn.Sequential(
    nn.Linear(d_model, d_ff),
    nn.ReLU(),
    nn.Linear(d_ff, d_model),
    nn.Dropout(drop_prob)
)

x = x + dropout(ffn_output)
x = layer_norm(x)

import torch
import torch.nn as nn

class EncoderBlock(nn.Module):
    def __init__(self, d_model, n_head, d_ff, drop_prob=0.1):
        """
        Args:
            d_model: 嵌入維度（例如 512）
            n_head: 多頭數(shù)量（通常設(shè)為 8）
            d_ff: Feed Forward 網(wǎng)絡(luò)中間維度（通常為 2048）
            drop_prob: Dropout 概率
        """
        super(EncoderBlock, self).__init__()
        self.attention = MultiHeadAttention(d_model, n_head)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(drop_prob)
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(drop_prob)

    def forward(self, x, src_mask=None):
        # Step 1: Multi-Head Self-Attention
        sa_output, _ = self.attention(x, x, x, src_mask)
        x = self.norm1(x + self.dropout(sa_output))

        # Step 2: Position-wise FeedForward
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output))

        return x

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, d_model, n_head, d_ff, drop_prob=0.1):
        """
        Args:
            num_layers: encoder block 堆疊層數(shù)（原論文為 6）
            d_model: 模型維度（如 512）
            n_head: 注意力頭數(shù)（如 8）
            d_ff: FeedForward 網(wǎng)絡(luò)維度（如 2048）
        """
        super(TransformerEncoder, self).__init__()
        self.blocks = nn.ModuleList([
            EncoderBlock(d_model, n_head, d_ff, drop_prob)
            for _ in range(num_layers)
        ])

    def forward(self, x, mask=None):
        for block in self.blocks:
            x = block(x, mask)
        return x

# 創(chuàng)建輸入張量
x = torch.randn(batch_size=32, seq_len=20, d_model=512)  # [32, 20, 512]

# 構(gòu)建 Encoder
encoder = TransformerEncoder(num_layers=6, d_model=512, n_head=8, d_ff=2048)
output = encoder(x)

print("Encoder 輸出形狀:", output.shape)  # [32, 20, 512]

x = self.mha1(q=dec_out, k=dec_out, v=dec_out, mask=trg_mask)

x = self.mha2(q=x, k=enc_out, v=enc_out, mask=src_mask)

x = self.ffn(x)

x = self.ln(x_residual + dropout(sublayer_output))

class DecoderLayer(nn.Module):
    def __init__(self, d_model, ffn_hidden, n_head, drop_prob):
        super(DecoderLayer, self).__init__()
        # 第一個(gè) Multi-Head Attention: Masked Self-Attention
        self.mha1 = MultiHeadAttention(d_model, n_head)
        self.ln1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(p=drop_prob)

        # 第二個(gè) Multi-Head Attention: Encoder-Decoder Attention
        self.mha2 = MultiHeadAttention(d_model, n_head)
        self.ln2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(p=drop_prob)

        # 前饋網(wǎng)絡(luò)
        self.ffn = PositionwiseFeedForward(d_model, ffn_hidden)
        self.ln3 = nn.LayerNorm(d_model)
        self.dropout3 = nn.Dropout(p=drop_prob)

    def forward(self, dec_out, enc_out, trg_mask, src_mask):
        x_residual1 = dec_out
        # Step 1: Masked Self-Attention
        x = self.mha1(q=dec_out, k=dec_out, v=dec_out, mask=trg_mask)
        x = self.ln1(x_residual1 + self.dropout1(x))

        if enc_out is not None:
            # Step 2: Encoder-Decoder Attention
            x_residual2 = x
            x = self.mha2(q=x, k=enc_out, v=enc_out, mask=src_mask)
            x = self.ln2(x_residual2 + self.dropout2(x))

        # Step 3: Position-wise Feed Forward
        x_residual3 = x
        x = self.ffn(x)
        x = self.ln3(x_residual3 + self.dropout3(x))

        return x

class Decoder(nn.Module):
    def __init__(self, dec_voc_size, max_len, d_model, ffn_hidden, n_head, n_layers, drop_prob, device):
        super().__init__()
        # 輸入嵌入 + 位置編碼
        self.emb = TransformerEmbedding(
            d_model=d_model,
            drop_prob=drop_prob,
            max_len=max_len,
            vocab_size=dec_voc_size,
            device=device
        )

        # 堆疊多個(gè) Decoder Layer
        self.layers = nn.ModuleList([
            DecoderLayer(
                d_model=d_model,
                ffn_hidden=ffn_hidden,
                n_head=n_head,
                drop_prob=drop_prob
            ) for _ in range(n_layers)
        ])

        # 最終輸出層：映射到目標(biāo)詞匯表大小
        self.linear = nn.Linear(d_model, dec_voc_size)

    def forward(self, trg, src, trg_mask, src_mask):
        # trg: 目標(biāo)序列 [batch_size, trg_seq_len]
        # src: Encoder 輸出 [batch_size, src_seq_len, d_model]
        trg = self.emb(trg)

        for layer in self.layers:
            trg = layer(trg, src, trg_mask, src_mask)

        # 輸出：[batch_size, trg_seq_len, dec_voc_size]
        output = self.linear(trg)
        return output

import torch
from torch import nn

class Transformer(nn.Module):
    def __init__(self, src_pad_idx, trg_pad_idx, trg_sos_idx, enc_voc_size, dec_voc_size, 
                 d_model, n_head, max_len, ffn_hidden, n_layers, drop_prob, device):
        super().__init__()
        
        # 保存特殊標(biāo)記的索引
        self.src_pad_idx = src_pad_idx  # 源語言填充索引
        self.trg_pad_idx = trg_pad_idx  # 目標(biāo)語言填充索引
        self.trg_sos_idx = trg_sos_idx  # 目標(biāo)語言起始符號(hào)索引
        self.device = device  # 設(shè)備信息
        
        # 構(gòu)建 Encoder
        self.encoder = Encoder(
            d_model=d_model,  # 模型維度
            n_head=n_head,  # 注意力頭數(shù)量
            max_len=max_len,  # 最大序列長度
            ffn_hidden=ffn_hidden,  # 前饋網(wǎng)絡(luò)隱藏層維度
            enc_voc_size=enc_voc_size,  # 源語言詞匯表大小
            drop_prob=drop_prob,  # Dropout 概率
            n_layers=n_layers,  # Encoder 層數(shù)
            device=device  # 設(shè)備信息
        )
        
        # 構(gòu)建 Decoder
        self.decoder = Decoder(
            d_model=d_model,  # 模型維度
            n_head=n_head,  # 注意力頭數(shù)量
            max_len=max_len,  # 最大序列長度
            ffn_hidden=ffn_hidden,  # 前饋網(wǎng)絡(luò)隱藏層維度
            dec_voc_size=dec_voc_size,  # 目標(biāo)語言詞匯表大小
            drop_prob=drop_prob,  # Dropout 概率
            n_layers=n_layers,  # Decoder 層數(shù)
            device=device  # 設(shè)備信息
        )

    def forward(self, src, trg):
        # 創(chuàng)建源序列的 padding mask
        src_mask = self.make_pad_mask(src, src, self.src_pad_idx, self.src_pad_idx)
        
        # 創(chuàng)建目標(biāo)序列到源序列的 mask
        src_trg_mask = self.make_pad_mask(trg, src, self.trg_pad_idx, self.src_pad_idx)
        
        # 創(chuàng)建目標(biāo)序列的 padding mask 和因果mask的組合
        trg_mask = self.make_pad_mask(trg, trg, self.trg_pad_idx, self.trg_pad_idx) * \
                   self.make_no_peak_mask(trg, trg)

        # 編碼器前向傳播
        enc_src = self.encoder(src, src_mask)
        
        # 解碼器前向傳播
        output = self.decoder(trg, enc_src, trg_mask, src_trg_mask)
        
        return output

    def make_pad_mask(self, q, k, q_pad_idx, k_pad_idx):
        # 獲取輸入序列長度
        len_q, len_k = q.size(1), k.size(1)

        # 創(chuàng)建針對(duì) key 的 mask
        # batch_size x 1 x 1 x len_k
        k_mask = k.ne(k_pad_idx).unsqueeze(1).unsqueeze(2)
        # batch_size x 1 x len_q x len_k
        k_mask = k_mask.repeat(1, 1, len_q, 1)

        # 創(chuàng)建針對(duì) query 的 mask
        # batch_size x 1 x len_q x 1
        q_mask = q.ne(q_pad_idx).unsqueeze(1).unsqueeze(3)
        # batch_size x 1 x len_q x len_k
        q_mask = q_mask.repeat(1, 1, 1, len_k)
        
        # 組合兩個(gè) mask
        mask = k_mask & q_mask
        
        return mask

    def make_no_peak_mask(self, q, k):
        # 創(chuàng)建因果mask，防止解碼器看到未來信息
        len_q, len_k = q.size(1), k.size(1)
        
        # 創(chuàng)建下三角矩陣，保證解碼器只能關(guān)注當(dāng)前詞及之前的詞
        mask = torch.tril(torch.ones(len_q, len_k)).type(torch.BoolTensor).to(self.device)
        
        return mask