PyTorch & HuggingFace Snippets

Key code patterns for loading, embedding, training, and deploying models

0. Quick Reference -- Common Syntax

NumPy

import numpy as np

# Creating arrays
a = np.array([1, 2, 3])                  # from list
np.ones((2, 3))                           # 2x3 of ones
np.random.randn(5, 3)                     # 5x3 standard normal
np.random.randint(0, 10, size=(3,))       # 3 random ints in [0, 10)
np.arange(0, 10, 2)                       # [0, 2, 4, 6, 8]
np.linspace(0, 1, 5)                      # [0, 0.25, 0.5, 0.75, 1.0]
np.eye(3)                                 # 3x3 identity

# Shape
a.reshape(1, 3)                           # reshape to (1, 3)
a[:, None]                                # add dim: (3,) -> (3, 1)
np.squeeze(a)                             # remove dims of size 1

# Indexing
mask = a > 1                              # [False, True, True]
a[mask]                                   # [2, 3] -- boolean indexing
np.where(a > 1, a, 0)                     # [0, 2, 3] -- conditional
np.argsort(a)                             # indices that sort the array
np.argmax(a)                              # index of max

PyTorch

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

# Creating tensors
t = torch.tensor([1.0, 2.0, 3.0])        # from list
torch.randn(5, 3)                         # 5x3 standard normal
torch.from_numpy(np_array)                # numpy -> tensor
t.numpy()                                 # tensor -> numpy (shares memory)

# Shape
t.view(1, -1)                             # reshape (-1 infers size)
t.unsqueeze(0)                            # add dim: (3,) -> (1, 3)
t.squeeze()                               # remove size-1 dims

# Common operations
t.argmax(dim=-1)                          # index of max
F.softmax(t, dim=-1)                      # softmax
F.normalize(t, p=2, dim=-1)              # L2 normalize
F.cosine_similarity(a, b, dim=-1)         # cosine sim
torch.topk(t, k=3, dim=-1)               # top-k values + indices

# Gradient control
with torch.no_grad():                     # disable grad tracking
    out = model(x)
param.requires_grad = False               # freeze a parameter

# Common layers
nn.Linear(in_features, out_features)      # fully connected
nn.Embedding(vocab_size, embed_dim)       # embedding lookup
nn.Dropout(p=0.1)                         # dropout
nn.CrossEntropyLoss()                     # classification loss
nn.MSELoss()                              # regression loss

1. Data Preprocessing

Load, Split, Handle Imbalance

import pandas as pd
import torch
from sklearn.model_selection import train_test_split
from torch.utils.data import WeightedRandomSampler, DataLoader

# Load and inspect
df = pd.read_csv("dataset.csv")
print(df.shape, df.dtypes, df.isnull().sum())
print(df["label"].value_counts())

# 70/15/15 stratified split
train_df, temp_df = train_test_split(df, test_size=0.3, random_state=42, stratify=df["label"])
val_df, test_df = train_test_split(temp_df, test_size=0.5, random_state=42, stratify=temp_df["label"])
print(f"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}")

# Handle class imbalance -- Option A: weighted loss
class_counts = [8000, 2000]  # [negative, positive]
weights = torch.tensor([1.0 / c for c in class_counts])
weights = weights / weights.sum()
loss_fn = torch.nn.CrossEntropyLoss(weight=weights)

# Handle class imbalance -- Option B: oversample minority
labels = [0] * 8000 + [1] * 2000
sample_weights = [1.0 / class_counts[l] for l in labels]
sampler = WeightedRandomSampler(sample_weights, num_samples=len(labels), replacement=True)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
Always inspect first: df.describe(), df.isnull().sum(), df["label"].value_counts(). Two ways to handle imbalance: weighted loss (penalize mistakes on minority class more) or oversampling (show minority class more often).

2. Building a Classification Pipeline

End-to-End Spam Classifier with Pretrained Transformer

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from sklearn.metrics import precision_recall_fscore_support
from sklearn.model_selection import train_test_split

