# transformer.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    """
    Multi-Head Self-Attention

    Time: O(n² × d) where n=seq_len, d=d_model
    Space: O(n²) for attention matrix
    """

    def __init__(self, d_model=768, num_heads=12, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 64 per head

        # Linear projections
        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)

        self.dropout = nn.Dropout(dropout)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        """
        Scaled Dot-Product Attention

        Args:
            Q, K, V: [batch, heads, seq_len, d_k]
            mask: [batch, 1, 1, seq_len] for padding

        Returns:
            output: [batch, heads, seq_len, d_k]
            attention_weights: [batch, heads, seq_len, seq_len]
        """
        # scores: [batch, heads, seq_len, seq_len]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        # Apply mask (padding = -inf)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Attention weights
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Apply to values
        output = torch.matmul(attn_weights, V)
        return output, attn_weights

    def split_heads(self, x):
        """[batch, seq, d_model] → [batch, heads, seq, d_k]"""
        batch_size, seq_len, d_model = x.size()
        return x.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

    def combine_heads(self, x):
        """[batch, heads, seq, d_k] → [batch, seq, d_model]"""
        batch_size, _, seq_len, d_k = x.size()
        return x.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

    def forward(self, query, key, value, mask=None):
        # Linear projections and split heads
        Q = self.split_heads(self.W_q(query))
        K = self.split_heads(self.W_k(key))
        V = self.split_heads(self.W_v(value))

        # Attention
        attn_output, attn_weights = self.scaled_dot_product_attention(Q, K, V, mask)

        # Combine heads and final projection
        output = self.combine_heads(attn_output)
        output = self.W_o(output)

        return output, attn_weights

class PositionWiseFeedForward(nn.Module):
    """
    FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

    Applied independently to each position
    """

    def __init__(self, d_model=768, d_ff=3072, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        return self.linear2(self.dropout(F.gelu(self.linear1(x))))

class TransformerEncoderLayer(nn.Module):
    """Single encoder layer with self-attention + FFN"""

    def __init__(self, d_model=768, num_heads=12, d_ff=3072, dropout=0.1):
        super().__init__()

        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = PositionWiseFeedForward(d_model, d_ff, dropout)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Self-attention + residual + norm
        attn_output, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))

        # FFN + residual + norm
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout2(ffn_output))

        return x

class PositionalEncoding(nn.Module):
    """
    Sinusoidal positional encoding

    PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    """

    def __init__(self, d_model=768, max_len=512):
        super().__init__()

        # Create PE matrix [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()

        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)

        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len, :]

class TransformerEncoder(nn.Module):
    """Complete Transformer Encoder (BERT-style)"""

    def __init__(
        self,
        vocab_size=30522,
        d_model=768,
        num_layers=12,
        num_heads=12,
        d_ff=3072,
        max_len=512,
        dropout=0.1
    ):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_len)

        self.layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])

        self.dropout = nn.Dropout(dropout)
        self.d_model = d_model

    def forward(self, input_ids, attention_mask=None):
        """
        Args:
            input_ids: [batch, seq_len]
            attention_mask: [batch, seq_len] (1=real, 0=padding)

        Returns:
            [batch, seq_len, d_model]
        """
        # Token embeddings + scaling
        x = self.embedding(input_ids) * math.sqrt(self.d_model)

        # Add positional encoding
        x = self.pos_encoding(x)
        x = self.dropout(x)

        # Prepare attention mask [batch, 1, 1, seq_len]
        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)

        # Pass through layers
        for layer in self.layers:
            x = layer(x, attention_mask)

        return x

# Example
if __name__ == "__main__":
    model = TransformerEncoder(
        vocab_size=30522,  # BERT vocab
        d_model=768,
        num_layers=12,
        num_heads=12
    )

    input_ids = torch.randint(0, 30522, (2, 10))  # batch=2, seq=10
    mask = torch.ones(2, 10)

    output = model(input_ids, mask)
    print(f"Output shape: {output.shape}")  # [2, 10, 768]

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total_params:,}")  # ~110M

# bert_implementation.py
import torch
import torch.nn as nn
from transformers import BertTokenizer, BertModel, BertForSequenceClassification
from torch.utils.data import Dataset, DataLoader
import numpy as np

class BERTClassifier(nn.Module):
    """
    BERT for sequence classification

    Architecture:
    Input → BERT → [CLS] embedding → Dropout → Linear → Softmax
    """

    def __init__(self, num_classes=2, dropout=0.1):
        super().__init__()

        # Load pre-trained BERT
        self.bert = BertModel.from_pretrained('bert-base-uncased')

        # Classification head
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(768, num_classes)

    def forward(self, input_ids, attention_mask):
        """
        Args:
            input_ids: [batch, seq_len]
            attention_mask: [batch, seq_len]

        Returns:
            logits: [batch, num_classes]
        """
        # Get BERT outputs
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

        # Use [CLS] token representation (first token)
        pooled_output = outputs.pooler_output  # [batch, 768]

        # Classification
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        return logits

class TextDataset(Dataset):
    """Dataset for BERT fine-tuning"""

    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]

        # Tokenize
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,  # Add [CLS] and [SEP]
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(label, dtype=torch.long)
        }

def fine_tune_bert(train_texts, train_labels, val_texts, val_labels):
    """
    Fine-tune BERT for classification

    Training strategy:
    1. Freeze BERT layers initially (optional)
    2. Use smaller learning rate for BERT (2e-5)
    3. Gradient accumulation for larger effective batch
    4. Warmup + linear decay scheduler
    """

    # Initialize
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BERTClassifier(num_classes=2)

    # Datasets
    train_dataset = TextDataset(train_texts, train_labels, tokenizer)
    val_dataset = TextDataset(val_texts, val_labels, tokenizer)

    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=32)

    # Optimizer with differential learning rates
    optimizer = torch.optim.AdamW([
        {'params': model.bert.parameters(), 'lr': 2e-5},  # Lower LR for BERT
        {'params': model.classifier.parameters(), 'lr': 1e-3}  # Higher for head
    ], weight_decay=0.01)

    # Loss
    criterion = nn.CrossEntropyLoss()

    # Training loop
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    num_epochs = 3  # BERT typically needs 2-4 epochs

    for epoch in range(num_epochs):
        model.train()
        total_loss = 0

        for batch in train_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['label'].to(device)

            # Forward pass
            logits = model(input_ids, attention_mask)
            loss = criterion(logits, labels)

            # Backward pass
            optimizer.zero_grad()
            loss.backward()

            # Gradient clipping (important for BERT)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()

            total_loss += loss.item()

        # Validation
        model.eval()
        val_correct = 0
        val_total = 0

        with torch.no_grad():
            for batch in val_loader:
                input_ids = batch['input_ids'].to(device)
                attention_mask = batch['attention_mask'].to(device)
                labels = batch['label'].to(device)

                logits = model(input_ids, attention_mask)
                predictions = torch.argmax(logits, dim=1)

                val_correct += (predictions == labels).sum().item()
                val_total += labels.size(0)

        val_accuracy = val_correct / val_total
        avg_loss = total_loss / len(train_loader)

        print(f"Epoch {epoch+1}/{num_epochs}")
        print(f"  Train Loss: {avg_loss:.4f}")
        print(f"  Val Accuracy: {val_accuracy:.4f}")

    return model

# Example: Inference
def predict_sentiment(model, tokenizer, text, device='cpu'):
    """Predict sentiment for single text"""
    model.eval()

    encoding = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=128,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt'
    )

    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        logits = model(input_ids, attention_mask)
        probabilities = torch.softmax(logits, dim=1)
        prediction = torch.argmax(probabilities, dim=1)

    return prediction.item(), probabilities[0].cpu().numpy()

# Example usage
if __name__ == "__main__":
    # Sample data
    train_texts = ["This movie is great!", "Terrible film, waste of time."]
    train_labels = [1, 0]  # 1=positive, 0=negative

    val_texts = ["Loved it!", "Boring."]
    val_labels = [1, 0]

    # Fine-tune
    model = fine_tune_bert(train_texts, train_labels, val_texts, val_labels)

    # Predict
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    sentiment, probs = predict_sentiment(
        model, tokenizer, "This is amazing!"
    )
    print(f"Sentiment: {'Positive' if sentiment == 1 else 'Negative'}")
    print(f"Probabilities: {probs}")

from gensim.models import Word2Vec, KeyedVectors

# Train Word2Vec
model = Word2Vec(sentences, vector_size=100, window=5)

# Load pre-trained GloVe
glove = KeyedVectors.load_word2vec_format('glove.txt')

# Analogies: king - man + woman ≈ queen
model.wv.most_similar(positive=['king', 'woman'], negative=['man'])

from sklearn.feature_extraction.text import TfidfVectorizer

vectorizer = TfidfVectorizer(max_features=5000, ngram_range=(1, 2))
X = vectorizer.fit_transform(documents)

# attention_mechanisms.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaledDotProductAttention(nn.Module):
    """
    Basic attention mechanism

    Time: O(n·m·d) where n=query_len, m=key_len
    Space: O(n·m) for attention matrix
    """

    def __init__(self, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

    def forward(self, Q, K, V, mask=None):
        """
        Args:
            Q: [batch, n, d_k] queries
            K: [batch, m, d_k] keys
            V: [batch, m, d_v] values
            mask: [batch, n, m] or broadcastable

        Returns:
            output: [batch, n, d_v]
            attention_weights: [batch, n, m]
        """
        d_k = Q.size(-1)

        # Compute attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

        # Apply mask (set masked positions to -inf)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Attention weights (softmax over keys dimension)
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # Apply attention to values
        output = torch.matmul(attention_weights, V)

        return output, attention_weights

class AdditiveAttention(nn.Module):
    """
    Bahdanau Attention (additive)

    score(q, k) = v^T tanh(W_q q + W_k k)

    Older mechanism, used in early seq2seq models
    """

    def __init__(self, hidden_dim):
        super().__init__()
        self.W_q = nn.Linear(hidden_dim, hidden_dim)
        self.W_k = nn.Linear(hidden_dim, hidden_dim)
        self.v = nn.Linear(hidden_dim, 1)

    def forward(self, query, keys, values, mask=None):
        """
        Args:
            query: [batch, d_q]
            keys: [batch, seq_len, d_k]
            values: [batch, seq_len, d_v]

        Returns:
            context: [batch, d_v]
            attention_weights: [batch, seq_len]
        """
        # Expand query for broadcasting
        query = query.unsqueeze(1)  # [batch, 1, d_q]

        # Compute scores
        scores = self.v(torch.tanh(
            self.W_q(query) + self.W_k(keys)
        )).squeeze(-1)  # [batch, seq_len]

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Attention weights
        attention_weights = F.softmax(scores, dim=-1)

        # Context vector (weighted sum of values)
        context = torch.matmul(
            attention_weights.unsqueeze(1), values
        ).squeeze(1)

        return context, attention_weights

class MultiHeadAttention(nn.Module):
    """Multi-Head Attention (used in Transformers)"""

    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // 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)

        self.attention = ScaledDotProductAttention(dropout)

    def split_heads(self, x, batch_size):
        """Split into multiple heads"""
        x = x.view(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(1, 2)  # [batch, heads, seq, d_k]

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # Linear projections and split heads
        Q = self.split_heads(self.W_q(query), batch_size)
        K = self.split_heads(self.W_k(key), batch_size)
        V = self.split_heads(self.W_v(value), batch_size)

        # Apply attention
        attn_output, attn_weights = self.attention(Q, K, V, mask)

        # Concatenate heads
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, -1, self.d_model)

        # Final linear
        output = self.W_o(attn_output)

        return output, attn_weights

class CrossAttention(nn.Module):
    """
    Cross-Attention for encoder-decoder

    Query from decoder, Keys/Values from encoder
    Used in: Translation, image captioning, text-to-image
    """

    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, num_heads)

    def forward(self, decoder_hidden, encoder_outputs, mask=None):
        """
        Args:
            decoder_hidden: [batch, dec_len, d_model] (queries)
            encoder_outputs: [batch, enc_len, d_model] (keys & values)
            mask: Optional padding mask for encoder

        Returns:
            output: [batch, dec_len, d_model]
        """
        # Q from decoder, K/V from encoder
        output, attn_weights = self.attention(
            query=decoder_hidden,
            key=encoder_outputs,
            value=encoder_outputs,
            mask=mask
        )

        return output, attn_weights

# Example: Visualize Attention
def visualize_attention_example():
    """Example showing how attention focuses on relevant words"""
    # Simple example: translating "I love cats" to French

    # Encoder outputs (simplified)
    encoder_out = torch.randn(1, 3, 512)  # 3 words: I, love, cats

    # Decoder at step 1 (generating "J'")
    decoder_hidden = torch.randn(1, 1, 512)

    cross_attn = CrossAttention(d_model=512)
    output, attn_weights = cross_attn(decoder_hidden, encoder_out)

    # Attention weights shape: [1, num_heads, 1, 3]
    # Shows how much decoder attends to each encoder word

    print("Attention weights (decoder step 1):")
    print(attn_weights[0, 0, 0, :])
    # Example output: [0.85, 0.10, 0.05]
    # → Decoder focuses on "I" when generating "J'"

if __name__ == "__main__":
    # Test attention mechanisms
    batch_size = 2
    seq_len = 10
    d_model = 512

    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)

    # Scaled dot-product
    attn = ScaledDotProductAttention()
    output, weights = attn(Q, K, V)
    print(f"Output shape: {output.shape}")  # [2, 10, 512]
    print(f"Attention weights shape: {weights.shape}")  # [2, 10, 10]

    # Multi-head attention
    mha = MultiHeadAttention(d_model=512, num_heads=8)
    output, weights = mha(Q, K, V)
    print(f"Multi-head output: {output.shape}")  # [2, 10, 512]

English: "The cat sat on the mat"
French:  "Le chat s'est assis sur le tapis"

When generating "chat":
- High attention on "cat" (0.9)
- Low attention on other words (0.1 distributed)

from transformers import pipeline

ner = pipeline("ner", model="dslim/bert-base-NER")
results = ner("Apple was founded by Steve Jobs in California")
# [{'entity': 'B-ORG', 'word': 'Apple'}, ...]

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
tokens = tokenizer.tokenize("unbelievable")  # ['un', '##bel', '##iev', '##able']

from transformers import pipeline

classifier = pipeline("sentiment-analysis")
result = classifier("I love this product!")
# [{'label': 'POSITIVE', 'score': 0.999}]

GPT Attention Pattern (Causal):
Token 1: Can see [1]
Token 2: Can see [1, 2]
Token 3: Can see [1, 2, 3]
→ Upper triangular mask

BERT Attention Pattern (Bidirectional):
Token 1: Can see [1, 2, 3]
Token 2: Can see [1, 2, 3]
Token 3: Can see [1, 2, 3]
→ No mask (full context)

# gpt_architecture.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class CausalSelfAttention(nn.Module):
    """
    Causal self-attention for GPT

    Key difference from BERT: Upper triangular mask prevents
    attending to future tokens
    """

    def __init__(self, d_model=768, num_heads=12, max_len=1024):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_qkv = nn.Linear(d_model, 3 * d_model)
        self.W_o = nn.Linear(d_model, d_model)

        # Causal mask: upper triangular (prevent attending to future)
        self.register_buffer(
            'causal_mask',
            torch.tril(torch.ones(max_len, max_len)).view(1, 1, max_len, max_len)
        )

    def forward(self, x):
        batch_size, seq_len, d_model = x.size()

        # Compute Q, K, V
        qkv = self.W_qkv(x)  # [batch, seq, 3*d_model]
        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.d_k)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, batch, heads, seq, d_k]

        Q, K, V = qkv[0], qkv[1], qkv[2]

        # Attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        # Apply causal mask
        scores = scores.masked_fill(
            self.causal_mask[:, :, :seq_len, :seq_len] == 0,
            float('-inf')
        )

        # Attention weights
        attn_weights = F.softmax(scores, dim=-1)

        # Apply to values
        out = torch.matmul(attn_weights, V)

        # Concatenate heads
        out = out.transpose(1, 2).contiguous()
        out = out.view(batch_size, seq_len, d_model)

        return self.W_o(out)

class GPTBlock(nn.Module):
    """Single GPT transformer block"""

    def __init__(self, d_model=768, num_heads=12):
        super().__init__()

        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, num_heads)

        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model)
        )

    def forward(self, x):
        # Pre-LN (GPT-2 style)
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

class GPT(nn.Module):
    """GPT model for text generation"""

    def __init__(
        self,
        vocab_size=50257,  # GPT-2 vocab
        d_model=768,
        num_layers=12,
        num_heads=12,
        max_len=1024
    ):
        super().__init__()

        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)

        self.blocks = nn.ModuleList([
            GPTBlock(d_model, num_heads)
            for _ in range(num_layers)
        ])

        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        # Tie weights (embedding = lm_head)
        self.lm_head.weight = self.token_embedding.weight

    def forward(self, input_ids):
        """
        Args:
            input_ids: [batch, seq_len]

        Returns:
            logits: [batch, seq_len, vocab_size]
        """
        batch_size, seq_len = input_ids.size()

        # Token + position embeddings
        positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0)

        x = self.token_embedding(input_ids) + self.position_embedding(positions)

        # Transformer blocks
        for block in self.blocks:
            x = block(x)

        # Final layer norm
        x = self.ln_f(x)

        # Project to vocabulary
        logits = self.lm_head(x)

        return logits

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=50, temperature=1.0, top_k=None):
        """
        Autoregressive generation

        Args:
            input_ids: [batch, seq_len] prompt tokens
            max_new_tokens: Number of tokens to generate
            temperature: Sampling temperature (higher = more random)
            top_k: If set, only sample from top-k tokens

        Returns:
            generated: [batch, seq_len + max_new_tokens]
        """
        for _ in range(max_new_tokens):
            # Get logits for last position
            logits = self(input_ids)  # [batch, seq, vocab]
            logits = logits[:, -1, :] / temperature  # [batch, vocab]

            # Top-k filtering
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')

            # Sample next token
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            # Append to sequence
            input_ids = torch.cat([input_ids, next_token], dim=1)

        return input_ids

# Example usage
if __name__ == "__main__":
    model = GPT(
        vocab_size=50257,
        d_model=768,
        num_layers=12,
        num_heads=12
    )

    # Generate text
    prompt = torch.randint(0, 50257, (1, 10))  # Random prompt
    generated = model.generate(prompt, max_new_tokens=20, temperature=0.8)

    print(f"Prompt shape: {prompt.shape}")
    print(f"Generated shape: {generated.shape}")

    total_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {total_params:,}")  # ~117M for GPT-2 Small

Translate to French: "I love cats"
→ "J'aime les chats"

Translate to French:
English: "Hello" → French: "Bonjour"
English: "I love cats" →
→ French: "J'aime les chats"

English: "Hello" → French: "Bonjour"
English: "Goodbye" → French: "Au revoir"
English: "Thank you" → French: "Merci"
English: "I love cats" →
→ French: "J'aime les chats"

from sumy.summarizers.lex_rank import LexRankSummarizer

from transformers import pipeline

summarizer = pipeline("summarization", model="t5-base")
summary = summarizer(long_text, max_length=100)

Input → [Encoder] → Context → [Decoder] → Output

User Query: "What is the capital of France?"
       ↓
┌──────────────────┐
│  1. Query        │ → Embed query: [768-dim vector]
│     Embedding    │
└────────┬─────────┘
         ↓
┌──────────────────┐
│  2. Vector       │ → Search similar docs (cosine sim)
│     Search       │    Top-k retrieval (k=3-10)
└────────┬─────────┘
         ↓
┌──────────────────┐
│  3. Re-ranking   │ → Optional: Cross-encoder rerank
│     (Optional)   │    Improve relevance
└────────┬─────────┘
         ↓
┌──────────────────┐
│  4. Context      │ → "Paris is the capital..."
│     Retrieved    │    "France's capital city..."
└────────┬─────────┘
         ↓
┌──────────────────┐
│  5. Prompt       │ → Context: {retrieved_docs}
│     Construction │    Query: {user_query}
└────────┬─────────┘
         ↓
┌──────────────────┐
│  6. LLM          │ → GPT-4, Claude, etc.
│     Generation   │
└────────┬─────────┘
         ↓
Response: "The capital of France is Paris."

# positional_encoding.py
import torch
import torch.nn as nn
import math
import matplotlib.pyplot as plt

class SinusoidalPositionalEncoding(nn.Module):
    """
    Original Transformer positional encoding (Vaswani et al., 2017)

    Used in: BERT, GPT-2, T5
    Time: O(n × d) to compute
    Space: O(n × d) cached
    """

    def __init__(self, d_model=768, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # Create PE matrix [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        # Compute div_term: 10000^(2i/d_model) for i in [0, d_model/2)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float) *
            (-math.log(10000.0) / d_model)
        )

        # Apply sin to even indices, cos to odd indices
        pe[:, 0::2] = torch.sin(position * div_term)  # Even: 0, 2, 4, ...
        pe[:, 1::2] = torch.cos(position * div_term)  # Odd: 1, 3, 5, ...

        pe = pe.unsqueeze(0)  # [1, max_len, d_model]

        # Register as buffer (not a parameter, but saved in state_dict)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Args:
            x: [batch, seq_len, d_model] - token embeddings
        Returns:
            [batch, seq_len, d_model] - embeddings + PE
        """
        seq_len = x.size(1)

        # Add positional encoding (broadcasting over batch)
        x = x + self.pe[:, :seq_len, :]
        return self.dropout(x)

