From Chaos to Focus: Understanding Transformers & Attention 🎯

# Simple Attention in a Classroom
students = {
    "Alice": "I love science and helped with the volcano",
    "Bob": "I was playing games during the project",
    "Charlie": "I brought the materials for the experiment",
    "Diana": "I was absent that day",
    "Eve": "I explained how volcanoes work"
}

question = "Who helped with the science project?"

# Attention scores (how relevant each student is)
attention_scores = {
    "Alice": 0.9,    # Very relevant! 
    "Bob": 0.1,      # Not relevant
    "Charlie": 0.8,  # Quite relevant!
    "Diana": 0.0,    # Not relevant at all
    "Eve": 0.85      # Very relevant!
}

import numpy as np
import matplotlib.pyplot as plt

class SimpleAttention:
    """The simplest attention mechanism ever!"""
    
    def __init__(self):
        self.history = []  # To visualize later
    
    def calculate_attention(self, query, keys, values):
        """
        Query: What we're looking for
        Keys: What each position contains  
        Values: The actual information
        """
        # Step 1: Calculate similarity scores
        scores = []
        for key in keys:
            similarity = self.similarity(query, key)
            scores.append(similarity)
        
        # Step 2: Convert to percentages (softmax)
        attention_weights = self.softmax(scores)
        
        # Step 3: Weighted sum of values
        result = np.zeros_like(values[0], dtype=float)  # <-- Fix here!
        for i, value in enumerate(values):
            result += attention_weights[i] * value
        
        # Save for visualization
        self.history.append(attention_weights)
        
        return result, attention_weights
                    
    
    def similarity(self, a, b):
        """Simple dot product similarity"""
        return np.dot(a, b)
    
    def softmax(self, scores):
        """Convert scores to probabilities"""
        exp_scores = np.exp(scores - np.max(scores))  # Subtract max for stability
        return exp_scores / exp_scores.sum()
    
    def visualize(self, words):
        """Show attention as a heatmap"""
        if not self.history:
            return
            
        plt.figure(figsize=(8, 6))
        plt.imshow([self.history[-1]], cmap='YlOrRd', aspect='auto')
        plt.colorbar(label='Attention Weight')
        plt.xticks(range(len(words)), words)
        plt.yticks([0], ['Query'])
        plt.title('Where is the model paying attention?')
        
        # Add values on cells
        for i, weight in enumerate(self.history[-1]):
            plt.text(i, 0, f'{weight:.2f}', ha='center', va='center')
        
        plt.tight_layout()
        plt.show()

# Let's try it!
attention = SimpleAttention()

# Example: "The cat sat on the mat" - looking for what the cat did
words = ["The", "cat", "sat", "on", "the", "mat"]

# Simple word vectors (in real transformers, these are learned)
word_vectors = {
    "The": np.array([1, 0, 0, 0]),
    "cat": np.array([0, 1, 0, 0]),
    "sat": np.array([0, 0, 1, 0]),
    "on": np.array([0, 0, 0, 1]),
    "mat": np.array([1, 0, 0, 1])
}

# Query: "What did the cat do?"
query = word_vectors["cat"] + word_vectors["sat"] * 0.5  # Looking for cat + action

# Keys and values are our words
keys = [word_vectors.get(w, np.zeros(4)) for w in words]
values = keys  # In simple attention, keys and values can be the same

# Calculate attention!
result, weights = attention.calculate_attention(query, keys, values)

print("🎯 Attention weights for 'What did the cat do?':")
for word, weight in zip(words, weights):
    bar = "█" * int(weight * 20)
    print(f"{word:5} {bar} {weight:.2%}")

attention.visualize(words)

class LibraryAttention:
    """Attention explained as a library system"""
    
    def __init__(self):
        self.library = {
            # Book (Key) -> Information (Value)
            "Harry Potter": "A story about a wizard boy",
            "Cooking 101": "How to make delicious food",
            "Space Atlas": "Facts about planets and stars",
            "Python Guide": "Learn to code step by step"
        }
    
    def find_books(self, query):
        """
        Query: What you're looking for
        Keys: Book titles
        Values: Book contents
        """
        print(f"🔍 Query: '{query}'")
        print("\n📚 Checking each book:")
        
        scores = {}
        for book_title, book_content in self.library.items():
            # How relevant is this book?
            score = self.calculate_relevance(query, book_title, book_content)
            scores[book_title] = score
            print(f"   {book_title}: {score:.2f} relevance")
        
        # Get the most relevant books
        total_score = sum(scores.values())
        print("\n📖 Reading from books:")
        
        final_answer = ""
        for book, score in scores.items():
            if score > 0:
                weight = score / total_score
                print(f"   {book}: {weight:.1%} of attention")
                final_answer += f"({weight:.1%} from {book}) "
        
        return final_answer
    
    def calculate_relevance(self, query, key, value):
        """Simple word matching for relevance"""
        query_words = query.lower().split()
        key_words = key.lower().split()
        value_words = value.lower().split()
        
        score = 0
        for word in query_words:
            if word in key_words:
                score += 2  # Title match is important
            if word in value_words:
                score += 1  # Content match
        
        return score

