Transformer & Attention Mechanism

import numpy as np
import matplotlib.pyplot as plt

def positional_encoding(max_len, d_model):
    """Sinusoidal Positional Encoding."""
    PE = np.zeros((max_len, d_model))
    positions = np.arange(max_len).reshape(-1, 1)
    
    # 10000^(2i/d_model) = exp(2i * log(10000) / d_model)
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    
    PE[:, 0::2] = np.sin(positions * div_term)  # Even dimensions
    PE[:, 1::2] = np.cos(positions * div_term)  # Odd dimensions
    
    return PE

# Generate PE
max_len = 100
d_model = 64
pe = positional_encoding(max_len, d_model)

# Visualisasi
fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# Plot 1: Full PE heatmap
axes[0, 0].imshow(pe.T, aspect='auto', cmap='RdBu')
axes[0, 0].set_xlabel('Position')
axes[0, 0].set_ylabel('Dimension')
axes[0, 0].set_title('Positional Encoding Heatmap')

# Plot 2: Individual dimensions
for dim in [0, 1, 2, 3, 10, 20]:
    axes[0, 1].plot(pe[:50, dim], label=f'dim {dim}')
axes[0, 1].set_xlabel('Position')
axes[0, 1].set_ylabel('PE Value')
axes[0, 1].set_title('PE Values across Positions')
axes[0, 1].legend(fontsize=8)
axes[0, 1].grid(True, alpha=0.3)

# Plot 3: Similarity matrix (cosine sim antar posisi)
from numpy.linalg import norm
def cosine_sim(a, b):
    return np.dot(a, b) / (norm(a) * norm(b) + 1e-8)

n_pos = 50
sim_matrix = np.zeros((n_pos, n_pos))
for i in range(n_pos):
    for j in range(n_pos):
        sim_matrix[i, j] = cosine_sim(pe[i], pe[j])

axes[1, 0].imshow(sim_matrix, cmap='viridis')
axes[1, 0].set_xlabel('Position j')
axes[1, 0].set_ylabel('Position i')
axes[1, 0].set_title('Cosine Similarity antar Posisi')
plt.colorbar(axes[1, 0].images[0], ax=axes[1, 0])

# Plot 4: Info
axes[1, 1].text(0.5, 0.5, 
    f'Positional Encoding\n\n'
    f'Max Length: {max_len}\n'
    f'Model Dimension: {d_model}\n'
    f'Frequency Range: sin & cos\n'
    f'Method: Sinusoidal\n\n'
    f'Properties:\n'
    f'• PE(pos+k) bisa direpresentasikan\n'
    f'  sebagai fungsi linear dari PE(pos)\n'
    f'• Setiap posisi unik\n'
    f'• Bisa menangani sequence yang\n'
    f'  lebih panjang dari training',
    ha='center', va='center', fontsize=11,
    transform=axes[1, 1].transAxes,
    bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.3))
axes[1, 1].axis('off')

plt.suptitle('Positional Encoding — Transformer', fontsize=14)
plt.tight_layout()
plt.show()

from transformers import BertTokenizer, BertModel, BertForSequenceClassification
import torch

# =============================================
# 1. LOAD PRE-TRAINED BERT
# =============================================
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

print(f"Model: BERT-base-uncased")
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
print(f"Vocabulary size: {tokenizer.vocab_size:,}")
print(f"Max position: {model.config.max_position_embeddings}")

# =============================================
# 2. TOKENIZATION
# =============================================
text = "Transformer is a revolutionary architecture in NLP."
tokens = tokenizer(text, return_tensors='pt', padding=True, truncation=True)

print(f"\nOriginal text: {text}")
print(f"Token IDs: {tokens['input_ids']}")
print(f"Tokens: {tokenizer.convert_ids_to_tokens(tokens['input_ids'][0])}")
print(f"Attention mask: {tokens['attention_mask']}")

# BERT tokenization menggunakan WordPiece:
# "revolutionary" → ["revolution", "##ary"]
# "#" = subword continuation token

# =============================================
# 3. GET EMBEDDINGS
# =============================================
with torch.no_grad():
    outputs = model(**tokens)

# outputs.last_hidden_state: (batch, seq_len, hidden_size=768)
# outputs.pooler_output: (batch, hidden_size=768) — [CLS] token

print(f"\nLast hidden state shape: {outputs.last_hidden_state.shape}")
print(f"Pooler output (CLS) shape: {outputs.pooler_output.shape}")

# =============================================
# 4. FINE-TUNE UNTUK SENTIMENT ANALYSIS
# =============================================
from transformers import BertForSequenceClassification
import torch.nn as nn

# Load BERT untuk klasifikasi (2 labels: positive/negative)
model_clf = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=2
)

