Transformers & Attention — The AI Revolution

import numpy as np

def softmax(x, axis=-1):
    """Numerically stable softmax."""
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Scaled Dot-Product Attention from 'Attention Is All You Need'.
    
    Args:
        Q: Queries  (batch, n, d_k) or (n, d_k)
        K: Keys     (batch, m, d_k) or (m, d_k)
        V: Values   (batch, m, d_v) or (m, d_v)
        mask: Optional mask (n, m) — 0 for positions to attend, -inf for masked
    
    Returns:
        output: Weighted values (batch, n, d_v) or (n, d_v)
        attention_weights: (batch, n, m) or (n, m)
    """
    d_k = Q.shape[-1]
    
    # Step 1: QK^T dot products
    scores = np.matmul(Q, K.swapaxes(-2, -1))  # (n, m)
    
    # Step 2: Scale by sqrt(d_k) — the variance argument!
    scores = scores / np.sqrt(d_k)
    
    # Step 3: Apply mask (for decoder / padding)
    if mask is not None:
        scores = scores + mask  # mask has -inf for blocked positions
    
    # Step 4: Softmax to get attention weights
    attention_weights = softmax(scores, axis=-1)
    
    # Step 5: Weighted sum of values
    output = np.matmul(attention_weights, V)
    
    return output, attention_weights

# ─── DEMO: Self-Attention on 4 tokens ───
np.random.seed(42)
seq_len, d_k, d_v = 4, 8, 8

# Random input embeddings
X = np.random.randn(seq_len, d_k)

# Learnable projection matrices (random for demo)
W_Q = np.random.randn(d_k, d_k) * 0.1
W_K = np.random.randn(d_k, d_k) * 0.1
W_V = np.random.randn(d_k, d_v) * 0.1

# Project to Q, K, V
Q = X @ W_Q
K = X @ W_K
V = X @ W_V

# Compute attention
output, weights = scaled_dot_product_attention(Q, K, V)

print("Input shape:", X.shape)
print("Output shape:", output.shape)
print("\nAttention weights (each row sums to 1):")
print(np.round(weights, 3))
print("\nRow sums:", np.round(weights.sum(axis=-1), 6))
        

import numpy as np

class MultiHeadAttention:
    """
    Multi-Head Attention from scratch.
    Splits input into h heads, runs attention in parallel, concatenates.
    """
    def __init__(self, d_model, num_heads):
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # Initialize projection matrices (Xavier initialization)
        scale = np.sqrt(2.0 / (d_model + self.d_k))
        self.W_Q = np.random.randn(d_model, d_model) * scale
        self.W_K = np.random.randn(d_model, d_model) * scale
        self.W_V = np.random.randn(d_model, d_model) * scale
        self.W_O = np.random.randn(d_model, d_model) * scale
    
    def split_heads(self, x):
        """Reshape (batch, seq, d_model) → (batch, heads, seq, d_k)"""
        batch_size = x.shape[0]
        seq_len = x.shape[1]
        # Reshape: (batch, seq, d_model) → (batch, seq, heads, d_k)
        x = x.reshape(batch_size, seq_len, self.num_heads, self.d_k)
        # Transpose: → (batch, heads, seq, d_k)
        return x.transpose(0, 2, 1, 3)
    
    def forward(self, Q, K, V, mask=None):
        """
        Args:
            Q, K, V: (batch, seq_len, d_model)
            mask: optional (seq_len, seq_len)
        Returns:
            output: (batch, seq_len, d_model)
        """
        batch_size = Q.shape[0]
        
        # Step 1: Linear projections
        Q_proj = Q @ self.W_Q  # (batch, n, d_model)
        K_proj = K @ self.W_K
        V_proj = V @ self.W_V
        
        # Step 2: Split into heads
        Q_heads = self.split_heads(Q_proj)  # (batch, h, n, d_k)
        K_heads = self.split_heads(K_proj)
        V_heads = self.split_heads(V_proj)
        
        # Step 3: Scaled dot-product attention per head
        d_k = self.d_k
        scores = np.matmul(Q_heads, K_heads.transpose(0, 1, 3, 2))
        scores = scores / np.sqrt(d_k)
        
        if mask is not None:
            scores = scores + mask
        
        attn_weights = self._softmax(scores)
        head_outputs = np.matmul(attn_weights, V_heads)  # (batch, h, n, d_k)
        
        # Step 4: Concatenate heads
        # (batch, h, n, d_k) → (batch, n, h, d_k) → (batch, n, d_model)
        concat = head_outputs.transpose(0, 2, 1, 3)
        concat = concat.reshape(batch_size, -1, self.d_model)
        
        # Step 5: Final linear projection
        output = concat @ self.W_O
        
        return output, attn_weights
    
    def _softmax(self, x, axis=-1):
        e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
        return e_x / np.sum(e_x, axis=axis, keepdims=True)

# ─── DEMO ───
np.random.seed(42)
batch_size, seq_len, d_model, num_heads = 2, 6, 64, 8

mha = MultiHeadAttention(d_model, num_heads)
X = np.random.randn(batch_size, seq_len, d_model)

output, weights = mha.forward(X, X, X)  # Self-attention

print(f"Input shape:  {X.shape}")        # (2, 6, 64)
print(f"Output shape: {output.shape}")    # (2, 6, 64)
print(f"Weight shape: {weights.shape}")   # (2, 8, 6, 6)
print(f"\nHead 0, Batch 0 attention (6×6):")
print(np.round(weights[0, 0], 3))
        

import numpy as np

def sinusoidal_positional_encoding(max_len, d_model):
    """
    Compute sinusoidal positional encoding as in 'Attention Is All You Need'.
    
    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    """
    PE = np.zeros((max_len, d_model))
    
    positions = np.arange(max_len)[:, np.newaxis]      # (max_len, 1)
    dim_indices = np.arange(0, d_model, 2)              # (d_model/2,)
    
    # Compute the denominator: 10000^(2i/d_model)
    # = exp(2i * ln(10000) / d_model)
    div_term = np.exp(dim_indices * (-np.log(10000.0) / d_model))
    
    # Apply sin to even indices, cos to odd indices
    PE[:, 0::2] = np.sin(positions * div_term)
    PE[:, 1::2] = np.cos(positions * div_term)
    
    return PE

# ─── DEMO ───
PE = sinusoidal_positional_encoding(max_len=50, d_model=16)
print("PE shape:", PE.shape)
print("\nPosition 0:", np.round(PE[0], 3))
print("Position 1:", np.round(PE[1], 3))
print("Position 2:", np.round(PE[2], 3))

# Verify the rotation property: PE(pos+k) is a linear transform of PE(pos)
pos, k, dim = 5, 3, 0
omega = 1.0 / (10000 ** (2 * dim / 16))
# PE(pos+k, 2*dim) should equal sin(omega*(pos+k))
expected = np.sin(omega * (pos + k))
# Using rotation: sin(w*pos)*cos(w*k) + cos(w*pos)*sin(w*k)
from_rotation = PE[pos, 2*dim] * np.cos(omega*k) + PE[pos, 2*dim+1] * np.sin(omega*k)
print(f"\nRotation property check:")
print(f"  PE({pos+k}, {2*dim}) = {PE[pos+k, 2*dim]:.6f}")
print(f"  Via rotation:        = {from_rotation:.6f}")
print(f"  Match: {np.isclose(PE[pos+k, 2*dim], from_rotation)}")
        

import numpy as np

class LayerNorm:
    """Layer Normalization."""
    def __init__(self, d_model, eps=1e-6):
        self.gamma = np.ones(d_model)
        self.beta = np.zeros(d_model)
        self.eps = eps
    
    def forward(self, x):
        mean = np.mean(x, axis=-1, keepdims=True)
        var = np.var(x, axis=-1, keepdims=True)
        x_norm = (x - mean) / np.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta

class FeedForward:
    """Position-wise Feed-Forward Network: FFN(x) = ReLU(xW1+b1)W2+b2"""
    def __init__(self, d_model, d_ff):
        scale = np.sqrt(2.0 / d_model)
        self.W1 = np.random.randn(d_model, d_ff) * scale
        self.b1 = np.zeros(d_ff)
        self.W2 = np.random.randn(d_ff, d_model) * scale
        self.b2 = np.zeros(d_model)
    
    def forward(self, x):
        hidden = np.maximum(0, x @ self.W1 + self.b1)  # ReLU
        return hidden @ self.W2 + self.b2

class TransformerBlock:
    """
    One complete Transformer encoder block:
    1. Multi-Head Self-Attention + Residual + LayerNorm
    2. Feed-Forward Network + Residual + LayerNorm
    """
    def __init__(self, d_model, num_heads, d_ff):
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.norm1 = LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm2 = LayerNorm(d_model)
    
    def forward(self, x, mask=None):
        # Sub-layer 1: Multi-Head Attention
        attn_output, attn_weights = self.attention.forward(x, x, x, mask)
        x = self.norm1.forward(x + attn_output)  # Residual + LayerNorm
        
        # Sub-layer 2: Feed-Forward
        ffn_output = self.ffn.forward(x)
        x = self.norm2.forward(x + ffn_output)    # Residual + LayerNorm
        
        return x, attn_weights

class TransformerEncoder:
    """Stack of N Transformer encoder blocks."""
    def __init__(self, num_layers, d_model, num_heads, d_ff, max_len=512):
        self.layers = [
            TransformerBlock(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ]
        self.PE = sinusoidal_positional_encoding(max_len, d_model)
    
    def forward(self, x):
        """x: (batch, seq_len, d_model)"""
        seq_len = x.shape[1]
        x = x + self.PE[:seq_len]  # Add positional encoding
        
        all_weights = []
        for layer in self.layers:
            x, weights = layer.forward(x)
            all_weights.append(weights)
        
        return x, all_weights

# ─── DEMO: 6-layer Transformer Encoder ───
np.random.seed(42)
batch, seq_len, d_model = 2, 10, 64
num_heads, d_ff, num_layers = 8, 256, 6

encoder = TransformerEncoder(num_layers, d_model, num_heads, d_ff)
X = np.random.randn(batch, seq_len, d_model)

output, all_weights = encoder.forward(X)

print(f"Input:  {X.shape}")          # (2, 10, 64)
print(f"Output: {output.shape}")      # (2, 10, 64)
print(f"Layers: {len(all_weights)}")  # 6
print(f"Attn weights per layer: {all_weights[0].shape}")  # (2, 8, 10, 10)
        

import tensorflow as tf
from transformers import TFBertModel, BertTokenizer
import numpy as np

# ─── 1. Load Pre-trained BERT ───
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
bert_model = TFBertModel.from_pretrained('bert-base-uncased')

# ─── 2. Build Sentiment Classifier ───
class BERTSentimentClassifier(tf.keras.Model):
    def __init__(self, num_classes=3, dropout_rate=0.3):
        super().__init__()
        self.bert = TFBertModel.from_pretrained('bert-base-uncased')
        self.dropout = tf.keras.layers.Dropout(dropout_rate)
        self.classifier = tf.keras.layers.Dense(
            num_classes, activation='softmax'
        )
    
    def call(self, input_ids, attention_mask, training=False):
        # Get BERT outputs
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            training=training
        )
        # Use [CLS] token representation (first token)
        cls_output = outputs.last_hidden_state[:, 0, :]  # (batch, 768)
        cls_output = self.dropout(cls_output, training=training)
        logits = self.classifier(cls_output)
        return logits

# ─── 3. Tokenize Dataset ───
def tokenize_data(texts, labels, max_length=128):
    """Tokenize texts for BERT input."""
    encodings = tokenizer(
        texts, 
        max_length=max_length,
        truncation=True,
        padding='max_length',
        return_tensors='tf'
    )
    dataset = tf.data.Dataset.from_tensor_slices((
        {
            'input_ids': encodings['input_ids'],
            'attention_mask': encodings['attention_mask']
        },
        labels
    ))
    return dataset

# ─── 4. Sample Data & Training ───
texts = [
    "This movie was absolutely wonderful!",
    "Terrible experience, worst film ever.",
    "It was okay, nothing special.",
    "I loved every minute of this masterpiece!",
    "Complete waste of time and money.",
    "Average movie with some good moments.",
]
labels = [2, 0, 1, 2, 0, 1]  # 0=negative, 1=neutral, 2=positive

dataset = tokenize_data(texts, labels)
dataset = dataset.batch(2).prefetch(tf.data.AUTOTUNE)

# ─── 5. Compile and Train ───
model = BERTSentimentClassifier(num_classes=3)

# Key: Use very small learning rate for BERT fine-tuning!
optimizer = tf.keras.optimizers.Adam(learning_rate=2e-5)
model.compile(
    optimizer=optimizer,
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# Fine-tune for 3 epochs (in practice, 3-5 is sufficient)
# model.fit(dataset, epochs=3)
print("Model built. Ready for fine-tuning!")
print(f"BERT parameters: ~110M")
print(f"Classifier head: {768 * 3 + 3} = {768*3+3} parameters")
        

import tensorflow as tf
import numpy as np

class MiniGPT(tf.keras.Model):
    """
    A minimal GPT-style decoder-only Transformer for text generation.
    Implements causal (masked) self-attention.
    """
    def __init__(self, vocab_size, d_model=128, num_heads=4,
                 d_ff=512, num_layers=4, max_len=256):
        super().__init__()
        self.d_model = d_model
        self.max_len = max_len
        
        # Token + Position Embeddings
        self.token_embed = tf.keras.layers.Embedding(vocab_size, d_model)
        self.pos_embed = tf.keras.layers.Embedding(max_len, d_model)
        
        # Transformer Decoder Blocks
        self.blocks = [
            self._decoder_block(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ]
        
        # Output projection
        self.ln_f = tf.keras.layers.LayerNormalization()
        self.head = tf.keras.layers.Dense(vocab_size)
    
    def _decoder_block(self, d_model, num_heads, d_ff):
        """Single decoder block with causal attention."""
        return {
            'attn': tf.keras.layers.MultiHeadAttention(
                num_heads=num_heads, key_dim=d_model // num_heads
            ),
            'ln1': tf.keras.layers.LayerNormalization(),
            'ffn': tf.keras.Sequential([
                tf.keras.layers.Dense(d_ff, activation='gelu'),
                tf.keras.layers.Dense(d_model),
            ]),
            'ln2': tf.keras.layers.LayerNormalization(),
        }
    
    def _causal_mask(self, seq_len):
        """Create causal mask: prevent attending to future tokens."""
        mask = tf.linalg.band_part(
            tf.ones((seq_len, seq_len)), -1, 0
        )
        return mask  # Lower triangular
    
    def call(self, x, training=False):
        batch_size, seq_len = tf.shape(x)[0], tf.shape(x)[1]
        
        # Embeddings
        positions = tf.range(seq_len)
        tok_emb = self.token_embed(x)           # (batch, seq, d_model)
        pos_emb = self.pos_embed(positions)      # (seq, d_model)
        h = tok_emb + pos_emb
        
        # Causal mask
        causal_mask = self._causal_mask(seq_len)
        
        # Pass through decoder blocks
        for block in self.blocks:
            # Pre-norm architecture (GPT-2 style)
            h_norm = block['ln1'](h)
            attn_out = block['attn'](
                query=h_norm, key=h_norm, value=h_norm,
                attention_mask=causal_mask, training=training
            )
            h = h + attn_out  # Residual
            
            h_norm = block['ln2'](h)
            ffn_out = block['ffn'](h_norm, training=training)
            h = h + ffn_out   # Residual
        
        h = self.ln_f(h)
        logits = self.head(h)  # (batch, seq, vocab_size)
        return logits
    
    def generate(self, start_tokens, max_new_tokens=50, temperature=0.8):
        """Autoregressive text generation."""
        tokens = tf.constant([start_tokens])
        
        for _ in range(max_new_tokens):
            # Crop to max_len
            crop = tokens[:, -self.max_len:]
            logits = self(crop, training=False)
            
            # Get logits for last position
            next_logits = logits[:, -1, :] / temperature
            
            # Sample from distribution
            next_token = tf.random.categorical(next_logits, 1)
            tokens = tf.concat([tokens, next_token], axis=1)
        
        return tokens.numpy()[0]

# ─── DEMO ───
vocab_size = 5000
model = MiniGPT(vocab_size=vocab_size, d_model=128, num_heads=4,
                d_ff=512, num_layers=4, max_len=256)

# Test forward pass
dummy_input = tf.constant([[1, 42, 100, 7, 88]])
logits = model(dummy_input)
print(f"Input shape:  {dummy_input.shape}")
print(f"Output shape: {logits.shape}")   # (1, 5, 5000)
print(f"Parameters:   {model.count_params():,}")

# Generate text (random tokens since untrained)
generated = model.generate([1, 42, 100], max_new_tokens=10)
print(f"Generated token IDs: {generated}")
        

import tensorflow as tf

class TransformerEncoderLayer(tf.keras.layers.Layer):
    """Production-quality Transformer encoder layer."""
    
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.mha = tf.keras.layers.MultiHeadAttention(
            num_heads=num_heads,
            key_dim=d_model // num_heads,
            dropout=dropout
        )
        self.ffn = tf.keras.Sequential([
            tf.keras.layers.Dense(d_ff, activation='relu'),
            tf.keras.layers.Dropout(dropout),
            tf.keras.layers.Dense(d_model),
            tf.keras.layers.Dropout(dropout),
        ])
        self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = tf.keras.layers.Dropout(dropout)
    
    def call(self, x, training=False, mask=None):
        # Multi-Head Self-Attention + Residual + Norm
        attn_output = self.mha(x, x, x, attention_mask=mask, training=training)
        attn_output = self.dropout1(attn_output, training=training)
        x = self.norm1(x + attn_output)
        
        # Feed-Forward + Residual + Norm
        ffn_output = self.ffn(x, training=training)
        x = self.norm2(x + ffn_output)
        
        return x

# Build 6-layer encoder
d_model, num_heads, d_ff = 512, 8, 2048
encoder_layers = [
    TransformerEncoderLayer(d_model, num_heads, d_ff)
    for _ in range(6)
]

# Test
x = tf.random.normal((2, 20, 512))  # batch=2, seq=20
for layer in encoder_layers:
    x = layer(x, training=True)
print(f"6-layer encoder output: {x.shape}")  # (2, 20, 512)
        

from sklearn.pipeline import Pipeline
from sklearn.svm import SVC
from sklearn.metrics import classification_report
from sklearn.model_selection import cross_val_score
import numpy as np

# ─── BERT as Feature Extractor for sklearn ───

class BERTFeatureExtractor:
    """Extract [CLS] embeddings from BERT for use with sklearn."""
    
    def __init__(self, model_name='bert-base-uncased', max_length=128):
        from transformers import BertTokenizer, TFBertModel
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = TFBertModel.from_pretrained(model_name)
        self.max_length = max_length
    
    def transform(self, texts):
        """Convert texts to BERT [CLS] embeddings (768-dim vectors)."""
        encodings = self.tokenizer(
            texts, max_length=self.max_length,
            truncation=True, padding='max_length',
            return_tensors='tf'
        )
        outputs = self.model(encodings, training=False)
        # Extract [CLS] token embedding
        cls_embeddings = outputs.last_hidden_state[:, 0, :].numpy()
        return cls_embeddings

# ─── Use with sklearn ───
# (Pseudocode — requires transformers & tensorflow installed)
"""
# Extract features
extractor = BERTFeatureExtractor()
X_train = extractor.transform(train_texts)  # (n, 768)
X_test = extractor.transform(test_texts)