# Try our library attention!
library = LibraryAttention()
library.find_books("How to cook in space")
print("-" * 50)
library.find_books("Learn magic spells")

class MultiHeadAttention:
    """Multiple attention heads working together"""
    
    def __init__(self, num_heads=4):
        self.num_heads = num_heads
        self.heads = []
        
        # Create specialized attention heads
        head_types = ["WHO", "WHAT", "WHERE", "WHEN"]
        for i in range(num_heads):
            head = {
                "name": head_types[i % len(head_types)],
                "attention": SimpleAttention()
            }
            self.heads.append(head)
    
    def analyze_sentence(self, sentence):
        """Each head looks for different things"""
        words = sentence.split()
        
        print(f"📝 Analyzing: '{sentence}'")
        print(f"👥 Using {self.num_heads} attention heads:\n")
        
        # Each head focuses on different aspects
        for i, head in enumerate(self.heads):
            print(f"🕵️ Head {i+1} ({head['name']}) is looking for {head['name'].lower()} information:")
            
            # Simulate different focus patterns
            if head['name'] == "WHO":
                # Focus on nouns/names
                focus_words = [w for w in words if w[0].isupper()]
            elif head['name'] == "WHAT":
                # Focus on verbs/actions
                focus_words = [w for w in words if w.endswith('ed') or w.endswith('ing')]
            elif head['name'] == "WHERE":
                # Focus on locations
                focus_words = [w for w in words if w in ['in', 'on', 'at', 'under', 'over']]
            else:  # WHEN
                # Focus on time words
                focus_words = [w for w in words if w in ['yesterday', 'today', 'tomorrow', 'now']]
            
            if focus_words:
                print(f"   Found: {', '.join(focus_words)}")
            else:
                print(f"   No specific {head['name'].lower()} information found")
            print()
        
        return f"Analyzed by {self.num_heads} different perspectives!"

# Demo multi-head attention
multi_attention = MultiHeadAttention(num_heads=4)
multi_attention.analyze_sentence("Alice helped Bob yesterday in the library")
print("=" * 60)
multi_attention.analyze_sentence("The cat is sleeping on the mat")

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

class SimpleTransformerBlock(nn.Module):
    """A real (but simple) Transformer block!"""
    
    def __init__(self, d_model=64, num_heads=4, d_ff=256, dropout=0.1):
        super().__init__()
        
        # Multi-head attention
        self.attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True
        )
        
        # Feed-forward network (the "thinking" part)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        
        # Layer normalization (keeps numbers stable)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # Dropout (prevents overfitting)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # Step 1: Multi-head attention
        attn_output, attention_weights = self.attention(x, x, x, attn_mask=mask)
        
        # Step 2: Add & Norm (residual connection)
        x = self.norm1(x + self.dropout(attn_output))
        
        # Step 3: Feed-forward
        ff_output = self.feed_forward(x)
        
        # Step 4: Add & Norm again
        x = self.norm2(x + self.dropout(ff_output))
        
        return x, attention_weights

# Let's create and test our transformer!
def demo_transformer():
    # Create a simple transformer
    transformer = SimpleTransformerBlock(d_model=64, num_heads=4)
    
    # Create some fake input (batch_size=1, seq_len=6, d_model=64)
    words = ["The", "cat", "sat", "on", "the", "mat"]
    fake_embeddings = torch.randn(1, len(words), 64)
    
    # Run through transformer
    output, attention_weights = transformer(fake_embeddings)
    
    print("🤖 Transformer Block Demo:")
    print(f"   Input shape: {fake_embeddings.shape}")
    print(f"   Output shape: {output.shape}")
    print(f"   Attention shape: {attention_weights.shape}")
    
    # Visualize attention from one head
    avg_attention = attention_weights[0].detach().numpy()  # Shape: (6, 6)
    
    plt.figure(figsize=(8, 6))
    plt.imshow(avg_attention, cmap='Blues', aspect='auto')
    plt.colorbar(label='Attention Weight')
    plt.xticks(range(len(words)), words)
    plt.yticks(range(len(words)), words)
    plt.title('Transformer Attention Pattern')
    plt.xlabel('Attending to')
    plt.ylabel('From position')
    
    # Add grid
    for i in range(len(words)):
        for j in range(len(words)):
            plt.text(j, i, f'{avg_attention[i,j]:.2f}', 
                    ha='center', va='center', fontsize=8)
    
    plt.tight_layout()
    plt.show()

demo_transformer()