class LearnedPositionalEmbedding(nn.Module):
    """
    Learned positional embeddings (used in BERT, GPT-2)

    Pros: Can learn task-specific position patterns
    Cons: Cannot extrapolate to longer sequences

    Used in: BERT, GPT-2, GPT-3
    """

    def __init__(self, max_len=512, d_model=768, dropout=0.1):
        super().__init__()
        self.position_embeddings = nn.Embedding(max_len, d_model)
        self.dropout = nn.Dropout(dropout)

        # Register position IDs (0, 1, 2, ..., max_len-1)
        self.register_buffer(
            'position_ids',
            torch.arange(max_len).expand((1, -1))
        )

    def forward(self, x):
        """
        Args:
            x: [batch, seq_len, d_model]
        Returns:
            [batch, seq_len, d_model]
        """
        seq_len = x.size(1)

        # Get position embeddings for current sequence
        position_ids = self.position_ids[:, :seq_len]
        position_embeds = self.position_embeddings(position_ids)

        return self.dropout(x + position_embeds)

class RoPE(nn.Module):
    """
    Rotary Positional Embedding (RoPE) - Su et al., 2021

    Key idea: Rotate query/key vectors by angle proportional to position

    Advantages:
    - Naturally encodes relative positions
    - Extrapolates to longer sequences
    - Better performance on long sequences

    Used in: LLaMA, GPT-NeoX, PaLM, GPT-J
    """

    def __init__(self, d_model=768, max_len=2048, base=10000):
        super().__init__()
        self.d_model = d_model
        self.max_len = max_len
        self.base = base

        # Precompute rotation frequencies
        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)

        # Precompute cos and sin for max_len positions
        t = torch.arange(max_len, dtype=torch.float)
        freqs = torch.outer(t, inv_freq)  # [max_len, d_model/2]
        emb = torch.cat([freqs, freqs], dim=-1)  # [max_len, d_model]

        self.register_buffer('cos_cached', emb.cos()[None, :, None, :])
        self.register_buffer('sin_cached', emb.sin()[None, :, None, :])

    def rotate_half(self, x):
        """Rotates half the hidden dims of the input"""
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return torch.cat([-x2, x1], dim=-1)

    def forward(self, q, k):
        """
        Apply rotary embeddings to query and key

        Args:
            q, k: [batch, num_heads, seq_len, head_dim]
        Returns:
            q_rot, k_rot: [batch, num_heads, seq_len, head_dim]
        """
        seq_len = q.size(2)

        # Get cached cos and sin for current sequence
        cos = self.cos_cached[:, :seq_len, :, :]
        sin = self.sin_cached[:, :seq_len, :, :]

        # Apply rotation
        q_rot = (q * cos) + (self.rotate_half(q) * sin)
        k_rot = (k * cos) + (self.rotate_half(k) * sin)

        return q_rot, k_rot