# Train any sklearn classifier on BERT features!
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression

models = {
    'SVM': SVC(kernel='rbf', C=1.0),
    'LogReg': LogisticRegression(max_iter=1000),
    'RF': RandomForestClassifier(n_estimators=100),
}

for name, clf in models.items():
    scores = cross_val_score(clf, X_train, y_train, cv=5)
    print(f"{name}: {scores.mean():.4f} ± {scores.std():.4f}")
"""

# ─── Simulated Demo (no GPU needed) ───
np.random.seed(42)
n_samples = 200
X_simulated = np.random.randn(n_samples, 768)  # Simulated BERT features
y_simulated = (X_simulated[:, 0] + X_simulated[:, 1] > 0).astype(int)

from sklearn.svm import SVC
from sklearn.model_selection import cross_val_score

svm = SVC(kernel='rbf', C=1.0)
scores = cross_val_score(svm, X_simulated, y_simulated, cv=5)
print(f"SVM on BERT features: {scores.mean():.4f} ± {scores.std():.4f}")
        

"""
Mini Project 1: Hindi Sentiment Analysis using Multilingual BERT
"""
from transformers import (
    AutoTokenizer, TFAutoModelForSequenceClassification,
    DataCollatorWithPadding
)
import tensorflow as tf
import numpy as np

# ─── 1. Load Hindi-capable Model ───
MODEL_NAME = "ai4bharat/indic-bert"  # or "bert-base-multilingual-cased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = TFAutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=3  # positive, negative, neutral
)

# ─── 2. Sample Hindi Data ───
hindi_reviews = [
    {"text": "यह फिल्म बहुत अच्छी थी, मुझे बहुत पसंद आई", "label": 2},
    {"text": "बेकार फिल्म, समय की बर्बादी", "label": 0},
    {"text": "कहानी ठीक थी लेकिन अभिनय कमजोर था", "label": 1},
    {"text": "शानदार अभिनय और बेहतरीन संगीत", "label": 2},
    {"text": "इतनी खराब फिल्म मैंने कभी नहीं देखी", "label": 0},
    {"text": "औसत फिल्म, एक बार देख सकते हैं", "label": 1},
    {"text": "दिल को छू लेने वाली कहानी", "label": 2},
    {"text": "बोरिंग और लंबी फिल्म", "label": 0},
]

# ─── 3. Tokenize ───
texts = [r["text"] for r in hindi_reviews]
labels = [r["label"] for r in hindi_reviews]

encodings = tokenizer(
    texts, max_length=128, truncation=True,
    padding="max_length", return_tensors="tf"
)

dataset = tf.data.Dataset.from_tensor_slices((
    dict(encodings), labels
)).batch(4)

# ─── 4. Fine-tune ───
optimizer = tf.keras.optimizers.Adam(learning_rate=2e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=optimizer, loss=loss, metrics=["accuracy"])

print("Model ready for Hindi sentiment analysis!")
print(f"Vocab size: {tokenizer.vocab_size}")
print(f"Model parameters: {model.count_params():,}")

# ─── 5. Inference Function ───
def predict_sentiment(text):
    """Predict sentiment of Hindi text."""
    inputs = tokenizer(text, return_tensors="tf", 
                       max_length=128, truncation=True, padding="max_length")
    outputs = model(inputs)
    probs = tf.nn.softmax(outputs.logits, axis=-1)
    label_map = {0: "नकारात्मक (Negative)", 
                 1: "तटस्थ (Neutral)", 
                 2: "सकारात्मक (Positive)"}
    pred = tf.argmax(probs, axis=-1).numpy()[0]
    return label_map[pred], probs.numpy()[0]

# Test (before training — predictions will be random)
test = "यह फिल्म बहुत शानदार है"
label, probs = predict_sentiment(test)
print(f"\nInput: {test}")
print(f"Prediction: {label}")
print(f"Probabilities: {probs}")
        

"""
Mini Project 2: Character-Level Language Model using Transformer Decoder
Inspired by Andrej Karpathy's nanoGPT
"""
import tensorflow as tf
import numpy as np

# ─── 1. Data Preparation ───
text = """
India is a vast country with diverse cultures, languages, and traditions.
The Indian constitution recognizes 22 official languages.
Artificial intelligence is transforming India's technology landscape.
From Bengaluru to Mumbai, startups are building innovative AI solutions.
The future of AI in India is bright, with millions of developers.
"""

# Character-level tokenization
chars = sorted(list(set(text)))
vocab_size = len(chars)
char_to_idx = {ch: i for i, ch in enumerate(chars)}
idx_to_char = {i: ch for ch, i in char_to_idx.items()}

def encode(s): return [char_to_idx[c] for c in s]
def decode(l): return ''.join([idx_to_char[i] for i in l])

data = np.array(encode(text))
print(f"Vocabulary size: {vocab_size}")
print(f"Text length: {len(data)} characters")
print(f"Sample encoding: '{text[:10]}' → {encode(text[:10])}")

# ─── 2. Create Training Sequences ───
block_size = 32  # Context window
batch_size = 8

def create_dataset(data, block_size):
    X, Y = [], []
    for i in range(len(data) - block_size):
        X.append(data[i:i+block_size])
        Y.append(data[i+1:i+block_size+1])
    return np.array(X), np.array(Y)

X, Y = create_dataset(data, block_size)
dataset = tf.data.Dataset.from_tensor_slices((X, Y))
dataset = dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

# ─── 3. Build Mini Transformer LM ───
class CharTransformerLM(tf.keras.Model):
    def __init__(self, vocab_size, d_model=64, num_heads=4,
                 num_layers=3, d_ff=128, max_len=128):
        super().__init__()
        self.token_emb = tf.keras.layers.Embedding(vocab_size, d_model)
        self.pos_emb = tf.keras.layers.Embedding(max_len, d_model)
        
        self.blocks = []
        for _ in range(num_layers):
            self.blocks.append({
                'attn': tf.keras.layers.MultiHeadAttention(
                    num_heads=num_heads, key_dim=d_model//num_heads
                ),
                'ln1': tf.keras.layers.LayerNormalization(),
                'ffn': tf.keras.Sequential([
                    tf.keras.layers.Dense(d_ff, activation='gelu'),
                    tf.keras.layers.Dense(d_model)
                ]),
                'ln2': tf.keras.layers.LayerNormalization(),
            })
        
        self.ln_f = tf.keras.layers.LayerNormalization()
        self.head = tf.keras.layers.Dense(vocab_size)
    
    def call(self, x, training=False):
        B, T = tf.shape(x)[0], tf.shape(x)[1]
        
        tok = self.token_emb(x)
        pos = self.pos_emb(tf.range(T))
        h = tok + pos
        
        # Causal mask
        mask = tf.linalg.band_part(tf.ones((T, T)), -1, 0)
        
        for block in self.blocks:
            h_n = block['ln1'](h)
            attn = block['attn'](h_n, h_n, h_n,
                                  attention_mask=mask, training=training)
            h = h + attn
            h_n = block['ln2'](h)
            h = h + block['ffn'](h_n)
        
        h = self.ln_f(h)
        return self.head(h)
    
    def generate(self, start_text, max_tokens=100, temperature=0.8):
        tokens = encode(start_text)
        for _ in range(max_tokens):
            x = tf.constant([tokens[-block_size:]])
            logits = self(x, training=False)
            next_logits = logits[0, -1, :] / temperature
            next_token = tf.random.categorical(
                next_logits[tf.newaxis, :], 1
            )[0, 0].numpy()
            tokens.append(next_token)
        return decode(tokens)

# ─── 4. Train ───
model = CharTransformerLM(vocab_size)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam(1e-3)

model.compile(optimizer=optimizer, loss=loss_fn)
print(f"\nModel parameters: {model.count_params():,}")

# Train for a few epochs
# model.fit(dataset, epochs=50, verbose=1)

# Generate text
# generated = model.generate("India is", max_tokens=50)
# print(f"Generated: {generated}")
print("\nModel built successfully! Ready for training.")