class ToyTransformer(nn.Module):
    """A complete but tiny Transformer for learning!"""
    
    def __init__(self, vocab_size=100, d_model=64, num_heads=4, num_layers=2):
        super().__init__()
        
        # Token embeddings (words to vectors)
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        
        # Positional encoding (so model knows word order)
        self.positional_encoding = self.create_positional_encoding(1000, d_model)
        
        # Stack of encoder layers
        self.encoder_layers = nn.ModuleList([
            SimpleTransformerBlock(d_model, num_heads)
            for _ in range(num_layers)
        ])
        
        # Output projection
        self.output_projection = nn.Linear(d_model, vocab_size)
        
        # Dropout
        self.dropout = nn.Dropout(0.1)
    
    def create_positional_encoding(self, max_len, d_model):
        """Create position encodings (the magic that tells position!)"""
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        
        # Create div_term for the sinusoidal pattern
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                           -(torch.log(torch.tensor(10000.0)) / d_model))
        
        # Apply sin to even indices
        pe[:, 0::2] = torch.sin(position * div_term)
        
        # Apply cos to odd indices  
        pe[:, 1::2] = torch.cos(position * div_term)
        
        return pe
    
    def forward(self, x):
        # Get sequence length
        seq_len = x.shape[1]
        
        # Step 1: Convert tokens to embeddings
        x = self.token_embedding(x)
        
        # Step 2: Add positional encoding
        x = x + self.positional_encoding[:seq_len, :].unsqueeze(0)
        x = self.dropout(x)
        
        # Step 3: Pass through encoder layers
        attention_maps = []
        for layer in self.encoder_layers:
            x, attn = layer(x)
            attention_maps.append(attn)
        
        # Step 4: Project to vocabulary
        output = self.output_projection(x)
        
        return output, attention_maps
    
    def generate_text(self, start_tokens, max_length=20):
        """Generate text autoregressively"""
        self.eval()
        generated = start_tokens.clone()
        
        with torch.no_grad():
            for _ in range(max_length):
                # Get predictions
                output, _ = self(generated)
                
                # Get next token (greedy decoding)
                next_token = output[0, -1, :].argmax().unsqueeze(0).unsqueeze(0)
                
                # Append to sequence
                generated = torch.cat([generated, next_token], dim=1)
        
        return generated

# Demo: Create a tiny transformer
print("🏗️ Building a Tiny Transformer...")

# Create vocabulary (simple number-based)
vocab_size = 50
model = ToyTransformer(vocab_size=vocab_size, d_model=32, num_heads=2, num_layers=2)

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
print(f"📊 Total parameters: {total_params:,}")

# Show model architecture
print("\n🏛️ Model Architecture:")
for name, module in model.named_children():
    print(f"   {name}: {module.__class__.__name__}")

class SequenceReverser:
    """Train a transformer to reverse sequences"""
    
    def __init__(self, max_len=10):
        self.max_len = max_len
        self.model = ToyTransformer(
            vocab_size=20,  # 0-9 digits + special tokens
            d_model=32,
            num_heads=2,
            num_layers=2
        )
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        self.criterion = nn.CrossEntropyLoss()
        
    def generate_data(self, num_samples=100):
        """Generate sequence reversal data"""
        data = []
        
        for _ in range(num_samples):
            # Random length sequence
            length = torch.randint(3, self.max_len, (1,)).item()
            
            # Random sequence of digits
            sequence = torch.randint(0, 10, (length,))
            
            # Reversed sequence
            reversed_seq = torch.flip(sequence, [0])
            
            data.append((sequence, reversed_seq))
        
        return data
    
    def train(self, epochs=50):
        """Train the transformer"""
        print("🏃 Training Sequence Reverser...")
        
        losses = []
        
        for epoch in range(epochs):
            epoch_loss = 0
            data = self.generate_data(100)
            
            for input_seq, target_seq in data:
                # Add batch dimension
                input_seq = input_seq.unsqueeze(0)
                target_seq = target_seq.unsqueeze(0)
                
                # Forward pass
                output, _ = self.model(input_seq)
                
                # Calculate loss
                loss = self.criterion(
                    output.view(-1, output.size(-1)),
                    target_seq.view(-1)
                )
                
                # Backward pass
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                
                epoch_loss += loss.item()
            
            avg_loss = epoch_loss / len(data)
            losses.append(avg_loss)
            
            if (epoch + 1) % 10 == 0:
                print(f"   Epoch {epoch+1}: Loss = {avg_loss:.4f}")
                self.test()
        
        return losses
    
    def test(self):
        """Test the model"""
        self.model.eval()
        
        # Test sequence
        test_seq = torch.tensor([1, 2, 3, 4, 5])
        print(f"   Test: {test_seq.tolist()} → ", end="")
        
        with torch.no_grad():
            output, attention = self.model(test_seq.unsqueeze(0))
            predicted = output.argmax(dim=-1).squeeze()
            print(f"{predicted.tolist()}")
        
        self.model.train()

# Train our reverser!
reverser = SequenceReverser()
losses = reverser.train(epochs=100)

# Plot training progress
plt.figure(figsize=(10, 4))
plt.plot(losses)
plt.title('Training Progress: Learning to Reverse')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.grid(True)
plt.show()

