Evaluating machine translation with chrF and COMET

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import io
import json
import math
import os
import re
import time
import unicodedata
import urllib.request
import zipfile
from collections import Counter

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import load_file
from tfs import decoding
from tfs.model import Transformer
from tfs.viz import heatmap
from tokenizers import Tokenizer
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# sacreBLEU, the reference implementation of BLEU and chrF
if importlib.util.find_spec("sacrebleu") is None:
    pip = [sys.executable, "-m", "pip", "install", "-q"]
    subprocess.check_call(pip + ["sacrebleu==2.6.0"])
import sacrebleu

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you evaluate in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson runs large models: switch to a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

path = hf_hub_download("xmarva/tfs-tatoeba-en-ru", "data/test.parquet", repo_type="dataset")
test = pd.read_parquet(path)

HERO = "Are you coming alone?"
if SMOKE:  # 150 English sentences with all their translations, and the one from the picture
    chosen = set(test.en.drop_duplicates().sample(150, random_state=SEED)) | {HERO}
    test = test[test.en.isin(chosen)].reset_index(drop=True)
print("test pairs:", len(test), "distinct English sentences:", test.en.nunique())

card = hf_hub_download("Helsinki-NLP/opus-mt-en-ru", "README.md")
rows = []
with open(card, encoding="utf-8") as f:
    for line in f:
        cells = [cell.strip() for cell in line.strip().strip("|").split("|")]
        if len(cells) == 3 and cells[1].replace(".", "").isdigit():  # a row of scores
            scores_here = {"BLEU": float(cells[1]), "chrF": float(cells[2])}
            rows.append({"test set": cells[0], **scores_here})
table(pd.DataFrame(rows), digits=3)

LEADERBOARD = (
    "https://raw.githubusercontent.com/Helsinki-NLP/OPUS-MT-leaderboard/"
    "dcc11795b3653d36865af674ffa784808d721a12/models/OPUS-MT-models/en-ru/"
)
scores_file = urllib.request.urlopen(LEADERBOARD + "opus-2020-02-11.scores.txt")
rows = []
for line in scores_file.read().decode().splitlines():
    pair, test_set, chrf, bleu, release, sentences, words = line.split("\t")
    if test_set.startswith("tatoeba"):
        scores_here = {"BLEU": float(bleu), "chrF": float(chrf)}
        rows.append({"test set": test_set, "sentences": int(sentences), **scores_here})
published = pd.DataFrame(rows)
table(published, digits=3)

QUOTES = {"«": '"', "»": '"', "„": '"', "“": '"', "”": '"', "‘": "'", "’": "'", "–": "—"}


def normalize(text):
    text = unicodedata.normalize("NFKC", text)
    kept = [char for char in text if unicodedata.category(char) != "Mn"]  # no stress marks
    text = "".join(kept)
    for fancy, plain in QUOTES.items():
        text = text.replace(fancy, plain)
    text = re.sub(r"\s+-\s+", " — ", text)  # " - " is a dash
    text = re.sub(r"\s+", " ", text)
    return text.strip()

TESTSETS = (
    "https://raw.githubusercontent.com/Helsinki-NLP/OPUS-MT-testsets/"
    "ac0f19c0be3cb4e0d9450fcb2d7a9847c4bdd180/testsets/eng-rus/"
)
v2021_en = urllib.request.urlopen(TESTSETS + "tatoeba-test-v2021-08-07.eng").read().decode()
v2021_ru = urllib.request.urlopen(TESTSETS + "tatoeba-test-v2021-08-07.rus").read().decode()
v2021 = pd.DataFrame({"en": v2021_en.splitlines(), "ru": v2021_ru.splitlines()})

our_pairs = set(zip(test.en, test.ru))
found = 0
for en, ru in zip(v2021.en, v2021.ru):
    if (normalize(en), normalize(ru)) in our_pairs:
        found += 1
print("v2021-08-07 pairs:", len(v2021), "of them in our test set:", found)

opus_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ru")
opus = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ru").to(device).eval()
opus_cache = {}


