03. Cross-Encoder Reranker Fine-Tuning — ms-marco-MiniLM for MangaAssist

graph TB
    subgraph "Input Construction"
        A["Query: 'dark isekai manga'"] --> C["[CLS] dark isekai manga [SEP]<br>Overlord - A dark fantasy isekai<br>where the protagonist... [SEP]"]
        B["Doc: 'Overlord - A dark fantasy...'"] --> C
    end

    subgraph "Tokenization"
        C --> D["WordPiece Tokens (max 128)<br>[CLS] dark is ##ek ##ai manga [SEP]<br>over ##lord - a dark fantasy<br>is ##ek ##ai where the... [SEP]"]
    end

    subgraph "Embedding (384-dim, ~12M params)"
        D --> E["Token Emb + Position Emb + Segment Emb<br>Segment A: query tokens<br>Segment B: document tokens"]
    end

    subgraph "6 Transformer Layers"
        E --> F0["Layer 0: Basic token features<br>LR: 3.2e-6 | Grad: 0.008×"]
        F0 --> F1["Layer 1-2: Syntactic interactions<br>Query-doc token alignment begins"]
        F1 --> F3["Layer 3-4: Semantic matching<br>'dark' ↔ 'dark fantasy' link formed<br>LR: 6.4e-6 to 1e-5 | Grad: 0.18-0.38×"]
        F3 --> F5["Layer 5: High-level relevance<br>'isekai' + 'dark fantasy' → relevant<br>LR: 1.6e-5 | Grad: 0.65×"]
    end

    subgraph "Relevance Scoring"
        F5 --> G["[CLS] Pooling<br>384-dim relevance vector"]
        G --> H["Linear(384 → 1)<br>LR: 2e-5 | Grad: 1.0×"]
        H --> I["Relevance Score: 0.94"]
    end

    style F0 fill:#e8f5e9
    style F1 fill:#fff9c4
    style F3 fill:#fff3e0
    style F5 fill:#ffccbc
    style H fill:#ef9a9a

graph TD
    subgraph "Before Fine-Tuning (Layer 5 Head 3)"
        direction LR
        Q1["dark<br>(query)"] -->|"0.15"| D1a["dark<br>(doc)"]
        Q1 -->|"0.12"| D1b["fantasy<br>(doc)"]
        Q1 -->|"0.10"| D1c["comedy<br>(doc)"]
        Q1 -->|"0.08"| D1d["adventure<br>(doc)"]
        Q1 -->|"0.55"| D1e["other tokens<br>(distributed)"]
    end

    subgraph "After Fine-Tuning (Layer 5 Head 3)"
        direction LR
        Q2["dark<br>(query)"] -->|"0.38"| D2a["dark<br>(doc)"]
        Q2 -->|"0.28"| D2b["fantasy<br>(doc)"]
        Q2 -->|"0.02"| D2c["comedy<br>(doc)"]
        Q2 -->|"0.04"| D2d["adventure<br>(doc)"]
        Q2 -->|"0.28"| D2e["other tokens"]
    end

graph LR
    subgraph "Standard Pairwise Loss"
        A["Pair (pos=1, neg=2)<br>Gradient: 0.50"] --> B["All pairs get<br>equal treatment"]
        C["Pair (pos=5, neg=8)<br>Gradient: 0.50"] --> B
    end

    subgraph "LambdaRank"
        D["Pair (pos=1, neg=2)<br>|ΔNDCG| = 0.37<br>Gradient: 0.50 × 0.37 = 0.185"] --> E["Top positions get<br>12× more gradient<br>than bottom"]
        F["Pair (pos=5, neg=8)<br>|ΔNDCG| = 0.03<br>Gradient: 0.50 × 0.03 = 0.015"] --> E
    end

    subgraph "Effect on MangaAssist"
        E --> G["Position 1 error:<br>Wrong manga at top<br>= user leaves"]
        E --> H["Position 8 error:<br>Wrong order beyond top-3<br>= user never sees it"]
    end