class AttentionVisualizer:
    """Interactive attention visualization"""
    
    def __init__(self):
        self.sentences = [
            "The quick brown fox jumps over the lazy dog",
            "Attention is all you need for understanding",
            "Transformers revolutionized natural language processing",
            "The cat sat on the mat and looked around"
        ]
    
    def create_attention_pattern(self, sentence, pattern_type="self"):
        """Create different attention patterns"""
        words = sentence.split()
        n = len(words)
        
        if pattern_type == "self":
            # Each word attends to itself mostly
            attention = np.eye(n) * 0.7 + np.ones((n, n)) * 0.3 / n
            
        elif pattern_type == "forward":
            # Words attend to previous words (causal)
            attention = np.tril(np.ones((n, n)))
            # Normalize rows
            attention = attention / attention.sum(axis=1, keepdims=True)
            
        elif pattern_type == "backward":
            # Words attend to future words
            attention = np.triu(np.ones((n, n)))
            attention = attention / attention.sum(axis=1, keepdims=True)
            
        elif pattern_type == "middle":
            # Everything attends to middle word
            attention = np.zeros((n, n))
            middle = n // 2
            attention[:, middle] = 1.0
            
        elif pattern_type == "edges":
            # Attend to first and last words
            attention = np.zeros((n, n))
            attention[:, 0] = 0.5
            attention[:, -1] = 0.5
            
        return attention, words
    
    def visualize_all_patterns(self):
        """Show all attention patterns"""
        patterns = ["self", "forward", "backward", "middle", "edges"]
        
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        axes = axes.flatten()
        
        sentence = self.sentences[0]
        
        for idx, pattern in enumerate(patterns):
            attention, words = self.create_attention_pattern(sentence, pattern)
            
            ax = axes[idx]
            im = ax.imshow(attention, cmap='YlOrRd', aspect='auto', vmin=0, vmax=1)
            
            # Labels
            ax.set_xticks(range(len(words)))
            ax.set_yticks(range(len(words)))
            ax.set_xticklabels(words, rotation=45, ha='right')
            ax.set_yticklabels(words)
            ax.set_title(f'{pattern.capitalize()} Attention')
            
            # Add text annotations for small matrices
            if len(words) < 10:
                for i in range(len(words)):
                    for j in range(len(words)):
                        text = ax.text(j, i, f'{attention[i, j]:.2f}',
                                     ha="center", va="center", fontsize=8)
        
        # Remove empty subplot
        axes[-1].axis('off')
        
        plt.suptitle('Different Types of Attention Patterns', fontsize=16)
        plt.tight_layout()
        plt.show()
    
    def interactive_attention(self):
        """Create an interactive attention demo"""
        print("🎮 Interactive Attention Demo!")
        print("="*50)
        
        for i, sentence in enumerate(self.sentences):
            print(f"\n{i+1}. {sentence}")
        
        choice = int(input("\nChoose a sentence (1-4): ")) - 1
        sentence = self.sentences[choice]
        
        print("\nAttention Patterns:")
        print("1. Self (diagonal)")
        print("2. Forward (causal)")  
        print("3. Backward (anti-causal)")
        print("4. Middle (center focus)")
        print("5. Edges (boundary focus)")
        
        pattern_choice = int(input("\nChoose pattern (1-5): "))
        patterns = ["self", "forward", "backward", "middle", "edges"]
        pattern = patterns[pattern_choice - 1]
        
        attention, words = self.create_attention_pattern(sentence, pattern)
        
        # Visualize
        plt.figure(figsize=(10, 8))
        plt.imshow(attention, cmap='YlOrRd', aspect='auto')
        plt.colorbar(label='Attention Weight')
        plt.xticks(range(len(words)), words, rotation=45, ha='right')
        plt.yticks(range(len(words)), words)
        plt.title(f'{pattern.capitalize()} Attention Pattern')
        plt.xlabel('Attending to')
        plt.ylabel('From word')
        
        # Add values
        for i in range(len(words)):
            for j in range(len(words)):
                plt.text(j, i, f'{attention[i, j]:.2f}',
                        ha="center", va="center", fontsize=8)
        
        plt.tight_layout()
        plt.show()

# Create visualizer
viz = AttentionVisualizer()
viz.visualize_all_patterns()

def compare_architectures():
    """Compare different architectures"""
    
    print("🏁 Architecture Race!\n")
    
    architectures = {
        "RNN": {
            "speed": "🐌 Slow (sequential)",
            "memory": "😅 Forgets long sequences",
            "parallelization": "❌ Can't parallelize",
            "long_range": "😵 Struggles with long text",
            "training": "😴 Slow training"
        },
        "CNN": {
            "speed": "🏃 Fast (parallel)",
            "memory": "🤷 Fixed context window",
            "parallelization": "✅ Highly parallel",
            "long_range": "😐 Limited range",
            "training": "😊 Fast training"
        },
        "Transformer": {
            "speed": "🚀 Super fast (parallel)",
            "memory": "🧠 Sees everything",
            "parallelization": "✅✅ Fully parallel",
            "long_range": "💪 Handles long text easily",
            "training": "⚡ Fast with right hardware"
        }
    }
    
    for arch, props in architectures.items():
        print(f"📊 {arch}:")
        for prop, value in props.items():
            print(f"   {prop}: {value}")
        print()
    
    # Visual comparison
    categories = ['Speed', 'Memory', 'Parallel', 'Range', 'Training']
    
    rnn_scores = [2, 3, 1, 2, 2]
    cnn_scores = [4, 3, 5, 3, 4]
    transformer_scores = [5, 5, 5, 5, 4]
    
    x = np.arange(len(categories))
    width = 0.25
    
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(x - width, rnn_scores, width, label='RNN', color='coral')
    ax.bar(x, cnn_scores, width, label='CNN', color='skyblue')
    ax.bar(x + width, transformer_scores, width, label='Transformer', color='gold')
    
    ax.set_xlabel('Capabilities')
    ax.set_ylabel('Score (1-5)')
    ax.set_title('Architecture Comparison')
    ax.set_xticks(x)
    ax.set_xticklabels(categories)
    ax.legend()
    ax.grid(axis='y', alpha=0.3)
    
    plt.tight_layout()
    plt.show()