@torch.no_grad()
def translate_opus(sentences, batch_size=64):
    todo = [s for s in dict.fromkeys(sentences) if s not in opus_cache]
    todo.sort(key=len)  # similar lengths in one batch: little padding
    for start in range(0, len(todo), batch_size):
        chunk = todo[start : start + batch_size]
        batch = opus_tokenizer(chunk, return_tensors="pt", padding=True).to(device)
        output = opus.generate(**batch)  # its generation_config: 4 beams
        texts = opus_tokenizer.batch_decode(output, skip_special_tokens=True)
        for sentence, text in zip(chunk, texts):
            opus_cache[sentence] = text
    return [opus_cache[s] for s in sentences]

if SMOKE:
    v2021 = v2021.head(200)

start = time.perf_counter()
v2021_hyps = translate_opus(list(v2021.en))
print("translated", len(v2021), "sentences in", round(time.perf_counter() - start), "s")

bleu = sacrebleu.BLEU()
our_bleu = bleu.corpus_score(v2021_hyps, [list(v2021.ru)])
our_chrf = sacrebleu.corpus_chrf(v2021_hyps, [list(v2021.ru)])
row = published[published["test set"] == "tatoeba-test-v2021-08-07"].iloc[0]
print("leaderboard:  BLEU", row.BLEU, " chrF", row.chrF)
our_chrf_share = round(our_chrf.score / 100, 4)  # the leaderboard writes chrF from 0 to 1
print("we measure:   BLEU", round(our_bleu.score, 1), " chrF", our_chrf_share)
print("signature:   ", bleu.get_signature())

