Practice Questions & Tips

def train_bpe(corpus: list[str], num_merges: int) -> dict:
    """Train a BPE tokenizer on a corpus.

    Args:
        corpus: List of words (pre-tokenized)
        num_merges: Number of merge operations to learn

    Returns:
        Dictionary of merge rules {(token_a, token_b): merged_token}
    """
    # Step 1: Initialize vocabulary as character-level tokens
    # Each word becomes a tuple of characters + end-of-word marker
    vocab = {}
    for word in corpus:
        chars = tuple(list(word) + ["</w>"])
        vocab[chars] = vocab.get(chars, 0) + 1

    merges = {}

    for i in range(num_merges):
        # Step 2: Count all adjacent pairs
        pairs = {}
        for word, freq in vocab.items():
            for j in range(len(word) - 1):
                pair = (word[j], word[j + 1])
                pairs[pair] = pairs.get(pair, 0) + freq

        if not pairs:
            break

        # Step 3: Find the most frequent pair
        best_pair = max(pairs, key=pairs.get)
        merges[best_pair] = best_pair[0] + best_pair[1]

        # Step 4: Merge that pair everywhere in the vocabulary
        new_vocab = {}
        for word, freq in vocab.items():
            new_word = []
            j = 0
            while j < len(word):
                if j < len(word) - 1 and (word[j], word[j + 1]) == best_pair:
                    new_word.append(best_pair[0] + best_pair[1])
                    j += 2
                else:
                    new_word.append(word[j])
                    j += 1
            new_vocab[tuple(new_word)] = freq
        vocab = new_vocab

    return merges


# Test it
corpus = ["low"] * 5 + ["lowest"] * 2 + ["newer"] * 6 + ["wider"] * 3
merges = train_bpe(corpus, num_merges=10)
print("Learned merges:", merges)

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        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)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Project to Q, K, V
        Q = self.W_q(x)  # (batch, seq_len, d_model)
        K = self.W_k(x)
        V = self.W_v(x)

        # Reshape for multi-head: (batch, num_heads, seq_len, d_k)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

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

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

        # Softmax + weighted sum
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V)

        # Reshape back: (batch, seq_len, d_model)
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

        return self.W_o(output)


# Test it
attn = SelfAttention(d_model=512, num_heads=8)
x = torch.randn(2, 10, 512)  # batch=2, seq_len=10, d_model=512
output = attn(x)
print(f"Input shape: {x.shape}, Output shape: {output.shape}")
# Both should be torch.Size([2, 10, 512])

from sentence_transformers import SentenceTransformer
import numpy as np

class SimpleRAG:
    def __init__(self, documents: list[str]):
        self.documents = documents
        self.encoder = SentenceTransformer("all-MiniLM-L6-v2")
        # Encode all documents at initialization
        self.doc_embeddings = self.encoder.encode(documents, normalize_embeddings=True)

    def retrieve(self, query: str, top_k: int = 3) -> list[tuple[str, float]]:
        """Retrieve top-k most relevant documents for a query."""
        query_embedding = self.encoder.encode([query], normalize_embeddings=True)

        # Cosine similarity (embeddings are normalized, so dot product = cosine)
        similarities = np.dot(self.doc_embeddings, query_embedding.T).flatten()

        # Get top-k indices
        top_indices = np.argsort(similarities)[::-1][:top_k]

        return [(self.documents[i], float(similarities[i])) for i in top_indices]

    def generate_prompt(self, query: str, top_k: int = 3) -> str:
        """Create a RAG prompt with retrieved context."""
        retrieved = self.retrieve(query, top_k)

        context = "\n\n".join([f"[Source {i+1}]: {doc}"
                               for i, (doc, score) in enumerate(retrieved)])

        prompt = f"""Answer the question based ONLY on the provided context.
If the context does not contain the answer, say "I don't have enough information."

Context:
{context}

Question: {query}

Answer:"""
        return prompt


# Usage
docs = [
    "BERT uses masked language modeling where 15% of tokens are masked.",
    "GPT models use causal language modeling with left-to-right attention.",
    "LoRA adds low-rank matrices to transformer weights for efficient fine-tuning.",
    "RAG combines retrieval with generation to ground LLM outputs in facts.",
]
rag = SimpleRAG(docs)
prompt = rag.generate_prompt("How does BERT train?")
print(prompt)