# Fine-tuning example
texts = ["This movie is amazing!", "I hated this film."]
labels = torch.tensor([1, 0])  # 1=positive, 0=negative

inputs = tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
outputs = model_clf(**inputs, labels=labels)

print(f"\nFine-tuning Loss: {outputs.loss.item():.4f}")
print(f"Logits shape: {outputs.logits.shape}")

# Prediction
predictions = torch.softmax(outputs.logits, dim=1)
print(f"Predictions: {predictions}")
# [positive_prob, negative_prob] untuk setiap teks

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaledDotProductAttention(nn.Module):
    """Scaled Dot-Product Attention."""
    def __init__(self, d_k):
        super().__init__()
        self.d_k = d_k
    
    def forward(self, Q, K, V, mask=None):
        # Q, K, V shape: (batch, heads, seq_len, d_k)
        
        # Step 1: Hitung attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        # scores shape: (batch, heads, seq_len, seq_len)
        
        # Step 2: Apply mask (untuk decoder - cegah melihat masa depan)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # Step 3: Softmax
        attn_weights = F.softmax(scores, dim=-1)
        
        # Step 4: Weighted sum
        output = torch.matmul(attn_weights, V)
        # output shape: (batch, heads, seq_len, d_k)
        
        return output, attn_weights


class MultiHeadAttention(nn.Module):
    """Multi-Head Attention."""
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # dimensi per head
        
        # Linear projections
        self.W_q = nn.Linear(d_model, d_model)  # Query projection
        self.W_k = nn.Linear(d_model, d_model)  # Key projection
        self.W_v = nn.Linear(d_model, d_model)  # Value projection
        self.W_o = nn.Linear(d_model, d_model)  # Output projection
        
        self.attention = ScaledDotProductAttention(self.d_k)
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 1. Linear projection
        Q = self.W_q(Q)  # (batch, seq_len, d_model)
        K = self.W_k(K)
        V = self.W_v(V)
        
        # 2. Split into multiple heads
        # (batch, seq_len, d_model) → (batch, n_heads, seq_len, d_k)
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        
        # 3. Scaled dot-product attention
        output, attn_weights = self.attention(Q, K, V, mask)
        
        # 4. Concatenate heads
        # (batch, n_heads, seq_len, d_k) → (batch, seq_len, d_model)
        output = output.transpose(1, 2).contiguous().view(
            batch_size, -1, self.d_model
        )
        
        # 5. Final linear projection
        output = self.W_o(output)
        
        return output, attn_weights


class TransformerBlock(nn.Module):
    """Single Transformer Encoder Block."""
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        
        # Multi-Head Self-Attention
        self.attention = MultiHeadAttention(d_model, n_heads)
        
        # Feed-Forward Network
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        
        # Layer Normalization
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        
        # Dropout
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # 1. Self-Attention + Residual + LayerNorm
        attn_out, attn_weights = self.attention(x, x, x, mask)
        x = self.ln1(x + self.dropout(attn_out))
        
        # 2. FFN + Residual + LayerNorm
        ffn_out = self.ffn(x)
        x = self.ln2(x + self.dropout(ffn_out))
        
        return x, attn_weights


# =============================================
# TEST: Jalankan Multi-Head Attention
# =============================================
print("=" * 60)
print("TESTING MULTI-HEAD SELF-ATTENTION")
print("=" * 60)

# Hyperparameters
d_model = 512
n_heads = 8
d_ff = 2048
seq_len = 10
batch_size = 2

# Create model
block = TransformerBlock(d_model, n_heads, d_ff)

# Random input (simulating token embeddings)
x = torch.randn(batch_size, seq_len, d_model)

# Forward pass
output, attn_weights = block(x)

print(f"Input shape: {x.shape}")             # (2, 10, 512)
print(f"Output shape: {output.shape}")        # (2, 10, 512)
print(f"Attention weights shape: {attn_weights.shape}")  # (2, 8, 10, 10)

# Attention weight visualization
print(f"\nAttention weights (Head 0, Sample 0):")
print(f"Shape: {attn_weights[0, 0].shape} (seq_len × seq_len)")
print(f"Row sum: {attn_weights[0, 0].sum(dim=-1).tolist()}")  # ≈ 1.0 per row

# Parameter count
params = sum(p.numel() for p in block.parameters())
print(f"\nParameters in Transformer block: {params:,}")

# Multi-head attention shapes summary
print("\n=== Shape Summary ===")
print(f"d_model = {d_model}")
print(f"n_heads = {n_heads}")
print(f"d_k = d_v = {d_model // n_heads}")
print(f"Q, K, V per head: ({seq_len}, {d_model // n_heads})")
print(f"Attention matrix: ({seq_len}, {seq_len})")
print(f"Output per head: ({seq_len}, {d_model // n_heads})")
print(f"Concat output: ({seq_len}, {d_model})")