🧠 Attention Mechanics & Multi-Head Linear Algebra Playground¶

import os
import sys
import sys
from pathlib import Path

try:
    import google.colab
    IN_COLAB = True
except ImportError:
    IN_COLAB = False

if IN_COLAB:
    print("🤖 Running in Google Colab. Setting up environment...")
    !pip install -q numpy scikit-learn matplotlib loguru
    
    # Clone repository to access custom local source modules
    if not os.path.exists("production-ai-systems"):
        !git clone https://github.com/RogerioLS/production-ai-systems.git
    
    sys.path.append("production-ai-systems")
    print("🎉 Colab environment successfully initialized!")
# Ensure project root is in the path when running locally
if not IN_COLAB:
    # Localiza dinamicamente a raiz do repositório procurando por 'pyproject.toml'
    current_dir = Path.cwd()
    root_dir = current_dir
    for parent in [current_dir] + list(current_dir.parents):
        if (parent / "pyproject.toml").exists() or (parent / "projects").exists():
            root_dir = parent
            break

    if str(root_dir) not in sys.path:
        sys.path.insert(0, str(root_dir))

    print(f"📁 Project root added to sys.path: {root_dir}")

# Ensure project root is in the path when running locally
if not IN_COLAB:
    import sys
    import os
    sys.path.append(os.path.abspath("../../../"))

import torch
import matplotlib.pyplot as plt

from projects.a_llm_basics.src.lab_03_attention.attention import (
    ScaledDotProductAttention,
    MultiHeadAttention,
)
print("Modules imported successfully!")

# Setup synthetic sequence
tokens = ["The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"]
seq_len = len(tokens)
d_model = 16

torch.manual_seed(42)
q = torch.randn(1, seq_len, d_model)
k = torch.randn(1, seq_len, d_model)
v = torch.randn(1, seq_len, d_model)

attention = ScaledDotProductAttention()

# 1. Unmasked Attention
out_bidir, weights_bidir = attention(q, k, v)
print("Bidirectional output shape:", out_bidir.shape)
print("Bidirectional weights shape:", weights_bidir.shape)

# 2. Causal Masked Attention
mask = torch.tril(torch.ones(seq_len, seq_len)).bool()
out_causal, weights_causal = attention(q, k, v, mask=mask)
print("Causal output shape:", out_causal.shape)
print("Causal weights shape:", weights_causal.shape)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

# Bidirectional
im1 = ax1.imshow(weights_bidir.squeeze(0).detach().numpy(), cmap="viridis", vmin=0, vmax=1)
ax1.set_title("Bidirectional Self-Attention Weights", fontsize=14, fontweight="bold")
ax1.set_xticks(range(seq_len))
ax1.set_yticks(range(seq_len))
ax1.set_xticklabels(tokens, rotation=45, ha="right")
ax1.set_yticklabels(tokens)
fig.colorbar(im1, ax=ax1, fraction=0.046, pad=0.04)

# Causal
im2 = ax2.imshow(weights_causal.squeeze(0).detach().numpy(), cmap="viridis", vmin=0, vmax=1)
ax2.set_title("Causal Masked Self-Attention Weights", fontsize=14, fontweight="bold")
ax2.set_xticks(range(seq_len))
ax2.set_yticks(range(seq_len))
ax2.set_xticklabels(tokens, rotation=45, ha="right")
ax2.set_yticklabels(tokens)
fig.colorbar(im2, ax=ax2, fraction=0.046, pad=0.04)

plt.tight_layout()
plt.show()