texts = ["Buy cheap meds now!", "Meeting at 3pm tomorrow", "You won $1M!", "Project update attached"]
labels = [1, 0, 1, 0]  # 1=spam, 0=ham
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.25, random_state=42)

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

class SpamDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.encodings = tokenizer(texts, truncation=True, padding="max_length", max_length=max_len, return_tensors="pt")
        self.labels = torch.tensor(labels)

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

    def __getitem__(self, idx):
        # pull input_ids and attention_mask for this sample
        item = {}
        for key, val in self.encodings.items():
            item[key] = val[idx]               # each is (seq_len,)
        return item, self.labels[idx]          # dict, scalar

train_loader = DataLoader(SpamDataset(train_texts, train_labels, tokenizer), batch_size=2, shuffle=True)
val_loader = DataLoader(SpamDataset(val_texts, val_labels, tokenizer), batch_size=2)

model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
loss_fn = torch.nn.CrossEntropyLoss()

for epoch in range(3):
    model.train()
    for batch, lbl in train_loader:                     # batch: dict of (B, seq_len), lbl: (B,)
        optimizer.zero_grad()
        out = model(**batch)                             # out.logits: (B, num_labels)
        loss = loss_fn(out.logits, lbl)
        loss.backward()
        optimizer.step()

    model.eval()
    all_preds, all_labels = [], []
    with torch.no_grad():
        for batch, lbl in val_loader:
            preds = model(**batch).logits.argmax(dim=-1)  # (B,)
            all_preds.extend(preds.tolist())
            all_labels.extend(lbl.tolist())
    p, r, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average="binary")
    print(f"Epoch {epoch+1}: P={p:.3f} R={r:.3f} F1={f1:.3f}")

Sklearn Baseline for Comparison

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report

pipe = Pipeline([
    ("tfidf", TfidfVectorizer(max_features=10000, ngram_range=(1, 2))),
    ("clf", LogisticRegression(max_iter=1000)),
])
pipe.fit(train_texts, train_labels)
preds = pipe.predict(val_texts)
print(classification_report(val_labels, preds, target_names=["ham", "spam"]))
Always build a simple baseline first (TF-IDF + LogReg). If the transformer doesn't beat it significantly, your data or labels may be the problem, not the model.

3. K-Means Clustering

K-Means from Scratch (NumPy)

import numpy as np

def kmeans(X, k, max_iters=100, seed=42):
    rng = np.random.RandomState(seed)
    n, d = X.shape                         # n=num points, d=num features

    # Random initialization: pick k random points as centroids
    indices = rng.choice(n, size=k, replace=False)  # (k,)
    centroids = X[indices]                           # (k, d)

    for _ in range(max_iters):
        # Assign each point to nearest centroid
        dists = np.linalg.norm(X[:, None] - centroids[None], axis=2)  # (n, k)
        assignments = np.argmin(dists, axis=1)                         # (n,)

        # Update centroids to mean of assigned points
        new_centroids = np.empty_like(centroids)     # (k, d)
        for j in range(k):
            members = X[assignments == j]            # (num_in_cluster_j, d)
            if len(members) > 0:
                new_centroids[j] = members.mean(axis=0)  # (d,)
            else:
                new_centroids[j] = centroids[j]

        if np.allclose(centroids, new_centroids):
            break
        centroids = new_centroids

    return centroids, assignments                    # (k, d), (n,)

X = np.random.randn(200, 5)                         # 200 points, 5 features
centroids, labels = kmeans(X, k=3)
for i in range(3):
    print(f"Cluster {i}: {np.sum(labels == i)} points")

Embedding Clustering

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("all-MiniLM-L6-v2")
texts = [
    "How to train a neural network", "Backpropagation explained", "Gradient descent variants",
    "Best Italian restaurants nearby", "Pizza dough recipe", "Pasta carbonara tips",
    "Stock market predictions", "Portfolio diversification", "Bond yields explained",
]
embeddings = model.encode(texts, convert_to_numpy=True)  # (9, 384)