sequenceDiagram
    participant Query as User Query
    participant BiEnc as Bi-Encoder<br>(Titan V2 + Adapter)
    participant OS as OpenSearch<br>(ANN Index)
    participant CrossEnc as Cross-Encoder<br>(MiniLM Reranker)
    participant User as User Sees<br>Top 3 Results

    Query->>BiEnc: "dark isekai manga"
    BiEnc->>OS: Query embedding (1024-dim)
    OS->>CrossEnc: Top 10 candidates with scores

    Note over CrossEnc: Score each of 10 candidates:<br>[CLS] query [SEP] doc_i [SEP]<br>10 forward passes × 5ms each = 50ms

    loop For each candidate (10x)
        CrossEnc->>CrossEnc: Concatenate query + doc<br>→ 6 attention layers<br>→ [CLS] → Linear → score
    end

    CrossEnc->>CrossEnc: Sort by cross-encoder score<br>Reorder: [7,2,5,...] → [2,5,7,...]

    CrossEnc->>User: Top 3 reranked results<br>1. Overlord (0.94)<br>2. Berserk (0.91)<br>3. Goblin Slayer (0.87)

graph TD
    subgraph "Source 1: Click-Through Logs (2000 pairs)"
        A["User searched X<br>Clicked product Y"] --> D["Positive pair:<br>(X, Y, relevance=3)"]
        B["User searched X<br>Saw product Z, didn't click"] --> E["Hard negative:<br>(X, Z, relevance=0)"]
    end

    subgraph "Source 2: Editorial (600 pairs)"
        F["Manga experts rate<br>query-product relevance<br>0=irrelevant, 1=marginal,<br>2=relevant, 3=perfect"] --> G["Graded relevance pairs"]
    end

    subgraph "Source 3: Synthetic (400 pairs)"
        H["Claude generates<br>'user searching for X<br>should/shouldn't find Y'"] --> I["Synthetic pairs<br>with quality filtering"]
    end

    D --> J["Combined Dataset<br>3000 query-doc pairs<br>with graded relevance"]
    E --> J
    G --> J
    I --> J

    J --> K["80/10/10 Split<br>Stratified by<br>relevance level"]

import json
from dataclasses import dataclass
from typing import List, Tuple
from sklearn.model_selection import train_test_split


@dataclass
class RankingExample:
    query: str
    document: str
    relevance: int  # 0=irrelevant, 1=marginal, 2=relevant, 3=perfect


def load_ranking_dataset(
    clickthrough_path: str,
    editorial_path: str,
    synthetic_path: str,
) -> List[RankingExample]:
    """
    Combine three sources into a unified ranking dataset.
    """
    examples = []

    # Click-through: binary relevance (clicked=3, not-clicked=0)
    with open(clickthrough_path) as f:
        for item in json.load(f):
            examples.append(RankingExample(
                query=item["query"],
                document=item["product_description"],
                relevance=3 if item["clicked"] else 0,
            ))

    # Editorial: graded relevance (0-3)
    with open(editorial_path) as f:
        for item in json.load(f):
            examples.append(RankingExample(
                query=item["query"],
                document=item["product_description"],
                relevance=item["relevance"],
            ))

    # Synthetic: binary (relevant=2, not-relevant=0)
    with open(synthetic_path) as f:
        for item in json.load(f):
            examples.append(RankingExample(
                query=item["query"],
                document=item["product_description"],
                relevance=2 if item["relevant"] else 0,
            ))

    return examples


def create_pairwise_samples(
    examples: List[RankingExample],
) -> List[Tuple[str, str, str]]:
    """
    Create (query, doc_better, doc_worse) triplets for pairwise training.
    Group by query, then form all pairs where relevance differs.
    """
    from collections import defaultdict

    by_query = defaultdict(list)
    for ex in examples:
        by_query[ex.query].append(ex)

    triplets = []
    for query, docs in by_query.items():
        for i in range(len(docs)):
            for j in range(len(docs)):
                if docs[i].relevance > docs[j].relevance:
                    triplets.append((
                        query,
                        docs[i].document,  # better
                        docs[j].document,  # worse
                    ))

    return triplets

import torch
import torch.nn as nn
import numpy as np