compare_architectures()

class TransformerApplications:
    """See Transformers in action!"""
    
    def __init__(self):
        self.applications = {
            "Translation": {
                "input": "Hello, how are you?",
                "output": "Hola, ¿cómo estás?",
                "attention_focus": ["Hello→Hola", "are→estás", "you→tú"]
            },
            "Summarization": {
                "input": "The cat sat on the mat. It was sunny. The cat was happy.",
                "output": "Happy cat sits on mat in sunshine.",
                "attention_focus": ["cat→cat", "happy→happy", "sunny→sunshine"]
            },
            "Question Answering": {
                "input": "Context: Paris is the capital of France. Question: What is the capital of France?",
                "output": "Paris",
                "attention_focus": ["capital→capital", "France→France", "Paris→Paris"]
            },
            "Code Generation": {
                "input": "# Function to calculate factorial",
                "output": "def factorial(n):\n    if n <= 1:\n        return 1\n    return n * factorial(n-1)",
                "attention_focus": ["factorial→factorial", "calculate→def", "Function→def"]
            }
        }
    
    def demonstrate_all(self):
        """Show all applications"""
        print("🌟 Transformer Applications Gallery\n")
        
        for app_name, details in self.applications.items():
            print(f"📱 {app_name}:")
            print(f"   Input:  '{details['input']}'")
            print(f"   Output: '{details['output']}'")
            print(f"   Key Attention: {' | '.join(details['attention_focus'])}")
            print()
    
    def attention_in_action(self, task="Translation"):
        """Visualize how attention works for specific task"""
        details = self.applications[task]
        
        print(f"\n🔍 Deep Dive: {task}")
        print("="*50)
        print(f"Input: {details['input']}")
        print(f"Output: {details['output']}")
        print("\n🎯 Where the model pays attention:")
        
        for attention_pair in details['attention_focus']:
            source, target = attention_pair.split('→')
            print(f"   When generating '{target}' → looks at '{source}'")

# Demo applications
apps = TransformerApplications()
apps.demonstrate_all()
apps.attention_in_action("Translation")

class SimpleChatBot:
    """A toy chatbot to understand transformers"""
    
    def __init__(self):
        self.knowledge_base = {
            "greeting": {
                "patterns": ["hello", "hi", "hey", "greetings"],
                "responses": ["Hello! How can I help you?", "Hi there!", "Greetings!"]
            },
            "weather": {
                "patterns": ["weather", "rain", "sunny", "cold"],
                "responses": ["It's a beautiful day!", "Perfect weather for learning!"]
            },
            "learning": {
                "patterns": ["learn", "study", "understand", "teach"],
                "responses": ["Learning is amazing!", "Let's learn together!"]
            },
            "transformer": {
                "patterns": ["transformer", "attention", "neural"],
                "responses": ["Transformers use attention to understand!", "Attention is all you need!"]
            }
        }
        
        # Simulate transformer components
        self.attention_scores = {}
    
    def calculate_attention(self, query, knowledge):
        """Calculate which knowledge to pay attention to"""
        scores = {}
        
        query_words = query.lower().split()
        
        for topic, data in knowledge.items():
            score = 0
            for pattern in data["patterns"]:
                if pattern in query.lower():
                    score += 1
            scores[topic] = score
        
        # Normalize scores (softmax-style)
        total = sum(scores.values()) + 0.001  # Avoid division by zero
        normalized_scores = {k: v/total for k, v in scores.items()}
        
        return normalized_scores
    
    def generate_response(self, query):
        """Generate response using attention mechanism"""
        print(f"\n🤔 Processing: '{query}'")
        
        # Calculate attention
        attention = self.calculate_attention(query, self.knowledge_base)
        
        # Show attention weights
        print("\n📊 Attention Distribution:")
        for topic, score in attention.items():
            bar = "█" * int(score * 20)
            print(f"   {topic:12} {bar} {score:.2%}")
        
        # Pick highest attention topic
        best_topic = max(attention.items(), key=lambda x: x[1])[0]
        
        if attention[best_topic] > 0.3:  # Threshold
            import random
            response = random.choice(self.knowledge_base[best_topic]["responses"])
        else:
            response = "I'm not sure about that. Can you tell me more?"
        
        print(f"\n💬 Response: {response}")
        return response
    
    def chat(self):
        """Interactive chat loop"""
        print("🤖 Simple Transformer ChatBot")
        print("Type 'quit' to exit\n")
        
        while True:
            user_input = input("You: ")
            if user_input.lower() == 'quit':
                print("Bot: Goodbye! Keep learning about transformers! 👋")
                break
            
            self.generate_response(user_input)
            print()

