Train an Embedding Model: Contrastive Fine-Tuning for Retrieval

"""Train an embedding model so Tavola's staff find the right help article from the way they actually ask.
The articles are a fixed corpus; you tune the encoder on (question, article) pairs with a contrastive objective."""
import numpy as np
import torch

import harness as H

DOCS, SPLITS = H.load()
DOC_TEXT = [d["text"] for d in DOCS]
DOC_ID = [d["id"] for d in DOCS]
DID_IX = {d: i for i, d in enumerate(DOC_ID)}


# ---------- Step 1: measure retrieval ----------

def scores(query_emb, doc_emb):
    """The similarity of every query to every document: query_emb @ doc_emb.T (both are already normalised)."""
    # TODO (Step 1): query_emb @ doc_emb.T.
    raise NotImplementedError("Step 1: write scores()")


def metrics(score_matrix, gold):
    """{"r@1", "r@5", "mrr"} for a (queries x docs) score matrix and gold[i] = the right document's column index.
    A query is a hit at k if its gold document is among the k highest-scoring; MRR averages 1 / rank."""
    # TODO (Step 1): per query, the rank of its gold column; r@1, r@5 and mean 1/rank.
    raise NotImplementedError("Step 1: write metrics()")


def evaluate(model, split):
    """Retrieval metrics of the model on one split: embed every document and every query, score, and score the ranks."""
    D = H.embed(model, DOC_TEXT)
    rows = SPLITS[split]
    Q = H.embed(model, [r["query"] for r in rows])
    return metrics(scores(Q, D).numpy(), [DID_IX[r["doc"]] for r in rows])


# ---------- Step 2: the contrastive loss ----------

def info_nce(q, d, temperature=0.05):
    """The in-batch contrastive loss (InfoNCE) for a batch of pairs: q[i] should be closest to d[i] and far from the
    other d's. Score every q against every d, divide by temperature, and take the cross-entropy that wants each row i
    to pick column i. Returns a scalar loss tensor."""
    raise NotImplementedError("info_nce() arrives in Step 2")


# ---------- Step 3: fine-tune ----------

def train_epoch(model, opt, pairs, epoch=0, temperature=0.05, batch_size=16):
    """One pass over the pairs in shuffled batches (seed the shuffle with epoch). For each batch: embed the queries
    and their gold documents with grad=True, compute info_nce, back-propagate, step. Return the mean loss."""
    raise NotImplementedError("train_epoch() arrives in Step 3")


# ---------- Step 4: hard negatives ----------

def hardest_negatives(model, pairs, k=1):
    """For each pair, the ids of the k documents that score highest but are not the pair's gold document: the ones the
    model is most likely to confuse it with. Returns a list of lists of doc ids."""
    raise NotImplementedError("hardest_negatives() arrives in Step 4")


# ---------- Step 5: keep the best checkpoint ----------

def best_epoch(history):
    """The epoch (an entry's "epoch") with the highest dev "mrr" in history; the earliest on a tie."""
    raise NotImplementedError("best_epoch() arrives in Step 5")