class ALiBi(nn.Module):
    """
    Attention with Linear Biases (ALiBi) - Press et al., 2021

    Key idea: Add linear bias to attention scores based on distance

    Formula: attention_score += -m × |i - j|
    where m is head-specific slope

    Advantages:
    - No position embeddings needed
    - Excellent extrapolation (trained on 1K, works on 10K+)
    - Memory efficient

    Used in: BLOOM, MPT, StableLM
    """

    def __init__(self, num_heads=12, max_len=2048):
        super().__init__()
        self.num_heads = num_heads

        # Compute slopes for each head: 2^(-8/n), 2^(-16/n), ...
        slopes = torch.Tensor(self._get_slopes(num_heads))
        self.register_buffer('slopes', slopes)

        # Precompute distance matrix
        positions = torch.arange(max_len)
        distance = positions[None, :] - positions[:, None]  # [max_len, max_len]
        distance = distance.abs().unsqueeze(0).unsqueeze(0)  # [1, 1, max_len, max_len]
        self.register_buffer('distance', distance)

    def _get_slopes(self, num_heads):
        """Compute head-specific slopes"""
        # Geometric sequence: 2^(-8/n) to 2^(-8)
        def get_slopes_power_of_2(n):
            start = 2 ** (-2 ** -(math.log2(n) - 3))
            ratio = start
            return [start * (ratio ** i) for i in range(n)]

        if math.log2(num_heads).is_integer():
            return get_slopes_power_of_2(num_heads)
        else:
            # Closest power of 2
            closest_power = 2 ** math.floor(math.log2(num_heads))
            return (
                get_slopes_power_of_2(closest_power) +
                self._get_slopes(2 * closest_power)[0::2][:num_heads - closest_power]
            )

    def forward(self, attention_scores):
        """
        Add linear position bias to attention scores

        Args:
            attention_scores: [batch, num_heads, seq_len, seq_len]
        Returns:
            biased_scores: [batch, num_heads, seq_len, seq_len]
        """
        seq_len = attention_scores.size(-1)

        # Get distance matrix for current sequence
        distance = self.distance[:, :, :seq_len, :seq_len]

        # Apply head-specific slopes: -m × distance
        bias = -self.slopes.view(1, -1, 1, 1) * distance

        return attention_scores + bias