# Use our kmeans from above
centroids, labels = kmeans(embeddings, k=3)              # (3, 384), (9,)
for i in range(3):
    members = [texts[j] for j in range(len(texts)) if labels[j] == i]
    print(f"Cluster {i}: {members}")
K-Means algorithm: 1) pick k random centroids, 2) assign points to nearest centroid, 3) recompute centroids as cluster means, 4) repeat until convergence. For high-dimensional embeddings, apply PCA or UMAP first for better clustering and visualization.

4. Logistic & Linear Regression from Scratch

Logistic Regression (NumPy)

import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))            # maps any real number to (0, 1)

def logistic_regression(X, y, lr=0.1, epochs=1000):
    n, d = X.shape                         # n=samples, d=features
    W = np.zeros(d)                        # (d,) -- one weight per feature
    b = 0.0                                # scalar bias

    for _ in range(epochs):
        z = X @ W + b                      # (n,) -- linear combination
        preds = sigmoid(z)                 # (n,) -- probabilities between 0 and 1
        error = preds - y                  # (n,) -- how far off each prediction is

        dW = X.T @ error / n               # (d,) -- gradient for weights
        db = np.mean(error)                # scalar -- gradient for bias

        W = W - lr * dW
        b = b - lr * db

    return W, b                            # (d,), scalar

def predict(X, W, b):
    probs = sigmoid(X @ W + b)             # (n,)
    return (probs >= 0.5).astype(int)      # (n,) -- 0 or 1

np.random.seed(42)
X = np.random.randn(200, 3)               # 200 samples, 3 features
y = (X @ [1, -2, 0.5] + np.random.randn(200) * 0.5 > 0).astype(float)  # binary labels

W, b = logistic_regression(X, y, lr=0.5, epochs=2000)
acc = (predict(X, W, b) == y).mean()
print(f"Accuracy: {acc:.3f}, bias: {b:.3f}")

Linear Regression (NumPy)

import numpy as np

def linear_regression_gd(X, y, lr=0.01, epochs=1000):
    n, d = X.shape
    W = np.zeros(d)                        # (d,)
    b = 0.0                                # scalar

    for _ in range(epochs):
        preds = X @ W + b                  # (n,)
        error = preds - y                  # (n,)

        dW = X.T @ error / n               # (d,)
        db = np.mean(error)                # scalar

        W = W - lr * dW
        b = b - lr * db

    return W, b

# Closed-form (normal equation) -- for reference
def linear_regression_closed(X, y):
    """W = (X^T X)^{-1} X^T y -- solves in one step, no iterations."""
    # Here we use the hstack trick to fold bias into W
    ones = np.ones((X.shape[0], 1))        # (n, 1)
    X_b = np.hstack([ones, X])             # (n, d+1) -- prepend column of 1s
    W_b = np.linalg.pinv(X_b.T @ X_b) @ X_b.T @ y  # (d+1,)
    return W_b[1:], W_b[0]                 # weights (d,), bias scalar

np.random.seed(42)
X = np.random.randn(100, 2)               # 100 samples, 2 features
y = X @ [3, -1.5] + 2 + np.random.randn(100) * 0.3  # true: W=[3, -1.5], b=2

W_gd, b_gd = linear_regression_gd(X, y, lr=0.1, epochs=2000)
W_cf, b_cf = linear_regression_closed(X, y)
print(f"GD:          bias={b_gd:.2f}, weights={W_gd}")
print(f"Closed-form: bias={b_cf:.2f}, weights={W_cf}")
Key gradient formulas (same structure for both):
dW = X.T @ error / n    db = mean(error)
Where error = preds - y. Only difference: logistic wraps the prediction in sigmoid, linear does not.

5. Models & Tokenization

Load Model + Tokenize

from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")
model.eval()

# Single text
tokens = tokenizer("Hello world", return_tensors="pt", padding=True, truncation=True, max_length=128)
output = model(**tokens)
print(output.last_hidden_state.shape)  # (1, seq_len, 768)