class LambdaRankLoss(nn.Module):
    """
    Pairwise ranking loss with NDCG-aware gradient weighting.

    L = -log(σ(s_better - s_worse)) × |ΔNDCG|

    The |ΔNDCG| factor amplifies gradients for pairs whose swap
    would most affect the user-visible ranking (top positions).
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        scores_better: torch.Tensor,  # (batch,) — scores for better docs
        scores_worse: torch.Tensor,    # (batch,) — scores for worse docs
        ndcg_deltas: torch.Tensor,     # (batch,) — |ΔNDCG| for each pair
    ) -> torch.Tensor:
        # Pairwise logistic loss
        diff = scores_better - scores_worse
        pairwise_loss = -torch.log(torch.sigmoid(diff) + 1e-8)

        # Weight by NDCG delta
        weighted_loss = pairwise_loss * ndcg_deltas

        return weighted_loss.mean()


def compute_ndcg_delta(
    relevances: List[int],
    pos_i: int,
    pos_j: int,
    k: int = 3,
) -> float:
    """
    Compute |ΔNDCG@K| if items at positions i and j were swapped.
    """
    # Current DCG@K
    dcg = sum(
        (2 ** relevances[p] - 1) / np.log2(p + 2)
        for p in range(min(k, len(relevances)))
    )

    # Swapped DCG@K
    swapped = list(relevances)
    swapped[pos_i], swapped[pos_j] = swapped[pos_j], swapped[pos_i]
    dcg_swapped = sum(
        (2 ** swapped[p] - 1) / np.log2(p + 2)
        for p in range(min(k, len(swapped)))
    )

    # Ideal DCG@K
    ideal = sorted(relevances, reverse=True)
    idcg = sum(
        (2 ** ideal[p] - 1) / np.log2(p + 2)
        for p in range(min(k, len(ideal)))
    )

    if idcg == 0:
        return 0.0

    return abs((dcg - dcg_swapped) / idcg)

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from torch.optim import AdamW
from transformers import get_linear_schedule_with_warmup
from torch.utils.data import DataLoader, Dataset


class PairwiseRankingDataset(Dataset):
    """Dataset of (query, better_doc, worse_doc, ndcg_delta) tuples."""

    def __init__(self, triplets, tokenizer, max_length=128):
        self.triplets = triplets
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        query, better_doc, worse_doc = self.triplets[idx]

        # Tokenize query-better_doc pair
        better_encoding = self.tokenizer(
            query, better_doc,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )

        # Tokenize query-worse_doc pair
        worse_encoding = self.tokenizer(
            query, worse_doc,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )

        return {
            "better_input_ids": better_encoding["input_ids"].squeeze(0),
            "better_attention_mask": better_encoding["attention_mask"].squeeze(0),
            "worse_input_ids": worse_encoding["input_ids"].squeeze(0),
            "worse_attention_mask": worse_encoding["attention_mask"].squeeze(0),
        }


def train_reranker(
    train_triplets: list,
    val_triplets: list,
    num_epochs: int = 5,
    batch_size: int = 16,
    base_lr: float = 2e-5,
):
    """
    Fine-tune ms-marco-MiniLM cross-encoder for manga reranking.

    Uses pairwise ranking loss: for each (query, better_doc, worse_doc),
    the model should score better_doc higher than worse_doc.
    """
    model_name = "cross-encoder/ms-marco-MiniLM-L-6-v2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name, num_labels=1
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # Optimizer with discriminative learning rates
    optimizer_params = []

    # Relevance head — highest LR
    optimizer_params.append({
        "params": model.classifier.parameters(),
        "lr": base_lr,
    })

    # Encoder layers — discriminative LR
    num_layers = len(model.bert.encoder.layer)
    for layer_idx in range(num_layers - 1, -1, -1):
        layer_lr = base_lr * (0.8 ** (num_layers - 1 - layer_idx))
        optimizer_params.append({
            "params": model.bert.encoder.layer[layer_idx].parameters(),
            "lr": layer_lr,
        })

    # Embeddings — lowest LR
    optimizer_params.append({
        "params": model.bert.embeddings.parameters(),
        "lr": base_lr * (0.8 ** num_layers),
    })

    optimizer = AdamW(optimizer_params, weight_decay=0.01)

    # Training
    train_dataset = PairwiseRankingDataset(train_triplets, tokenizer)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    total_steps = len(train_loader) * num_epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=int(0.1 * total_steps),
        num_training_steps=total_steps,
    )

    loss_fn = nn.MarginRankingLoss(margin=1.0)

    best_ndcg = 0.0

    for epoch in range(num_epochs):
        model.train()
        epoch_loss = 0.0

        for batch in train_loader:
            # Score better documents
            better_scores = model(
                input_ids=batch["better_input_ids"].to(device),
                attention_mask=batch["better_attention_mask"].to(device),
            ).logits.squeeze(-1)

            # Score worse documents
            worse_scores = model(
                input_ids=batch["worse_input_ids"].to(device),
                attention_mask=batch["worse_attention_mask"].to(device),
            ).logits.squeeze(-1)

            # Target: better > worse (all ones)
            target = torch.ones_like(better_scores)

            loss = loss_fn(better_scores, worse_scores, target)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()

            epoch_loss += loss.item()

        # Validation
        ndcg = evaluate_ndcg(model, tokenizer, val_triplets, device)
        avg_loss = epoch_loss / len(train_loader)
        print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, NDCG@3={ndcg:.4f}")

        if ndcg > best_ndcg:
            best_ndcg = ndcg
            model.save_pretrained("./best_reranker")
            tokenizer.save_pretrained("./best_reranker")

    return model


def evaluate_ndcg(model, tokenizer, val_triplets, device, k=3):
    """Compute NDCG@K on validation set."""
    model.eval()
    from collections import defaultdict

    # Group by query
    query_scores = defaultdict(list)
    for query, better_doc, worse_doc in val_triplets:
        for doc, rel in [(better_doc, 1), (worse_doc, 0)]:
            encoding = tokenizer(
                query, doc,
                max_length=128,
                padding="max_length",
                truncation=True,
                return_tensors="pt",
            )
            with torch.no_grad():
                score = model(
                    input_ids=encoding["input_ids"].to(device),
                    attention_mask=encoding["attention_mask"].to(device),
                ).logits.squeeze().item()
            query_scores[query].append((score, rel))

    # Compute NDCG per query
    ndcg_scores = []
    for query, scored_docs in query_scores.items():
        scored_docs.sort(key=lambda x: x[0], reverse=True)
        rels = [rel for _, rel in scored_docs]

        dcg = sum(
            (2 ** rels[i] - 1) / np.log2(i + 2)
            for i in range(min(k, len(rels)))
        )
        ideal_rels = sorted(rels, reverse=True)
        idcg = sum(
            (2 ** ideal_rels[i] - 1) / np.log2(i + 2)
            for i in range(min(k, len(ideal_rels)))
        )
        if idcg > 0:
            ndcg_scores.append(dcg / idcg)

    return np.mean(ndcg_scores) if ndcg_scores else 0.0

import sagemaker
from sagemaker.huggingface import HuggingFaceModel


def deploy_reranker():
    """
    Deploy fine-tuned reranker to SageMaker real-time endpoint.
    ml.c5.xlarge: 4 vCPU, 8GB RAM — CPU inference is sufficient
    for the small model (33M params) at 10 candidates per query.
    """
    huggingface_model = HuggingFaceModel(
        model_data="s3://manga-ml-models/reranker/model.tar.gz",
        role=sagemaker.get_execution_role(),
        transformers_version="4.36",
        pytorch_version="2.1",
        py_version="py310",
        env={
            "SAGEMAKER_MODEL_SERVER_WORKERS": "2",
        },
    )

    predictor = huggingface_model.deploy(
        initial_instance_count=2,
        instance_type="ml.c5.xlarge",
        endpoint_name="manga-reranker-v1",
    )

    return predictor


def invoke_reranker(predictor, query: str, candidates: list) -> list:
    """
    Score all candidates and return sorted by relevance.

    Input: query + list of candidate product descriptions
    Output: candidates sorted by cross-encoder relevance score
    """
    payload = {
        "inputs": [
            {"text": query, "text_pair": doc}
            for doc in candidates
        ],
    }

    scores = predictor.predict(payload)

    # Sort candidates by score (descending)
    scored = list(zip(candidates, scores))
    scored.sort(key=lambda x: x[1], reverse=True)

    return scored

graph LR
    subgraph "Training Pipeline (Monthly)"
        A[Click-through<br>Logs] --> B[Pairwise Triplet<br>Construction]
        B --> C[Cross-Encoder<br>Fine-Tuning<br>(SageMaker)]
        C --> D[Model Artifact<br>S3]
    end

    subgraph "Validation Gate"
        D --> E[NDCG@3 > 0.80<br>on golden set]
        E -->|pass| F[SageMaker Endpoint<br>Blue-Green Deploy]
        E -->|fail| G[Reject + Alert]
    end

    subgraph "Inference"
        H[Bi-Encoder<br>Top 10] --> I[Cross-Encoder<br>Reranker<br>(2x ml.c5.xlarge)]
        I --> J[Top 3 Sorted<br>Results]
    end