tokenizer_path = hf_hub_download("xmarva/tfs-bpe-en-ru-16k", "tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_path)
MODEL_REPO = "xmarva/tfs-transformer-en-ru"
with open(hf_hub_download(MODEL_REPO, "config.json")) as f:
    CONFIG = json.load(f)
model = Transformer(**CONFIG)
model.load_state_dict(load_file(hf_hub_download(MODEL_REPO, "model.safetensors")))
model = model.to(device).eval()
print(sum(p.numel() for p in model.parameters()), "parameters")

PROBES = [
    "The researcher said she was exhausted.",
    "The teacher said she was tired.",
    "My neighbour said he was busy.",
    "Tom forgot to lock the door.",
    "I can't believe you did that.",
    HERO,
]
alone = []
for sentence in PROBES:
    alone.append(decoding.beam_search(model, tokenizer, [sentence])[0])
together = decoding.beam_search(model, tokenizer, PROBES)
assert together == alone
table(pd.DataFrame({"English": PROBES, "ours, beam of 4": together}))

sentences = list(dict.fromkeys(test.en))

start = time.perf_counter()
greedy = decoding.greedy(model, tokenizer, sentences)
print("ours, greedy:", round(time.perf_counter() - start), "s")

start = time.perf_counter()
beam = decoding.beam_search(model, tokenizer, sentences)
print("ours, beam of 4:", round(time.perf_counter() - start), "s")

start = time.perf_counter()
opus_raw = translate_opus(sentences)
print("opus-mt-en-ru:", round(time.perf_counter() - start), "s")
print(len(sentences), "sentences each")

opus_normalized = [normalize(text) for text in opus_raw]
changed = []
for raw, clean in zip(opus_raw, opus_normalized):
    if raw != clean:
        changed.append({"opus-mt-en-ru writes": raw, "normalized": clean})
print("translations changed by the normalization:", len(changed))
table(pd.DataFrame(changed).head(4))

SYSTEMS = ["ours, greedy", "ours, beam", "opus-mt-en-ru"]
outputs = {
    "ours, greedy": dict(zip(sentences, greedy)),
    "ours, beam": dict(zip(sentences, beam)),
    "opus-mt-en-ru": dict(zip(sentences, opus_normalized)),
}
for system in SYSTEMS:
    test[system] = test.en.map(outputs[system])  # a translation for every test pair
table(test.sample(6, random_state=SEED)[["en", "ru", *SYSTEMS]])

from sacrebleu.tokenizers.tokenizer_13a import Tokenizer13a

tokenize_13a = Tokenizer13a()
print(tokenize_13a("Том забыл, чтобы я забыл дверь.").split())

def ngram_counts(items, n):
    counts = Counter()
    for i in range(len(items) - n + 1):
        counts[tuple(items[i : i + n])] += 1
    return counts


def bleu_stats(hypothesis, references, max_n=4):
    words = tokenize_13a(hypothesis).split()
    reference_words = [tokenize_13a(r).split() for r in references]
    matched, total = [], []
    for n in range(1, max_n + 1):
        most = Counter()  # the largest count of each n-gram in any one reference
        for ref in reference_words:
            for gram, count in ngram_counts(ref, n).items():
                most[gram] = max(most[gram], count)
        clipped = 0
        for gram, count in ngram_counts(words, n).items():
            clipped += min(count, most[gram])
        matched.append(clipped)
        total.append(max(len(words) - n + 1, 0))
    lengths = [len(ref) for ref in reference_words]
    # the closest reference length; on a tie, the shorter one
    closest = min(lengths, key=lambda length: (abs(length - len(words)), length))
    return matched + total + [len(words), closest]

def bleu_from_stats(stats, max_n=4):
    stats = np.asarray(stats, dtype=float)
    if stats.ndim == 2:
        stats = stats.sum(axis=0)  # one row per sentence: sum over the test set
    matched, total = stats[:max_n], stats[max_n : 2 * max_n]
    c, r = stats[2 * max_n], stats[2 * max_n + 1]
    if matched.min() == 0:
        return 0.0
    log_precision = np.log(matched / total).mean()
    brevity = 1.0 if c > r else math.exp(1 - r / c)
    return 100 * brevity * math.exp(log_precision)

tiny = [
    ("Том забыл, чтобы я забыл дверь.", "Том забыл запереть дверь."),
    ("Ты еще слишком молода.", "Ты ещё слишком молода."),
    ("Скажи Тому, что я, возможно, опоздаю.", "Скажи Тому, что я, может быть, опоздаю."),
]
tiny_stats = []
for hypothesis, reference in tiny:
    stats = bleu_stats(hypothesis, [reference])
    tiny_stats.append(stats)
    print(hypothesis, "| matched", stats[:4], "of", stats[4:8], "| lengths", stats[8:])

ours = bleu_from_stats(tiny_stats)
theirs = sacrebleu.corpus_bleu([h for h, _ in tiny], [[r for _, r in tiny]]).score
assert abs(ours - theirs) < 1e-9
print("BLEU of the three:", round(ours, 2), "= sacreBLEU's", round(theirs, 2))

bleu_rows = {}
scores = []
for system in SYSTEMS:
    stats = []
    for hypothesis, reference in zip(test[system], test.ru):
        stats.append(bleu_stats(hypothesis, [reference]))
    bleu_rows[system] = np.array(stats)
    ours = bleu_from_stats(bleu_rows[system])
    theirs = sacrebleu.corpus_bleu(list(test[system]), [list(test.ru)]).score
    assert abs(ours - theirs) < 1e-6
    scores.append({"system": system, "BLEU": ours})
table(pd.DataFrame(scores), digits=1)

ratios = np.linspace(0.9, 1.06, 200)  # around 1, where the three systems are
penalty = np.where(ratios > 1, 1.0, np.exp(1 - 1 / ratios))

fig, ax = plt.subplots(figsize=(6.4, 3))
ax.plot(ratios, penalty, color=COLORS["muted"])
for system, color in zip(SYSTEMS, [COLORS["enc"], COLORS["dec"], COLORS["data"]]):
    totals = bleu_rows[system].sum(axis=0)
    ratio = totals[8] / totals[9]  # translation words / reference words
    brevity = 1.0 if ratio > 1 else math.exp(1 - 1 / ratio)
    ax.plot(ratio, brevity, "o", color=color, label=f"{system}: c/r = {ratio:.3f}")
ax.set_xlabel("c / r, total length of the translations over the references")
ax.set_ylabel("brevity penalty")
ax.set_title("The brevity penalty")
ax.legend(fontsize=9)
plt.show()

glued = None
for text in opus_raw:
    if "«" in text and len(text) < 60:
        glued = text
        break
if glued is None:  # the smoke slice may have no such sentence
    glued = "Он сказал: «Я не знаю…»"

tokenize_intl = sacrebleu.BLEU(tokenize="intl").tokenizer
rows = [tokenize_13a(glued).split(), tokenize_intl(glued).split()]
token_chips(rows, labels=["13a", "intl"], mark_spaces=False)
plt.show()

rows = []
for tokenize in ["13a", "intl", "char", "none"]:
    for lowercase in [False, True]:
        metric = sacrebleu.BLEU(tokenize=tokenize, lowercase=lowercase)
        row = {"tokenizer": tokenize, "lowercase": lowercase}
        for system in SYSTEMS:
            row[system] = metric.corpus_score(list(test[system]), [list(test.ru)]).score
        rows.append(row)
table(pd.DataFrame(rows), digits=1)

metric = sacrebleu.BLEU()
metric.corpus_score(list(test["ours, beam"]), [list(test.ru)])  # a signature needs a score
print(metric.get_signature())

raw_by_pair = test.en.map(dict(zip(sentences, opus_raw)))
raw_bleu = sacrebleu.corpus_bleu(list(raw_by_pair), [list(test.ru)]).score
print("opus-mt-en-ru, raw:       ", round(raw_bleu, 2))
print("opus-mt-en-ru, normalized:", round(bleu_from_stats(bleu_rows["opus-mt-en-ru"]), 2))

def has_yo(text):
    return "ё" in text or "Ё" in text


def without_yo(text):
    return text.replace("ё", "е").replace("Ё", "Е")


print(f"references with ё: {test.ru.map(has_yo).mean():.1%}")
rows = []
for system in SYSTEMS:
    plain = [without_yo(text) for text in test[system]]
    references = [without_yo(text) for text in test.ru]
    rows.append(
        {
            "system": system,
            "writes ё": f"{test[system].map(has_yo).mean():.1%}",
            "BLEU": bleu_from_stats(bleu_rows[system]),
            "BLEU, ё as е": sacrebleu.corpus_bleu(plain, [references]).score,
        }
    )
table(pd.DataFrame(rows), digits=3)

def char_ngram_counts(text, max_n):
    chars = "".join(text.split())  # spaces removed
    return [ngram_counts(chars, n) for n in range(1, max_n + 1)]


def chrf_from_stats(stats, max_n=6, beta=2):
    stats = np.asarray(stats, dtype=float)
    if stats.ndim == 2:
        stats = stats.sum(axis=0)  # one row per sentence: sum over the test set
    precisions, recalls = [], []
    for n in range(max_n):
        h, r, m = stats[3 * n : 3 * n + 3]
        if h > 0 and r > 0:  # orders without n-grams are left out
            precisions.append(m / h)
            recalls.append(m / r)
    if not precisions:
        return 0.0
    p, r = np.mean(precisions), np.mean(recalls)
    if p + r == 0:
        return 0.0
    return 100 * (1 + beta**2) * p * r / (beta**2 * p + r)

def chrf_stats(hypothesis, references, max_n=6):
    hyp_counts = char_ngram_counts(hypothesis, max_n)
    best, best_score = None, -1.0
    for reference in references:
        ref_counts = char_ngram_counts(reference, max_n)
        stats = []
        for n in range(max_n):
            shared = 0
            for gram, count in hyp_counts[n].items():
                shared += min(count, ref_counts[n][gram])
            ref_total = sum(ref_counts[n].values())
            # an order the reference is too short to have isn't held against the translation
            hyp_total = sum(hyp_counts[n].values()) if ref_total > 0 else 0
            stats += [hyp_total, ref_total, shared]
        score = chrf_from_stats(stats, max_n)
        if score > best_score:  # keep the reference that suits this translation best
            best, best_score = stats, score
    return best

pairs = [
    ("сказал", "сказала"),
    ("Исследователь сказал, что она устала.", "Исследовательница сказала, что она устала."),
]
for hypothesis, reference in pairs:
    ours = chrf_from_stats(chrf_stats(hypothesis, [reference]))
    theirs = sacrebleu.sentence_chrf(hypothesis, [reference]).score
    assert abs(ours - theirs) < 1e-9
    bleu = sacrebleu.sentence_bleu(hypothesis, [reference]).score
    print(f"{hypothesis} | chrF {ours:.1f}, sentence BLEU {bleu:.1f}")

chrf_rows = {}
for row in scores:
    system = row["system"]
    stats = []
    for hypothesis, reference in zip(test[system], test.ru):
        stats.append(chrf_stats(hypothesis, [reference]))
    chrf_rows[system] = np.array(stats)
    row["chrF"] = chrf_from_stats(chrf_rows[system])
    theirs = sacrebleu.corpus_chrf(list(test[system]), [list(test.ru)]).score
    assert abs(row["chrF"] - theirs) < 1e-6
    # chrF++: words and word pairs added to the characters (below)
    chrf_plus = sacrebleu.CHRF(word_order=2)
    row["chrF++"] = chrf_plus.corpus_score(list(test[system]), [list(test.ru)]).score
table(pd.DataFrame(scores), digits=1)

# all the references of each English sentence
references = test.groupby("en", sort=False).ru.apply(list)
counts = references.apply(len).value_counts().sort_index()

fig, ax = plt.subplots(figsize=(6.4, 2.8))
ax.bar(counts.index, counts.values, color=COLORS["data"])
ax.set_yscale("log")  # the many sentences with one reference next to the few with fifteen
ax.set_xlabel("references per English sentence")
ax.set_ylabel("sentences")
ax.set_title("Most sentences have one reference; some have many")
plt.show()

several = (references.apply(len) > 1).sum()
print("sentences with several references:", several, "of", len(references))

most = references.apply(len).idxmax()  # the first sentence with the most references
print(most)
rows = []
for k in range(0, len(references[most]), 3):
    rows.append(references[most][k : k + 3])
token_chips(rows, role="dec", mark_spaces=False)
plt.show()

streams = []
for k in range(references.apply(len).max()):
    stream = []
    for refs in references:
        stream.append(refs[k] if k < len(refs) else None)
    streams.append(stream)

rows = []
for system in SYSTEMS:
    hypotheses = [outputs[system][sentence] for sentence in references.index]
    multi_bleu = [bleu_stats(h, refs) for h, refs in zip(hypotheses, references)]
    multi_chrf = [chrf_stats(h, refs) for h, refs in zip(hypotheses, references)]
    bleu_all = bleu_from_stats(multi_bleu)
    chrf_all = chrf_from_stats(multi_chrf)
    assert abs(bleu_all - sacrebleu.corpus_bleu(hypotheses, streams).score) < 1e-6
    assert abs(chrf_all - sacrebleu.corpus_chrf(hypotheses, streams).score) < 1e-6
    rows.append(
        {
            "system": system,
            "BLEU, first reference": sacrebleu.corpus_bleu(hypotheses, streams[:1]).score,
            "BLEU, all references": bleu_all,
            "chrF, first reference": sacrebleu.corpus_chrf(hypotheses, streams[:1]).score,
            "chrF, all references": chrf_all,
        }
    )
table(pd.DataFrame(rows), digits=1)

COMET_REPO = "Unbabel/wmt22-comet-da"
with open(hf_hub_download(COMET_REPO, "hparams.yaml")) as f:
    print(f.read())

checkpoint_path = hf_hub_download(COMET_REPO, "checkpoints/model.ckpt")
weights = torch.load(checkpoint_path, map_location="cpu", weights_only=True)["state_dict"]

rows = []
for name, tensor in weights.items():
    if not name.startswith("encoder.") and "scalar_parameters" not in name:
        rows.append({"tensor": name, "shape": str(tuple(tensor.shape))})
layer_scores = []
for name in weights:
    if "scalar_parameters" in name:
        layer_scores.append(weights[name])
shape = f"{len(layer_scores)} × (1,)"
rows.append({"tensor": "layerwise_attention.scalar_parameters.*", "shape": shape})
table(pd.DataFrame(rows))

from transformers import AutoConfig, XLMRobertaModel

XLMR = "FacebookAI/xlm-roberta-large"
xlmr_tokenizer = AutoTokenizer.from_pretrained(XLMR)
encoder = XLMRobertaModel(AutoConfig.from_pretrained(XLMR), add_pooling_layer=False)

encoder_weights = {}
for name, tensor in weights.items():
    if name.startswith("encoder.model."):
        encoder_weights[name.removeprefix("encoder.model.")] = tensor
missing, unexpected = encoder.load_state_dict(encoder_weights, strict=False)
print("missing:", missing, " unexpected:", unexpected)
encoder = encoder.to(device).eval()
print(sum(p.numel() for p in encoder.parameters()), "parameters in the encoder")

def sparsemax(z):
    z_sorted, _ = z.sort(descending=True)
    k = torch.arange(1, len(z) + 1, dtype=z.dtype)
    cumulative = z_sorted.cumsum(dim=0)
    survivors = (1 + k * z_sorted > cumulative).sum()  # how many weights stay above zero
    tau = (cumulative[survivors - 1] - 1) / survivors
    return (z - tau).clamp(min=0)

def sparsemax_by_search(z, steps=100):
    low, high = z.min().item() - 1, z.max().item()
    for _ in range(steps):
        tau = (low + high) / 2
        if (z - tau).clamp(min=0).sum() > 1:  # threshold too low: the sum is over 1
            low = tau
        else:
            high = tau
    return (z - tau).clamp(min=0)


for z in [torch.tensor([2.0, 1.0, -1.0]), torch.tensor([0.5, 0.3, 0.1])]:
    assert torch.allclose(sparsemax(z), sparsemax_by_search(z), atol=1e-6)
    shown_sparsemax = [round(x, 3) for x in sparsemax(z).tolist()]
    shown_softmax = [round(x, 3) for x in z.softmax(dim=0).tolist()]
    shown_z = [round(x, 1) for x in z.tolist()]
    print(shown_z, "→ sparsemax", shown_sparsemax, " softmax", shown_softmax)

layer_scores = torch.cat(layer_scores)
layer_weights = sparsemax(layer_scores)
assert torch.allclose(layer_weights, sparsemax_by_search(layer_scores), atol=1e-5)
gamma = weights["layerwise_attention.gamma"]

fig, ax = plt.subplots(figsize=(6.4, 2.8))
ax.bar(range(len(layer_weights)), layer_weights, color=COLORS["enc"])
for layer, weight in enumerate(layer_weights.tolist()):
    if weight == 0:
        ax.text(layer, 0.002, "0", ha="center", fontsize=8, color=COLORS["dec"])
ax.set_xlabel("XLM-R output (0: the embeddings, 24: the last layer)")
ax.set_ylabel("weight")
ax.set_title("How much of each layer COMET uses")
plt.show()
print("layers with weight exactly 0:", (layer_weights == 0).nonzero().flatten().tolist())
print("the mix is multiplied by gamma =", round(gamma.item(), 3))

@torch.no_grad()
def embed(texts):
    batch = xlmr_tokenizer(texts, return_tensors="pt", padding=True, truncation=True,
                           max_length=510).to(device)
    layers = encoder(**batch, output_hidden_states=True).hidden_states  # 25 × (B, T, 1024)
    mixed = 0
    for weight, layer in zip(layer_weights.tolist(), layers):
        mixed = mixed + weight * layer
    mixed = gamma.item() * mixed
    mask = batch["attention_mask"].unsqueeze(-1).float()
    return (mixed * mask).sum(dim=1) / mask.sum(dim=1)  # the mean over real tokens


estimator = torch.nn.Sequential(
    torch.nn.Linear(6144, 3072),
    torch.nn.Tanh(),
    torch.nn.Linear(3072, 1024),
    torch.nn.Tanh(),
    torch.nn.Linear(1024, 1),
)
for ours, theirs in [(0, 0), (2, 3), (4, 6)]:
    estimator[ours].weight.data = weights[f"estimator.ff.{theirs}.weight"]
    estimator[ours].bias.data = weights[f"estimator.ff.{theirs}.bias"]
estimator = estimator.to(device).eval()
del weights, encoder_weights  # 2.3 GB of tensors we no longer need

embeddings = {}


def embed_all(texts, batch_size=64):
    todo = [text for text in dict.fromkeys(texts) if text not in embeddings]
    todo.sort(key=len)  # similar lengths in one batch: little padding
    for start in range(0, len(todo), batch_size):
        chunk = todo[start : start + batch_size]
        for text, vector in zip(chunk, embed(chunk).cpu()):
            embeddings[text] = vector


@torch.no_grad()
def comet(sources, translations, references, batch_size=512):
    embed_all(sources + translations + references)
    scores = []
    for start in range(0, len(translations), batch_size):
        end = start + batch_size
        s = torch.stack([embeddings[t] for t in sources[start:end]]).to(device)
        h = torch.stack([embeddings[t] for t in translations[start:end]]).to(device)
        r = torch.stack([embeddings[t] for t in references[start:end]]).to(device)
        features = torch.cat([h, r, h * r, (h - r).abs(), h * s, (h - s).abs()], dim=1)
        scores += estimator(features).squeeze(1).tolist()
    return scores

PACKAGE_SCORES = [  # (source, translation, reference, score from unbabel-comet 2.2.7)
    ("The researcher said she was exhausted.", "Исследователь сказал, что она устала.",
     "Исследовательница сказала, что вымоталась.", 0.873125),
    ("The researcher said she was exhausted.", "Исследовательница сказала, что устала.",
     "Исследовательница сказала, что вымоталась.", 0.906672),
    ("Tom forgot to lock the door.", "Том забыл, чтобы я забыл дверь.",
     "Том забыл запереть дверь.", 0.459986),
    ("Tom forgot to lock the door.", "", "Том забыл запереть дверь.", 0.347070),
    ("Tom forgot to lock the door.", "Tom forgot to lock the door.",
     "Том забыл запереть дверь.", 0.614086),
]
fixture_sources = [row[0] for row in PACKAGE_SCORES]
fixture_translations = [row[1] for row in PACKAGE_SCORES]
fixture_references = [row[2] for row in PACKAGE_SCORES]
ours = comet(fixture_sources, fixture_translations, fixture_references)
for row, mine in zip(PACKAGE_SCORES, ours):
    theirs = row[3]
    assert abs(mine - theirs) < 1e-4
    print(f"{mine:.4f} = {theirs:.4f}  {row[1]!r}")

start = time.perf_counter()
for row in scores:
    system = row["system"]
    test[f"COMET {system}"] = comet(list(test.en), list(test[system]), list(test.ru))
    row["COMET"] = test[f"COMET {system}"].mean()  # the system's score: the mean over pairs
print("embedded", len(embeddings), "texts in", round(time.perf_counter() - start), "s")
table(pd.DataFrame(scores), digits=3)

hero_refs = references[HERO]
hero_translation = outputs["opus-mt-en-ru"][HERO]
hero_comet = comet([HERO] * len(hero_refs), [hero_translation] * len(hero_refs), hero_refs)
rows = []
for reference, comet_score in zip(hero_refs, hero_comet):
    rows.append(
        {
            "reference": reference,
            "BLEU": sacrebleu.sentence_bleu(hero_translation, [reference]).score,
            "chrF": sacrebleu.sentence_chrf(hero_translation, [reference]).score,
            "COMET": comet_score,
        }
    )
print("opus-mt-en-ru:", hero_translation)
table(pd.DataFrame(rows).sort_values("COMET", ascending=False), digits=2)

MQM_REPO = "RicardoRei/wmt-mqm-human-evaluation"
mqm_path = hf_hub_download(MQM_REPO, "train.csv", repo_type="dataset")
mqm = pd.read_csv(mqm_path)
mqm = mqm[(mqm.lp == "en-ru") & (mqm.year == 2022)].reset_index(drop=True)
if SMOKE:
    kept = set(mqm.src.drop_duplicates().head(20))
    mqm = mqm[mqm.src.isin(kept)].reset_index(drop=True)
print("segments:", len(mqm), "systems:", mqm.system.nunique())
table(mqm[["src", "mt", "score", "system"]].sample(4, random_state=SEED))

start = time.perf_counter()
mqm["COMET"] = comet(list(mqm.src), list(mqm.mt), list(mqm.ref))
print("COMET on", len(mqm), "segments in", round(time.perf_counter() - start), "s")

rows = []
for system, part in mqm.groupby("system"):
    rows.append(
        {
            "system": system,
            "MQM": part.score.mean(),
            "BLEU": sacrebleu.corpus_bleu(list(part.mt), [list(part.ref)]).score,
            "chrF": sacrebleu.corpus_chrf(list(part.mt), [list(part.ref)]).score,
            "COMET": part.COMET.mean(),
        }
    )
judged = pd.DataFrame(rows).sort_values("MQM", ascending=False).reset_index(drop=True)
for metric in ["BLEU", "chrF", "COMET"]:
    judged[f"{metric} rank"] = judged[metric].rank(ascending=False).astype(int)
table(judged, digits=2)

LABELLED = ["Online-W", "bleu_bestmbr", "comet_bestmbr", "M2M100_1.2B-B4"]

fig, axes = plt.subplots(1, 3, figsize=(11, 3.4))
for ax, metric in zip(axes, ["BLEU", "chrF", "COMET"]):
    ax.scatter(judged.MQM, judged[metric], color=COLORS["muted"], s=24)
    for _, row in judged[judged.system.isin(LABELLED)].iterrows():
        ax.scatter(row.MQM, row[metric], color=COLORS["dec"], s=30)
        point = (row.MQM, row[metric])
        ax.annotate(row.system, point, xytext=(4, 4), textcoords="offset points",
                    fontsize=8)
    ax.set_xlabel("MQM, translators (0 = no errors)")
    ax.set_title(metric)
plt.tight_layout()
plt.show()

rows = []
for metric in ["BLEU", "chrF", "COMET"]:
    agree, pairs = 0, 0
    for i in range(len(judged)):
        for j in range(i + 1, len(judged)):
            pairs += 1
            human = judged.MQM[i] - judged.MQM[j]
            machine = judged[metric][i] - judged[metric][j]
            if human * machine > 0:  # same sign: same order
                agree += 1
    rows.append({"metric": metric, "pairs in the translators' order": f"{agree} of {pairs}",
                 "accuracy": agree / pairs})
table(pd.DataFrame(rows), digits=3)

def breakages(reference, english):
    broken = {}
    if re.search(r"(?i)(?<![а-яё])не\s", reference):  # «не» as a word: drop it
        without = re.sub(r"(?i)(?<![а-яё])не\s", "", reference, count=1)
        broken["«не» dropped"] = without[:1].upper() + without[1:]
    if "Том" in reference:  # Tom is a different person now, in any case form
        broken["Tom becomes John"] = reference.replace("Том", "Джон")
    words = reference.rstrip(".!?").split()
    if len(words) >= 4:
        ending = reference[len(reference.rstrip(".!?")) :]
        broken["last word dropped"] = " ".join(words[:-1]) + ending
    broken["left in English"] = english
    if has_yo(reference):  # not an error at all: the control
        broken["ё written as е"] = without_yo(reference)
    return broken


print(breakages("Том не знает, где Мэри живёт.", "Tom doesn't know where Mary lives."))

rows = []
for english, refs in references.items():
    distinct = list(dict.fromkeys(refs))
    if len(distinct) < 2:
        continue
    first, other = distinct[0], distinct[1]
    for kind, broken in breakages(first, english).items():
        row = {"en": english, "reference": first, "other": other, "kind": kind}
        rows.append({**row, "broken": broken})
stress = pd.DataFrame(rows)
print(stress.kind.value_counts().to_string())

stress["COMET other"] = comet(list(stress.en), list(stress.other), list(stress.reference))
stress["COMET broken"] = comet(list(stress.en), list(stress.broken), list(stress.reference))
for metric, score in [("BLEU", sacrebleu.sentence_bleu), ("chrF", sacrebleu.sentence_chrf)]:
    other, broken = [], []
    for row in stress.itertuples():
        other.append(score(row.other, [row.reference]).score)
        broken.append(score(row.broken, [row.reference]).score)
    stress[f"{metric} other"] = other
    stress[f"{metric} broken"] = broken

KINDS = ["«не» dropped", "Tom becomes John", "last word dropped", "left in English",
         "ё written as е"]  # the last one is no error at all
METRICS = ["BLEU", "chrF", "COMET"]
preferred = np.zeros((len(KINDS), len(METRICS)))
for i, kind in enumerate(KINDS):
    part = stress[stress.kind == kind]
    for j, metric in enumerate(METRICS):
        preferred[i, j] = (part[f"{metric} other"] > part[f"{metric} broken"]).mean()

heatmap(preferred, rows=KINDS, cols=METRICS, annotate=True, fmt="{:.0%}", vmax=1, cell=0.6,
        title="How often the metric prefers the other correct translation")
plt.show()