# Batch
texts = ["What is the weather?", "Remind me to call mom.", "Play my morning playlist."]
batch = tokenizer(texts, return_tensors="pt", padding="longest", truncation=True, max_length=64)
print(batch["input_ids"].shape)  # (3, longest_seq_len)

Sentence Embedding Model

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
sentences = ["Samsung Bixby is a voice assistant", "Memory systems store user preferences"]
embeddings = model.encode(sentences, convert_to_tensor=True)
print(embeddings.shape)  # (2, 384)
AutoModel -- base transformer, returns hidden states. For embeddings or custom heads.
AutoModelForSequenceClassification -- adds classification head on [CLS]. For sentiment, intent.
AutoModelForCausalLM -- LM head for next-token prediction. For text generation.
attention_mask -- 1 = real token, 0 = padding. Always pass to model.

6. Generating Embeddings

Get [CLS] Token Embedding from BERT

import torch
from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")
model.eval()

text = "User prefers jazz music in the evening."
tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)

with torch.no_grad():
    outputs = model(**tokens)

cls_embedding =outputs.last_hidden_state[:, 0, :]  # [CLS]
print(cls_embedding.shape)  # (1, 768)

Mean-Pooled Sentence Embedding (Proper Way)

import torch
from transformers import AutoModel, AutoTokenizer

def mean_pooling(model_output, attention_mask):
    """Mean pooling ignoring padding tokens."""
    hidden = model_output.last_hidden_state          # (batch, seq_len, dim)
    mask = attention_mask.unsqueeze(-1).float()       # (batch, seq_len, 1)
    mask = mask.expand(hidden.size())                 # (batch, seq_len, dim)
    summed = torch.sum(hidden * mask, dim=1)          # (batch, dim)
    counts = torch.clamp(mask.sum(dim=1), min=1e-9)   # (batch, dim)
    return summed / counts

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")
model.eval()

text = "Schedule a meeting with the design team."
tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)

with torch.no_grad():
    outputs = model(**tokens)

embedding = mean_pooling(outputs, tokens["attention_mask"])
print(embedding.shape)  # (1, 768)

Batch Embedding with torch.no_grad()

import torch
from transformers import AutoModel, AutoTokenizer

def embed_batch(texts, model, tokenizer, max_length=128, device="cpu"):
    """Embed texts using mean pooling. Returns (N, hidden_dim)."""
    tokens = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=max_length)
    tokens = tokens.to(device)
    model.to(device)
    model.eval()

    with torch.no_grad():
        outputs = model(**tokens)

    # Mean pooling: average token embeddings, ignoring padding
    hidden = outputs.last_hidden_state                        # (N, seq_len, dim)
    mask = tokens["attention_mask"].unsqueeze(-1).float()      # (N, seq_len, 1)
    mask = mask.expand(hidden.size())                          # (N, seq_len, dim)
    summed = torch.sum(hidden * mask, dim=1)                   # (N, dim)
    counts = torch.clamp(mask.sum(dim=1), min=1e-9)            # (N, dim)
    embeddings = summed / counts

    # L2 normalize
    return torch.nn.functional.normalize(embeddings, p=2, dim=1)

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

texts = ["Turn off the lights", "Set alarm for 7am", "What is the weather?"]
embs = embed_batch(texts, model, tokenizer)
print(embs.shape)  # (3, 768)

Using sentence-transformers .encode()

from sentence_transformers import SentenceTransformer
import torch

model = SentenceTransformer("all-MiniLM-L6-v2")

texts =[
    "Play relaxing music",
    "Navigate to the nearest coffee shop",
    "Read my unread messages",
]

