Fine-Tune a Small Text Classifier: Baselines, a Training Loop, Early Stopping, Abstention and a Shift Test

"""Train a small text classifier for Deployflow's team chat: label each message by the writer's intent
(blocker, status, request, praise, decision). The words overlap between intents, so the model has to learn the
intent, not the vocabulary."""
import numpy as np
import torch
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

import harness as H

LABELS = H.LABELS


def texts_labels(rows):
    return [r["text"] for r in rows], [r["label"] for r in rows]


def accuracy(pred, gold):
    """The share of predictions that match, as a float."""
    return float(np.mean([p == g for p, g in zip(pred, gold)])) if gold else 0.0


# ---------- Step 1: two baselines to beat ----------

def fit_bow(train):
    """A bag-of-words baseline: a TfidfVectorizer (1- and 2-grams) fitted on the training texts feeding a
    LogisticRegression. Return an object with a .predict(texts) that gives label strings."""
    # TODO (Step 1): TfidfVectorizer(1,2-grams) + LogisticRegression on the training texts; return an object with .predict.
    raise NotImplementedError("Step 1: write fit_bow()")


def fit_frozen(train):
    """A feature-extraction baseline: the frozen encoder's embeddings (H.encode_frozen) feeding a LogisticRegression.
    Return an object with .predict(texts)."""
    # TODO (Step 1): H.encode_frozen(texts) as features + LogisticRegression; return an object with .predict.
    raise NotImplementedError("Step 1: write fit_frozen()")


# ---------- Step 2: fine-tune the encoder ----------

def train_epoch(model, opt, rows, epoch=0, batch_size=16):
    """One pass over rows in shuffled batches (seed the shuffle with epoch so runs repeat). For each batch: zero the
    gradients, run the model on the tokenised texts with their labels, back-propagate the loss, step the optimiser.
    Return the mean loss over the batches."""
    raise NotImplementedError("train_epoch() arrives in Step 2")


def predict(model, texts, batch_size=32):
    """The predicted label string for each text."""
    raise NotImplementedError("predict() arrives in Step 2")


# ---------- Step 3: stop at the right epoch ----------

def best_epoch(history):
    """The epoch (an entry's "epoch") with the highest "dev_acc" in history (a list of {"epoch","train_acc",
    "dev_acc"}); the earliest on a tie. This is the checkpoint to keep, not the last one."""
    raise NotImplementedError("best_epoch() arrives in Step 3")


# ---------- Step 4: when the model is unsure ----------

def probabilities(model, texts, batch_size=32):
    """The softmax probability of each label for each text, as an array of shape (len(texts), len(LABELS))."""
    raise NotImplementedError("probabilities() arrives in Step 4")


def coverage_accuracy(probs, gold, tau):
    """(coverage, accuracy): the share of texts whose top probability is at least tau (the rest go to a person), and
    the accuracy on that kept share (1.0 if none are kept). gold is the list of true label strings."""
    raise NotImplementedError("coverage_accuracy() arrives in Step 4")


# ---------- Step 5: the shift test ----------

def per_class_f1(pred, gold, labels=None):
    """{label: F1} for each label. F1 is 2*precision*recall/(precision+recall), 0 when it is undefined."""
    raise NotImplementedError("per_class_f1() arrives in Step 5")