# Example: Using different positional encodings
def compare_positional_encodings():
    """Demonstrate different positional encoding methods"""
    batch_size, seq_len, d_model = 2, 128, 768
    num_heads = 12

    # Input token embeddings
    token_embeddings = torch.randn(batch_size, seq_len, d_model)

    print("=" * 60)
    print("Positional Encoding Comparison")
    print("=" * 60)

    # 1. Sinusoidal
    sinusoidal_pe = SinusoidalPositionalEncoding(d_model, max_len=512)
    output_sin = sinusoidal_pe(token_embeddings)
    print(f"\n1. Sinusoidal PE:")
    print(f"   Parameters: 0 (deterministic)")
    print(f"   Output shape: {output_sin.shape}")
    print(f"   Can extrapolate: Yes")

    # 2. Learned
    learned_pe = LearnedPositionalEmbedding(max_len=512, d_model=d_model)
    output_learned = learned_pe(token_embeddings)
    params_learned = sum(p.numel() for p in learned_pe.parameters())
    print(f"\n2. Learned PE:")
    print(f"   Parameters: {params_learned:,} (512 × 768)")
    print(f"   Output shape: {output_learned.shape}")
    print(f"   Can extrapolate: No (max_len=512)")

    # 3. RoPE (applied to Q, K in attention)
    rope = RoPE(d_model, max_len=2048)
    head_dim = d_model // num_heads
    q = token_embeddings.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
    k = token_embeddings.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
    q_rot, k_rot = rope(q, k)
    print(f"\n3. RoPE:")
    print(f"   Parameters: 0 (rotation angles)")
    print(f"   Output shape: {q_rot.shape}")
    print(f"   Can extrapolate: Yes (excellent)")

    # 4. ALiBi (applied to attention scores)
    alibi = ALiBi(num_heads, max_len=2048)
    attention_scores = torch.randn(batch_size, num_heads, seq_len, seq_len)
    biased_scores = alibi(attention_scores)
    print(f"\n4. ALiBi:")
    print(f"   Parameters: 0 (linear biases)")
    print(f"   Biased scores shape: {biased_scores.shape}")
    print(f"   Can extrapolate: Yes (best extrapolation)")
    print("=" * 60)

if __name__ == "__main__":
    compare_positional_encodings()