embeddings =model.encode(
    texts,
    convert_to_tensor=True,
    normalize_embeddings=True,
    batch_size=32,
    show_progress_bar=False,
)
print(embeddings.shape)  # (3, 384)
[CLS] vs Mean Pooling vs Max Pooling
[CLS] token -- fast, single vector. Works well when the model was pre-trained with [CLS] as a sentence-level representation (BERT's NSP). Can be noisy for retrieval.
Mean pooling -- average all token embeddings (excluding padding). Generally the best default for sentence embeddings. Used by sentence-transformers.
Max pooling -- take the max across the sequence dimension for each hidden dim. Can capture salient features but loses ordering information. Rarely the best choice.

7. Similarity Search with Embeddings

Cosine Similarity with torch.nn.functional

import torch
import torch.nn.functional as F

def pairwise_cosine_similarity(a, b):
    """a: (M, D), b: (N, D) -> (M, N) cosine similarities."""
    a_norm = F.normalize(a, p=2, dim=1)       # (M, D)
    b_norm = F.normalize(b, p=2, dim=1)       # (N, D)
    return torch.mm(a_norm, b_norm.t())        # (M, N)

query = torch.randn(1, 768)                   # 1 query, 768-dim
database = torch.randn(1000, 768)              # 1000 docs, 768-dim
similarities = pairwise_cosine_similarity(query, database)  # (1, 1000)
top_k_values, top_k_indices = torch.topk(similarities, k=5, dim=1)  # (1, 5)
print(top_k_indices)

Build a Simple FAISS Index (Flat L2 and IVF)

import numpy as np
import faiss

dim = 384
num_vectors = 100_000
embeddings = np.random.randn(num_vectors, dim).astype("float32")  # (100K, 384)
faiss.normalize_L2(embeddings)

# Flat index (exact brute-force search)
flat_index = faiss.IndexFlatIP(dim)
flat_index.add(embeddings)

# IVF index (approximate, partitions into Voronoi cells)
nlist = 256  # number of Voronoi cells
quantizer = faiss.IndexFlatIP(dim)
ivf_index = faiss.IndexIVFFlat(quantizer, dim, nlist, faiss.METRIC_INNER_PRODUCT)
ivf_index.train(embeddings)
ivf_index.add(embeddings)
ivf_index.nprobe = 16  # search 16 of 256 cells

query = np.random.randn(1, dim).astype("float32")  # (1, 384)
faiss.normalize_L2(query)

distances, indices = flat_index.search(query, k=5)  # distances: (1, 5), indices: (1, 5)
print(f"Flat results: {indices[0]}")

distances, indices = ivf_index.search(query, k=5)
print(f"IVF results: {indices[0]}")

Approximate Nearest Neighbor with FAISS (HNSW)

import numpy as np
import faiss

dim = 384
num_vectors = 100_000
embeddings = np.random.randn(num_vectors, dim).astype("float32")  # (100K, 384)
faiss.normalize_L2(embeddings)

hnsw_index = faiss.IndexHNSWFlat(dim, 32)  # M=32 neighbors per node
hnsw_index.hnsw.efConstruction = 200
hnsw_index.hnsw.efSearch = 64
hnsw_index.metric_type = faiss.METRIC_INNER_PRODUCT

hnsw_index.add(embeddings)

query = np.random.randn(1, dim).astype("float32")  # (1, 384)
faiss.normalize_L2(query)
distances, indices = hnsw_index.search(query, k=10)  # (1, 10), (1, 10)
print(f"HNSW top-10: {indices[0]}")
FAISS Index Types -- Trade-offs for On-Device
IndexFlatIP / IndexFlatL2 -- exact search, O(N) per query. Perfect for <100K vectors. No training needed.
IndexIVFFlat -- partitions space into Voronoi cells. Much faster search, slight accuracy loss. Needs training. Good for 100K-10M vectors.
IndexHNSWFlat -- graph-based. Excellent recall, fast search, but high memory (stores full vectors + graph). Good for on-device if memory allows.
IndexIVFPQ -- IVF + product quantization. Compresses vectors heavily (e.g., 384 floats -> 48 bytes). Best for on-device with tight RAM constraints on Galaxy phones.

8. Training Loop Essentials

Custom Dataset + DataLoader

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=128):
        self.encodings = tokenizer(texts, padding="max_length", truncation=True,
                                   max_length=max_length, return_tensors="pt")
        self.labels = torch.tensor(labels, dtype=torch.long)

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

    def __getitem__(self, idx):
        # returns input_ids (max_length,), attention_mask (max_length,), label (scalar)
        return self.encodings["input_ids"][idx], self.encodings["attention_mask"][idx], self.labels[idx]

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
dataset = TextDataset(["Set alarm", "Play jazz", "Weather?"], [0, 1, 2], tokenizer)
loader = DataLoader(dataset, batch_size=2, shuffle=True)