# Create and run chatbot
bot = SimpleChatBot()
# Uncomment to run interactive chat:
# bot.chat()

# Demo some conversations
demo_queries = [
    "Hello! How are you?",
    "What's the weather like?",
    "I want to learn about transformers",
    "Can you teach me attention mechanism?"
]

for query in demo_queries:
    bot.generate_response(query)
    print("="*50)

class AttentionMathExplained:
    """The math of attention, explained simply!"""
    
    def __init__(self):
        print("🧮 Attention Math Demystified!\n")
    
    def explain_attention_formula(self):
        """Break down the scary attention formula"""
        print("The Famous Attention Formula:")
        print("Attention(Q,K,V) = softmax(QK^T / √d_k)V")
        print("\nLet's break it down:\n")
        
        steps = {
            "Q (Query)": "What you're looking for",
            "K (Key)": "What each word contains", 
            "V (Value)": "The actual information",
            "QK^T": "How similar Q is to each K (dot product)",
            "√d_k": "Scaling factor (keeps numbers reasonable)",
            "softmax": "Convert to percentages (sum to 1)",
            "×V": "Weighted average of values"
        }
        
        for symbol, meaning in steps.items():
            print(f"   {symbol:8} = {meaning}")
    
    def step_by_step_example(self):
        """Walk through attention calculation"""
        print("\n\n📝 Step-by-Step Example:")
        print("Sentence: 'The cat sat'")
        
        # Simple 2D vectors for illustration
        words = ["The", "cat", "sat"]
        vectors = {
            "The": np.array([1, 0]),
            "cat": np.array([0, 1]),
            "sat": np.array([1, 1])
        }
        
        print("\n1️⃣ Word Vectors:")
        for word, vec in vectors.items():
            print(f"   {word}: {vec}")
        
        # Let's say we're looking from "cat"'s perspective
        query = vectors["cat"]
        keys = np.array([vectors[w] for w in words])
        values = keys  # Same as keys for simplicity
        
        print(f"\n2️⃣ Query (from 'cat'): {query}")
        
        # Calculate QK^T (dot products)
        scores = np.dot(keys, query)
        print(f"\n3️⃣ Similarity Scores (QK^T):")
        for i, (word, score) in enumerate(zip(words, scores)):
            print(f"   cat · {word} = {score}")
        
        # Scale by sqrt(d_k)
        d_k = 2  # Dimension of vectors
        scaled_scores = scores / np.sqrt(d_k)
        print(f"\n4️⃣ Scaled Scores (÷√{d_k}):")
        for word, score in zip(words, scaled_scores):
            print(f"   {word}: {score:.3f}")
        
        # Softmax
        exp_scores = np.exp(scaled_scores)
        softmax_scores = exp_scores / exp_scores.sum()
        print(f"\n5️⃣ Softmax (percentages):")
        for word, score in zip(words, softmax_scores):
            print(f"   {word}: {score:.2%}")
        
        # Final weighted sum
        output = np.sum(values * softmax_scores.reshape(-1, 1), axis=0)
        print(f"\n6️⃣ Final Output: {output}")
        
        return softmax_scores
    
    def visualize_calculation(self, attention_weights):
        """Visualize the calculation process"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        
        # Attention weights
        words = ["The", "cat", "sat"]
        ax1.bar(words, attention_weights)
        ax1.set_title("Attention Weights from 'cat'")
        ax1.set_ylabel("Weight")
        ax1.set_ylim([0, 1])
        
        # Attention mechanism diagram
        ax2.text(0.5, 0.9, "Attention Mechanism", ha='center', fontsize=14, weight='bold')
        
        # Draw boxes and arrows
        boxes = {
            'Q': (0.2, 0.7),
            'K': (0.5, 0.7), 
            'V': (0.8, 0.7),
            'QK^T': (0.35, 0.5),
            'Softmax': (0.35, 0.3),
            'Output': (0.5, 0.1)
        }
        
        for label, (x, y) in boxes.items():
            if label in ['Q', 'K', 'V']:
                ax2.add_patch(plt.Rectangle((x-0.05, y-0.05), 0.1, 0.1, 
                                          fill=True, facecolor='lightblue'))
            ax2.text(x, y, label, ha='center', va='center')
        
        # Draw arrows
        arrows = [
            ((0.25, 0.65), (0.32, 0.55)),  # Q to QK^T
            ((0.5, 0.65), (0.38, 0.55)),   # K to QK^T
            ((0.35, 0.45), (0.35, 0.35)),  # QK^T to Softmax
            ((0.35, 0.25), (0.45, 0.15)),  # Softmax to Output
            ((0.75, 0.65), (0.55, 0.15))   # V to Output
        ]
        
        for start, end in arrows:
            ax2.annotate('', xy=end, xytext=start,
                        arrowprops=dict(arrowstyle='->', lw=2))
        
        ax2.set_xlim([0, 1])
        ax2.set_ylim([0, 1])
        ax2.axis('off')
        
        plt.tight_layout()
        plt.show()

# Run the math explanation
math_demo = AttentionMathExplained()
math_demo.explain_attention_formula()
weights = math_demo.step_by_step_example()
math_demo.visualize_calculation(weights)

class PositionalEncodingDemo:
    """Understanding positional encoding"""
    
    def __init__(self):
        self.max_len = 20
        self.d_model = 64
    
    def explain_why_needed(self):
        """Why do we need positional encoding?"""
        print("🤔 Why Positional Encoding?\n")
        
        # Without position
        print("Without position info:")
        print("  'cat sat mat' = 'mat sat cat' = 'sat cat mat' 😵")
        print("  The model can't tell the difference!\n")
        
        # With position
        print("With position info:")
        print("  'cat[1] sat[2] mat[3]' ≠ 'mat[1] sat[2] cat[3]' ✅")
        print("  Now the model knows the order!")
    
    def visualize_encoding(self):
        """Visualize the sinusoidal pattern"""
        # Create positional encoding
        pe = torch.zeros(self.max_len, self.d_model)
        position = torch.arange(0, self.max_len).unsqueeze(1).float()
        
        div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * 
                           -(np.log(10000.0) / self.d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # Visualize
        plt.figure(figsize=(12, 8))
        
        # Show full encoding matrix
        plt.subplot(2, 1, 1)
        plt.imshow(pe.numpy(), cmap='RdBu', aspect='auto')
        plt.colorbar(label='Encoding Value')
        plt.xlabel('Encoding Dimension')
        plt.ylabel('Position in Sequence')
        plt.title('Positional Encoding Matrix (Sinusoidal Pattern)')
        
        # Show encoding for specific positions
        plt.subplot(2, 1, 2)
        positions_to_show = [0, 5, 10, 15]
        for pos in positions_to_show:
            plt.plot(pe[pos, :32], label=f'Position {pos}', alpha=0.8)
        
        plt.xlabel('Encoding Dimension')
        plt.ylabel('Value')
        plt.title('Positional Encoding for Different Positions')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.show()
    
    def demonstrate_uniqueness(self):
        """Show that each position has unique encoding"""
        print("\n🔍 Checking Uniqueness of Positions:")
        
        # Create simple encoding
        positions = torch.arange(5)
        
        # Simple encoding: just the position number
        simple_encoding = positions.unsqueeze(1).float()
        
        # Sinusoidal encoding (simplified)
        sin_encoding = torch.sin(positions.unsqueeze(1) * 0.5)
        cos_encoding = torch.cos(positions.unsqueeze(1) * 0.5)
        complex_encoding = torch.cat([sin_encoding, cos_encoding], dim=1)
        
        print("\nSimple Encoding (just position number):")
        for i, enc in enumerate(simple_encoding):
            print(f"   Position {i}: {enc.item():.1f}")
        
        print("\nSinusoidal Encoding (sin + cos):")
        for i, enc in enumerate(complex_encoding):
            print(f"   Position {i}: [{enc[0].item():.3f}, {enc[1].item():.3f}]")
        
        print("\n✨ Each position has a unique 'signature'!")

# Run positional encoding demo
pos_demo = PositionalEncodingDemo()
pos_demo.explain_why_needed()
pos_demo.visualize_encoding()
pos_demo.demonstrate_uniqueness()

class TransformerTrainingTips:
    """Pro tips for training transformers"""
    
    def __init__(self):
        self.tips = {
            "Learning Rate Warmup": {
                "why": "Transformers are sensitive at the start",
                "how": "Start with tiny LR, gradually increase",
                "visual": self.plot_warmup_schedule
            },
            "Gradient Clipping": {
                "why": "Prevents exploding gradients", 
                "how": "Clip gradients to max norm (e.g., 1.0)",
                "visual": self.plot_gradient_clipping
            },
            "Label Smoothing": {
                "why": "Prevents overconfidence",
                "how": "Replace hard labels with soft targets",
                "visual": self.plot_label_smoothing
            },
            "Dropout Everywhere": {
                "why": "Prevents overfitting",
                "how": "Add dropout to attention, FFN, embeddings",
                "visual": None
            }
        }
    
    def plot_warmup_schedule(self):
        """Visualize learning rate warmup"""
        steps = np.arange(0, 1000)
        warmup_steps = 100
        
        # Linear warmup
        lr = np.minimum(steps / warmup_steps, 1.0)
        
        # Then decay
        lr = lr * (1 / np.sqrt(np.maximum(steps, warmup_steps)))
        
        plt.figure(figsize=(10, 4))
        plt.plot(steps, lr)
        plt.axvline(x=warmup_steps, color='red', linestyle='--', label='End of warmup')
        plt.xlabel('Training Steps')
        plt.ylabel('Learning Rate')
        plt.title('Learning Rate Schedule with Warmup')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.show()
    
    def plot_gradient_clipping(self):
        """Visualize gradient clipping"""
        # Generate random gradients
        gradients = np.random.normal(0, 2, 1000)
        gradients[::100] = np.random.normal(0, 10, 10)  # Add some outliers
        
        # Clip gradients
        max_norm = 1.0
        clipped = np.clip(gradients, -max_norm, max_norm)
        
        plt.figure(figsize=(10, 4))
        plt.subplot(1, 2, 1)
        plt.hist(gradients, bins=50, alpha=0.7, color='red')
        plt.title('Original Gradients (with outliers)')
        plt.xlabel('Gradient Value')
        
        plt.subplot(1, 2, 2)
        plt.hist(clipped, bins=50, alpha=0.7, color='green')
        plt.title(f'Clipped Gradients (max={max_norm})')
        plt.xlabel('Gradient Value')
        
        plt.tight_layout()
        plt.show()
    
    def plot_label_smoothing(self):
        """Visualize label smoothing"""
        classes = ['Cat', 'Dog', 'Bird', 'Fish']
        
        # Hard labels
        hard_labels = [1.0, 0.0, 0.0, 0.0]
        
        # Smoothed labels (ε = 0.1)
        epsilon = 0.1
        smooth_labels = [(1 - epsilon) if x == 1 else epsilon/3 for x in hard_labels]
        
        x = np.arange(len(classes))
        width = 0.35
        
        plt.figure(figsize=(10, 5))
        plt.bar(x - width/2, hard_labels, width, label='Hard Labels', color='red', alpha=0.7)
        plt.bar(x + width/2, smooth_labels, width, label='Smooth Labels', color='blue', alpha=0.7)
        
        plt.xlabel('Classes')
        plt.ylabel('Probability')
        plt.title('Label Smoothing Effect')
        plt.xticks(x, classes)
        plt.legend()
        plt.grid(axis='y', alpha=0.3)
        plt.show()
    
    def show_all_tips(self):
        """Display all training tips"""
        print("🎯 Transformer Training Pro Tips:\n")
        
        for i, (tip_name, details) in enumerate(self.tips.items(), 1):
            print(f"{i}. {tip_name}")
            print(f"   Why: {details['why']}")
            print(f"   How: {details['how']}\n")
            
            if details['visual']:
                details['visual']()

# Show training tips
tips = TransformerTrainingTips()
tips.show_all_tips()

def common_mistakes():
    """Common transformer mistakes and solutions"""
    
    mistakes = {
        "🐛 Forgetting Positional Encoding": {
            "symptom": "Model treats 'cat sat mat' same as 'mat cat sat'",
            "fix": "Always add positional encoding to embeddings!",
            "code": "x = embeddings + positional_encoding"
        },
        "💥 Attention Explosion": {
            "symptom": "Out of memory with long sequences",
            "fix": "Attention is O(n²) - use shorter sequences or efficient attention",
            "code": "max_seq_length = 512  # Don't go crazy!"
        },
        "🎲 No Masking in Training": {
            "symptom": "Model cheats by looking at future words",
            "fix": "Use causal masking for autoregressive models",
            "code": "mask = torch.triu(torch.ones(n, n), diagonal=1)"
        },
        "📉 Learning Rate Too High": {
            "symptom": "Loss explodes or oscillates wildly",
            "fix": "Use warmup and smaller learning rates",
            "code": "lr = 1e-4  # Transformers like it small"
        },
        "🧊 Frozen Attention": {
            "symptom": "All positions get equal attention",
            "fix": "Check initialization and temperature scaling",
            "code": "attention_weights = softmax(scores / temperature)"
        }
    }
    
    print("⚠️ Common Transformer Pitfalls:\n")
    
    for mistake, details in mistakes.items():
        print(f"{mistake}")
        print(f"   Symptom: {details['symptom']}")
        print(f"   Fix: {details['fix']}")
        print(f"   Code: {details['code']}")
        print()

common_mistakes()

def create_learning_path():
    """Your personalized transformer learning journey"""
    
    learning_path = {
        "🌱 Beginner (You are here!)": [
            "✅ Understand attention mechanism",
            "✅ Build simple transformer blocks",
            "✅ Train on toy tasks",
            "⬜ Implement full encoder-decoder"
        ],
        "🌿 Intermediate": [
            "⬜ Study different attention variants (sparse, local, etc.)",
            "⬜ Implement BERT-style masked language modeling",
            "⬜ Fine-tune pretrained models",
            "⬜ Build a small GPT from scratch"
        ],
        "🌳 Advanced": [
            "⬜ Implement efficient attention (Linear, Performer, etc.)",
            "⬜ Multi-modal transformers (vision + text)",
            "⬜ Train your own language model",
            "⬜ Research novel architectures"
        ],
        "🏔️ Expert": [
            "⬜ Contribute to open source (HuggingFace, etc.)",
            "⬜ Publish research papers",
            "⬜ Build production systems",
            "⬜ Teach others!"
        ]
    }
    
    print("🗺️ Your Transformer Learning Journey:\n")
    
    for level, tasks in learning_path.items():
        print(f"{level}")
        for task in tasks:
            print(f"   {task}")
        print()
    
    print("💡 Remember: Every expert was once a beginner!")
    print("🚀 Keep building, keep learning, keep experimenting!")

create_learning_path()

print("""
🏆 FINAL CHALLENGE: Build Your Own Mini-GPT!

Your mission:
1. Create a character-level transformer
2. Train it on your favorite book
3. Generate new text in that style!

Starter code available at: [link to colab]

Share your creation with #MyFirstTransformer

Remember: The transformer that powers ChatGPT started 
exactly like the simple models we built today!

Happy coding! 🚀
""")