Training Loop with Validation and Early Stopping

import torch
import torch.nn as nn

def train_model(model, train_loader, val_loader, num_epochs=10, lr=2e-5, patience=3, device="cuda"):
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    best_val_loss = float("inf")
    patience_counter = 0

    for epoch in range(num_epochs):
        # Train
        model.train()
        train_loss = 0.0
        for input_ids, attention_mask, labels in train_loader:  # (B, seq_len), (B, seq_len), (B,)
            input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)

            optimizer.zero_grad()
            logits = model(input_ids, attention_mask)          # (B, num_classes)
            loss = criterion(logits, labels)                   # scalar
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            train_loss += loss.item()

        # Validate
        model.eval()
        val_loss, correct, total = 0.0, 0, 0
        with torch.no_grad():
            for input_ids, attention_mask, labels in val_loader:
                input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
                logits = model(input_ids, attention_mask)
                val_loss += criterion(logits, labels).item()
                correct += (logits.argmax(-1) == labels).sum().item()
                total += labels.size(0)

        avg_train = train_loss / len(train_loader)
        avg_val = val_loss / len(val_loader)
        print(f"Epoch {epoch+1}: train={avg_train:.4f} val={avg_val:.4f} acc={correct/total:.4f}")

        # Early stopping
        if avg_val < best_val_loss:
            best_val_loss = avg_val
            patience_counter = 0
            torch.save(model.state_dict(), "best_model.pt")
        else:
            patience_counter += 1
            if patience_counter >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break

    model.load_state_dict(torch.load("best_model.pt"))
    return model
Training loop pattern: zero_grad -> forward -> loss -> backward -> clip_grad -> step. Always validate after each epoch. Early stopping prevents overfitting: track best val loss and stop if it hasn't improved for patience epochs.

9. Two-Stage Retrieval (Bi-Encoder + Cross-Encoder)

Reranking with Cross-Encoder

import torch
from sentence_transformers import CrossEncoder

def retrieve_and_rerank(query, candidates, bi_encoder, cross_encoder, initial_k=20, final_k=5):
    """Two-stage retrieval: bi-encoder recall + cross-encoder reranking."""
    # Stage 1: bi-encoder retrieves top candidates
    query_emb = bi_encoder.encode([query], normalize_embeddings=True)   # (1, d)
    doc_embs = bi_encoder.encode(candidates, normalize_embeddings=True) # (N, d)
    scores = (query_emb @ doc_embs.T)[0]             # (N,)
    top_indices = scores.argsort()[::-1][:initial_k]  # (initial_k,)
    shortlist = [candidates[i] for i in top_indices]

    # Stage 2: cross-encoder reranks the shortlist
    pairs = [[query, doc] for doc in shortlist]       # list of [query, doc] pairs
    rerank_scores = cross_encoder.predict(pairs)      # (initial_k,) float scores

    # Sort by rerank score descending
    ranked = sorted(zip(shortlist, rerank_scores), key=lambda x: x[1], reverse=True)
    return ranked[:final_k]

from sentence_transformers import SentenceTransformer

bi_encoder = SentenceTransformer("all-MiniLM-L6-v2")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

docs = ["User likes jazz", "User likes rock", "User has a cat", "User works at Samsung"]
results = retrieve_and_rerank("What music does the user enjoy?", docs, bi_encoder, cross_encoder, initial_k=3, final_k=2)
for doc, score in results:
    print(f"{score:.4f} | {doc}")
Why Two-Stage Retrieval (Bi-Encoder + Cross-Encoder)
Bi-encoder -- encodes query and documents independently. Fast (precompute doc embeddings, just encode query at search time). But limited: no cross-attention between query and document tokens.
Cross-encoder -- concatenates query + document and processes jointly through the transformer. Much more accurate (full cross-attention), but O(N) per query since you cannot precompute.
Solution: Bi-encoder retrieves top-K candidates (e.g., 100 from 1M). Cross-encoder reranks those K. Best of both worlds: speed + accuracy.

10. Evaluation Metrics

Computing Retrieval Metrics (Recall@K, NDCG)

import torch
import numpy as np

def recall_at_k(retrieved_indices, relevant_indices, k):
    """Recall@K: fraction of relevant items found in top-K.
    retrieved_indices: (num_queries, num_retrieved) -- doc IDs per query
    relevant_indices:  list of lists -- ground truth relevant doc IDs per query
    """
    recalls = []
    for i in range(len(relevant_indices)):
        top_k = set(retrieved_indices[i, :k].tolist())  # set of k doc IDs
        relevant = set(relevant_indices[i])
        if len(relevant) == 0:
            continue
        hit_count = len(top_k & relevant)                # how many relevant found
        recalls.append(hit_count / len(relevant))
    return np.mean(recalls)

def ndcg_at_k(retrieved_indices, relevance_scores, k):
    """NDCG@K: normalized discounted cumulative gain.
    retrieved_indices: (num_queries, num_retrieved)
    relevance_scores:  list of dicts {doc_id: relevance_score} per query
    """
    ndcgs = []
    for i in range(len(relevance_scores)):
        # Compute DCG from retrieved order
        dcg = 0.0
        for rank in range(min(k, retrieved_indices.shape[1])):
            doc_id = retrieved_indices[i, rank].item()
            rel = relevance_scores[i].get(doc_id, 0.0)
            dcg += (2**rel - 1) / np.log2(rank + 2)

        # Compute ideal DCG (best possible ordering)
        ideal_rels = sorted(relevance_scores[i].values(), reverse=True)[:k]
        idcg = 0.0
        for j, r in enumerate(ideal_rels):
            idcg += (2**r - 1) / np.log2(j + 2)

        ndcgs.append(dcg / idcg if idcg > 0 else 0.0)
    return np.mean(ndcgs)

retrieved = torch.tensor([[3, 1, 4, 0, 2], [1, 0, 3, 2, 4]])
relevant = [[1, 3], [0, 1]]
print(f"Recall@3: {recall_at_k(retrieved, relevant, k=3):.4f}")

rel_scores = [{1: 2.0, 3: 1.0, 4: 0.5}, {0: 2.0, 1: 1.0}]
print(f"NDCG@3: {ndcg_at_k(retrieved, rel_scores, k=3):.4f}")

Perplexity Computation

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def compute_perplexity(model, tokenizer, text, device="cpu"):
    """Perplexity of text under a causal LM. Lower = more confident."""
    model.eval()
    model.to(device)
    encodings = tokenizer(text, return_tensors="pt").to(device)  # input_ids: (1, seq_len)

    with torch.no_grad():
        outputs = model(**encodings, labels=encodings["input_ids"])
        neg_log_likelihood = outputs.loss  # scalar -- avg cross-entropy over tokens

    return torch.exp(neg_log_likelihood).item()  # scalar

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

text = "Samsung Bixby is a voice assistant that remembers your preferences."
ppl = compute_perplexity(model, tokenizer, text)
print(f"Perplexity: {ppl:.2f}")
Recall@K -- did we find the relevant items in the top K? NDCG -- are relevant items ranked higher? Both are standard for evaluating retrieval/memory systems. Perplexity measures how confident the LM is in its predictions (lower = better).