Your notebook is not a product: a Python package with tests and CI

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 ast
import json
import os
import shutil
import time
import urllib.request
import zipfile
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download

if importlib.util.find_spec("sacrebleu") is None:  # Kaggle has it, Colab may not
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "sacrebleu"])
import sacrebleu

SACREBLEU_VERSION = "2.5.1"
PYPI_URL = f"https://pypi.org/pypi/sacrebleu/{SACREBLEU_VERSION}/json"
with urllib.request.urlopen(PYPI_URL) as f:
    release = json.load(f)
for item in release["urls"]:
    if item["filename"].endswith(".whl"):
        wheel_url, wheel_name = item["url"], item["filename"]
wheel_path = WORK_DIR / wheel_name
urllib.request.urlretrieve(wheel_url, wheel_path)

archive = zipfile.ZipFile(wheel_path)
size = wheel_path.stat().st_size / 1000
print(wheel_name, f"{size:.0f} kB,", len(archive.namelist()), "files")
for name in archive.namelist():
    if ".dist-info/" in name:
        print(name)

dist_info = next(name for name in archive.namelist() if ".dist-info/" in name).split("/")[0]
metadata = archive.read(f"{dist_info}/METADATA").decode()
for line in metadata.splitlines():
    if line.startswith(("Name:", "Version:", "Requires-Python:", "Requires-Dist:")):
        print(line)
print()
print(archive.read(f"{dist_info}/entry_points.txt").decode())

def repo_files(repo_id):
    rows = []
    for item in HfApi().model_info(repo_id, files_metadata=True).siblings:
        if not item.rfilename.startswith("."):
            rows.append({"repository": repo_id, "file": item.rfilename,
                         "kB": round(item.size / 1000)})
    return rows


rows = repo_files("Helsinki-NLP/opus-mt-en-ru")
rows += repo_files("xmarva/tfs-transformer-en-ru-tuned")
display(table(pd.DataFrame(rows)))

for repo_id in ["Helsinki-NLP/opus-mt-en-ru", "xmarva/tfs-transformer-en-ru-tuned"]:
    card_text = Path(hf_hub_download(repo_id, "README.md")).read_text(encoding="utf-8")
    print(f"--- {repo_id}: {len(card_text.splitlines())} lines")
    print("\n".join(card_text.splitlines()[:14]))
    print()

REPO = WORK_DIR / "course"
shutil.rmtree(REPO, ignore_errors=True)
for folder in ["perevod/src/perevod", "perevod/tests", ".github/workflows"]:
    (REPO / folder).mkdir(parents=True)
os.chdir(REPO)


def write(path, text):
    Path(path).write_text(text.lstrip("\n"), encoding="utf-8")


print("working in", Path.cwd())

first = subprocess.run([sys.executable, "-c", "import sys; print(repr(sys.path[0]))"],
                       capture_output=True, text=True).stdout.strip()
print("sys.path[0] when Python runs a command:", first)

write("perevod/pyproject.toml", """
[project]
name = "perevod"
version = "0.1.0"
description = "English-to-Russian translation with a small Transformer built from scratch"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
dependencies = [
    "torch>=2.1",
    "tokenizers>=0.15",
    "huggingface_hub>=0.20",
    "safetensors>=0.4",
]

[project.optional-dependencies]
test = ["pytest>=8", "ruff>=0.6", "sacrebleu>=2.4", "pandas>=2.0", "pyarrow>=14"]

[project.scripts]
perevod = "perevod.cli:main"

[project.urls]
Source = "https://github.com/xmarva/transformers-from-scratch/tree/main/perevod"
Model = "https://huggingface.co/xmarva/perevod-en-ru"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.pytest.ini_options]
markers = ["slow: downloads the published model and translates real sentences"]

[tool.ruff]
line-length = 92  # what a notebook cell shows: lesson 7 writes these files in its cells
""")

write("perevod/README.md", """
# perevod

English-to-Russian translation with the small Transformer built in the course
[Transformers: from scratch to production][course]. The model and its card are on the
Hugging Face Hub: [xmarva/perevod-en-ru](https://huggingface.co/xmarva/perevod-en-ru).

```bash
REPO=https://github.com/xmarva/transformers-from-scratch
pip install "perevod @ git+$REPO#subdirectory=perevod"
perevod "I think, therefore I am."
perevod --beam-size 1 < sentences.txt > translations.txt
```

```python
from perevod import Translator

translator = Translator.from_pretrained()
translator(["The doctor will see you now.", "Tom is reading."])
```

The model was trained on short everyday sentences from Tatoeba. Read its model card before
trusting it with anything else.

Tests: `pip install -e ".[test]"`, then `pytest -m "not slow"` (seconds, offline) or
`pytest` (downloads the model).

[course]: https://github.com/xmarva/transformers-from-scratch
""")

%%writefile perevod/src/perevod/model.py
# The encoder-decoder Transformer of lessons 2 and 5, with only the design the published
# model uses: pre-norm with RMSNorm, SwiGLU, RoPE. The parameter names are the lessons', so
# lesson 6's weights load as they are.

import math

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


def rope_angles(max_len, d_k, base=10_000):
    theta = base ** (-torch.arange(0, d_k, 2) / d_k)
    positions = torch.arange(max_len)
    return positions[:, None] * theta[None, :]


def apply_rope(x, angles):
    cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
    even, odd = x[..., 0::2], x[..., 1::2]
    rotated = torch.empty_like(x)
    rotated[..., 0::2] = even * cos - odd * sin
    rotated[..., 1::2] = even * sin + odd * cos
    return rotated

%%writefile -a perevod/src/perevod/model.py


class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))

    def forward(self, x):
        x32 = x.float()
        mean_square = x32.pow(2).mean(dim=-1, keepdim=True)
        normed = x32 * torch.rsqrt(mean_square + self.eps)
        return (normed * self.weight).to(x.dtype)


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        width = 8 * round(2 * d_ff / 3 / 8)  # three matrices with the weights of two
        self.w_gate = nn.Linear(d_model, width, bias=False)
        self.w_up = nn.Linear(d_model, width, bias=False)
        self.w_down = nn.Linear(width, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

%%writefile -a perevod/src/perevod/model.py


class Attention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads, self.d_k = n_heads, d_model // n_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        batch, length, _ = x.shape
        return x.view(batch, length, self.n_heads, self.d_k).transpose(1, 2)

    def forward(self, x_query, x_key_value, mask, angles=None):
        q = self.split_heads(self.w_q(x_query))
        k = self.split_heads(self.w_k(x_key_value))
        v = self.split_heads(self.w_v(x_key_value))
        if angles is not None:  # RoPE in self-attention; cross-attention has no positions
            q, k = apply_rope(q, angles), apply_rope(k, angles)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        batch, _, length, _ = out.shape
        return self.w_o(out.transpose(1, 2).reshape(batch, length, -1))

%%writefile -a perevod/src/perevod/model.py


class Block(nn.Module):
    def __init__(self, config, cross):
        super().__init__()
        self.self_attention = Attention(config.d_model, config.n_heads)
        self.cross_attention = Attention(config.d_model, config.n_heads) if cross else None
        self.feed_forward = SwiGLU(config.d_model, config.d_ff)
        self.norm1, self.norm2 = RMSNorm(config.d_model), RMSNorm(config.d_model)
        if cross:
            self.norm3 = RMSNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x, mask, angles, memory=None, memory_mask=None):
        h = self.norm1(x)
        x = x + self.dropout(self.self_attention(h, h, mask, angles))
        if self.cross_attention is None:
            return x + self.dropout(self.feed_forward(self.norm2(x)))
        h = self.norm2(x)
        x = x + self.dropout(self.cross_attention(h, memory, memory_mask))
        return x + self.dropout(self.feed_forward(self.norm3(x)))

%%writefile -a perevod/src/perevod/model.py


class Seq2Seq(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        d_k = config.d_model // config.n_heads
        self.register_buffer("angles", rope_angles(config.max_len, d_k))
        self.encoder = nn.ModuleList([Block(config, False) for _ in range(config.n_layers)])
        self.decoder = nn.ModuleList([Block(config, True) for _ in range(config.n_layers)])
        self.encoder_norm = RMSNorm(config.d_model)
        self.decoder_norm = RMSNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)

    def embed(self, ids):
        return self.dropout(self.embedding(ids) * math.sqrt(self.config.d_model))

    def encode(self, src):
        mask = (src != self.config.pad_id)[:, None, None, :]
        x = self.embed(src)
        for layer in self.encoder:
            x = layer(x, mask, self.angles[: src.shape[1]])
        return self.encoder_norm(x)

    def decode(self, tgt, memory, src):
        length = tgt.shape[1]
        causal = torch.ones(length, length, dtype=torch.bool, device=tgt.device).tril()
        mask = causal[None, None] & (tgt != self.config.pad_id)[:, None, None, :]
        memory_mask = (src != self.config.pad_id)[:, None, None, :]
        x = self.embed(tgt)
        for layer in self.decoder:
            x = layer(x, mask, self.angles[:length], memory, memory_mask)
        return self.decoder_norm(x) @ self.embedding.weight.T

    def forward(self, src, tgt):
        return self.decode(tgt, self.encode(src), src)

%%writefile perevod/src/perevod/decoding.py
# Greedy decoding and beam search of lesson 3, batched as in lesson 4, moved here from the
# course's notebooks. One change: a source longer than the model's positions is cut to fit,
# instead of failing inside the model.

import torch


def special_ids(tokenizer):
    pad_id = tokenizer.token_to_id("<pad>")
    return pad_id, tokenizer.token_to_id("<bos>"), tokenizer.token_to_id("<eos>")


def source_ids(tokenizer, sentence, limit):
    return tokenizer.encode(sentence).ids[:limit]


def pad(sequences, pad_id, device):
    longest = max(len(ids) for ids in sequences)
    batch = torch.full((len(sequences), longest), pad_id, device=device)
    for row, ids in enumerate(sequences):
        batch[row, : len(ids)] = torch.tensor(ids, device=device)
    return batch


def batches(sources, batch_size):
    # sentences in batches of similar length, so that little of a batch is padding
    order = sorted(range(len(sources)), key=lambda i: len(sources[i]))
    return [order[start : start + batch_size] for start in range(0, len(order), batch_size)]


def length_penalty(length, alpha=0.6):
    # Wu et al. (2016): ((5 + |y|) / 6) ** alpha
    return ((5 + length) / 6) ** alpha

%%writefile -a perevod/src/perevod/decoding.py


@torch.inference_mode()
def greedy(model, tokenizer, sentences, max_tokens=128, batch_size=256):
    model.eval()
    device = next(model.parameters()).device
    pad_id, bos, eos = special_ids(tokenizer)
    sources = []
    for sentence in sentences:
        sources.append(source_ids(tokenizer, sentence, model.config.max_len))
    max_tokens = min(max_tokens, model.config.max_len - 1)
    translations = [None] * len(sentences)
    for indices in batches(sources, batch_size):
        src = pad([sources[i] for i in indices], pad_id, device)
        memory = model.encode(src)
        written = torch.full((len(indices), 1), bos, device=device)
        done = torch.zeros(len(indices), dtype=torch.bool, device=device)
        for _ in range(max_tokens):
            next_ids = model.decode(written, memory, src)[:, -1].argmax(dim=-1)
            next_ids[done] = pad_id  # a finished sentence only collects padding
            written = torch.cat([written, next_ids[:, None]], dim=1)
            done |= next_ids == eos
            if done.all():
                break
        for row, i in enumerate(indices):
            ids = []
            for token_id in written[row, 1:].tolist():
                if token_id in (eos, pad_id):
                    break
                ids.append(token_id)
            translations[i] = tokenizer.decode(ids)
    return translations

%%writefile -a perevod/src/perevod/decoding.py


@torch.inference_mode()
def beam_search(model, tokenizer, sentences, beam_size=4, alpha=0.6, max_tokens=128,
                batch_size=64):
    model.eval()
    device = next(model.parameters()).device
    pad_id, bos, eos = special_ids(tokenizer)
    sources = []
    for sentence in sentences:
        sources.append(source_ids(tokenizer, sentence, model.config.max_len))
    max_tokens = min(max_tokens, model.config.max_len - 1)
    translations = [None] * len(sentences)
    for indices in batches(sources, batch_size):
        src = pad([sources[i] for i in indices], pad_id, device)
        memory = model.encode(src)
        # per sentence: a list of (tokens so far, total log-probability)
        beams = [[([bos], 0.0)] for _ in indices]
        finished = [[] for _ in indices]
        active = [True for _ in indices]

        for _ in range(max_tokens):
            rows, owners = [], []
            for k in range(len(indices)):
                if active[k]:
                    for tokens, _ in beams[k]:
                        rows.append(tokens)
                        owners.append(k)
            if not rows:
                break
            owner = torch.tensor(owners, device=device)
            tgt = torch.tensor(rows, device=device)  # all live beams have the same length
            scores = model.decode(tgt, memory[owner], src[owner])[:, -1]
            top = scores.float().log_softmax(dim=-1).topk(beam_size)
            values, ids = top.values.tolist(), top.indices.tolist()

            row = 0
            for k in range(len(indices)):
                if not active[k]:
                    continue
                candidates = []
                for tokens, total in beams[k]:
                    for log_p, token_id in zip(values[row], ids[row], strict=True):
                        candidates.append((tokens + [token_id], total + log_p))
                    row += 1
                candidates.sort(key=lambda candidate: candidate[1], reverse=True)
                beams[k] = []
                for tokens, total in candidates:
                    if tokens[-1] == eos:
                        finished[k].append((tokens, total))
                    else:
                        beams[k].append((tokens, total))
                    if len(beams[k]) == beam_size:
                        break
                if len(finished[k]) >= beam_size:
                    active[k] = False

        for k, i in enumerate(indices):

            def normalized(candidate):
                tokens, total = candidate
                return total / length_penalty(len(tokens) - 1, alpha)  # without <bos>

            best_tokens, _ = max(finished[k] or beams[k], key=normalized)
            translations[i] = tokenizer.decode(best_tokens[1:])  # decode drops <eos>
    return translations

from safetensors.torch import load_file
from tfs import decoding as notebook_decoding
from tfs.modern import Seq2Seq as NotebookSeq2Seq
from tokenizers import Tokenizer

LESSON6 = "xmarva/tfs-transformer-en-ru-tuned"
with open(hf_hub_download(LESSON6, "config.json")) as f:
    lesson6_config = json.load(f)
weights = load_file(hf_hub_download(LESSON6, "model.safetensors"))
TOKENIZER = hf_hub_download("xmarva/tfs-bpe-en-ru-16k", "tokenizer.json")
tokenizer = Tokenizer.from_file(TOKENIZER)

typo_config = dict(lesson6_config, setup={"norm": "rms", "ffn": "swiglu", "pos": "RoPE"})
typo_model = NotebookSeq2Seq(**typo_config)
print(typo_model.load_state_dict(weights))

notebook_model = NotebookSeq2Seq(**lesson6_config)
notebook_model.load_state_dict(weights)
probe = ["Do you dream in color?", "Mary has two brothers and a sister.",
         "We don't remember days, we remember moments."]
display(table(pd.DataFrame({
    "English": probe,
    '"pos": "rope"': notebook_decoding.beam_search(notebook_model.eval(), tokenizer, probe),
    '"pos": "RoPE"': notebook_decoding.beam_search(typo_model.eval(), tokenizer, probe),
})))

%%writefile perevod/src/perevod/config.py
from dataclasses import asdict, dataclass, fields

DESIGN = {"norm": "rms", "ffn": "swiglu", "pos": "rope"}  # the one design perevod runs


def from_dict(cls, data):
    known = []
    for field in fields(cls):
        known.append(field.name)
    unknown = sorted(set(data) - set(known))
    if unknown:
        raise ValueError(f"{cls.__name__}: unknown keys {unknown}, known are {known}")
    return cls(**data)


@dataclass(frozen=True)
class ModelConfig:
    setup: dict
    vocab_size: int
    d_model: int = 256
    n_heads: int = 4
    d_ff: int = 1024
    n_layers: int = 3
    dropout: float = 0.1
    pad_id: int = 0
    max_len: int = 256

    def __post_init__(self):
        problems = []
        if self.setup != DESIGN:
            problems.append(f"setup is {self.setup}, but perevod runs {DESIGN} only")
        if self.d_model % self.n_heads != 0:
            problems.append(f"{self.n_heads} heads don't split d_model={self.d_model}")
        elif self.d_model // self.n_heads % 2 != 0:
            problems.append("RoPE rotates pairs of numbers: d_model / n_heads must be even")
        if not 0 <= self.dropout < 1:
            problems.append(f"dropout={self.dropout} is not a probability below 1")
        if problems:
            raise ValueError("ModelConfig: " + "; ".join(problems))

    @classmethod
    def from_dict(cls, data):
        return from_dict(cls, data)

    def to_dict(self):
        return asdict(self)

%%writefile -a perevod/src/perevod/config.py
@dataclass(frozen=True)
class GenerationConfig:
    beam_size: int = 4
    alpha: float = 0.6
    max_tokens: int = 128

    def __post_init__(self):
        if self.beam_size < 1:
            raise ValueError(f"GenerationConfig: beam_size={self.beam_size} is below 1")

    @classmethod
    def from_dict(cls, data):
        return from_dict(cls, data)

    def to_dict(self):
        return asdict(self)

import runpy

ModelConfig = runpy.run_path("perevod/src/perevod/config.py")["ModelConfig"]
for bad in [typo_config, dict(lesson6_config, n_head=4), dict(lesson6_config, d_model=250)]:
    try:
        ModelConfig.from_dict(bad)
    except ValueError as error:
        print(error)

translator = Translator.from_pretrained()
translator("I think, therefore I am.")

%%writefile perevod/src/perevod/translator.py
import json
import warnings
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from tokenizers import Tokenizer

from . import decoding
from .config import GenerationConfig, ModelConfig
from .model import Seq2Seq

MODEL = "xmarva/perevod-en-ru"
REVISION = "4abc0d63e862ca99d2bd99cd1007e4a9fe47c358"  # tested with this commit
FILES = ["config.json", "generation_config.json", "model.safetensors", "tokenizer.json"]


class Translator:
    """Translates English into Russian with the course's Transformer.

    >>> translator = Translator.from_pretrained()
    >>> translator("I think, therefore I am.")
    """

    def __init__(self, model, tokenizer, generation):
        self.model = model.eval()
        self.tokenizer = tokenizer
        self.generation = generation

%%writefile -a perevod/src/perevod/translator.py

    @classmethod
    def from_pretrained(cls, model=MODEL, revision=None, device=None):
        """Load a model from a Hugging Face repository or from a local directory.

        The default model comes at the revision this version of perevod was tested with.
        """
        if model == MODEL and revision is None:
            revision = REVISION
        path = Path(model)
        if not path.is_dir():
            path = Path(snapshot_download(model, revision=revision, allow_patterns=FILES))
        config = ModelConfig.from_dict(json.loads((path / "config.json").read_text()))
        settings = json.loads((path / "generation_config.json").read_text())
        network = Seq2Seq(config)
        network.load_state_dict(load_file(path / "model.safetensors"))
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        tokenizer = Tokenizer.from_file(str(path / "tokenizer.json"))
        return cls(network.to(device), tokenizer, GenerationConfig.from_dict(settings))

try:
    notebook_decoding.beam_search(notebook_model, tokenizer, [""])
except RuntimeError as error:
    print("RuntimeError:", str(error)[:90])

pair = ["", "Nobody understands me."]
print(notebook_decoding.beam_search(notebook_model, tokenizer, pair))

%%writefile -a perevod/src/perevod/translator.py

    def translate(self, texts, beam_size=None):
        """Translate a string or a list of strings; beam_size=1 is greedy decoding."""
        single = isinstance(texts, str)
        if single:
            texts = [texts]
        cleaned = []
        for text in texts:
            cleaned.append(self.clean(text))
        # an empty input has an empty translation: the model never sees it
        to_translate = [text for text in cleaned if text]
        done = iter(self.decode(to_translate, beam_size or self.generation.beam_size))
        translations = []
        for text in cleaned:
            translations.append(next(done) if text else "")
        return translations[0] if single else translations

    __call__ = translate

    def clean(self, text):
        text = " ".join(text.split())  # line breaks and runs of spaces become one space
        limit = self.model.config.max_len
        if len(self.tokenizer.encode(text).ids) > limit:
            warnings.warn(f"input longer than {limit} tokens is cut to fit", stacklevel=3)
        return text

    def decode(self, sentences, beam_size):
        settings = self.generation
        if beam_size == 1:
            return decoding.greedy(self.model, self.tokenizer, sentences,
                                   settings.max_tokens)
        return decoding.beam_search(self.model, self.tokenizer, sentences, beam_size,
                                    settings.alpha, settings.max_tokens)

%%writefile perevod/src/perevod/__init__.py
__version__ = "0.1.0"


def __getattr__(name):
    # `import perevod` stays light, so that --help and --version answer at once;
    # PyTorch is loaded only when someone asks for the translator
    if name == "Translator":
        from .translator import Translator

        return Translator
    raise AttributeError(f"module 'perevod' has no attribute {name!r}")

%%writefile perevod/src/perevod/cli.py
import argparse
import sys

from . import __version__


def main(argv=None):
    parser = argparse.ArgumentParser(
        prog="perevod",
        description="Translate English into Russian, one line per argument, or, "
        "without arguments, one line per line of the standard input.",
    )
    parser.add_argument("text", nargs="*", help="sentences to translate")
    parser.add_argument("--beam-size", type=int, help="1 is greedy decoding")
    parser.add_argument("--model", help="a Hugging Face repository or a local directory")
    parser.add_argument("--revision", help="a branch, tag or commit of the repository")
    parser.add_argument("--device", help="cpu or cuda (default: cuda when there is one)")
    parser.add_argument("--version", action="version", version=f"perevod {__version__}")
    args = parser.parse_args(argv)

    texts = args.text
    if not texts:
        texts = [line.rstrip("\n") for line in sys.stdin]

    from .translator import MODEL, Translator  # PyTorch loads only when there is work to do

    translator = Translator.from_pretrained(args.model or MODEL, args.revision, args.device)
    for translation in translator.translate(texts, beam_size=args.beam_size):
        print(translation)
    return 0


if __name__ == "__main__":
    sys.exit(main())

!{sys.executable} -m pip install -q ./perevod
importlib.invalidate_caches()  # this Python started before perevod existed

import perevod
from perevod import Translator

print("perevod", perevod.__version__, "from", Path(perevod.__file__).parent)

translator = Translator.from_pretrained(device="cpu")
translator("I think, therefore I am.")

sentences = ["Do you dream in color?", "Nobody understands me.",
             "I've been learning Russian for three years.", "Where can I buy a dictionary?",
             "Tom is reading a book about the brain."]
from_notebooks = notebook_decoding.beam_search(notebook_model, tokenizer, sentences)
from_package = translator(sentences)
display(table(pd.DataFrame({"English": sentences, "perevod": from_package})))
assert from_package == from_notebooks
print("the package translates exactly like the notebook code")

def seconds(command, repeats=3):
    fastest = None
    for _ in range(repeats):
        start = time.perf_counter()
        subprocess.run(command, capture_output=True, check=True)
        took = time.perf_counter() - start
        if fastest is None or took < fastest:
            fastest = took
    return fastest


PEREVOD = shutil.which("perevod")
LOAD = "from perevod import Translator; Translator.from_pretrained(device='cpu')"
startup = {
    "import perevod": seconds([sys.executable, "-c", "import perevod"]),
    "perevod --version": seconds([PEREVOD, "--version"]),
    "import torch": seconds([sys.executable, "-c", "import torch"]),
    "import torch, load the model": seconds([sys.executable, "-c", LOAD]),
    "perevod 'Nobody understands me.'": seconds([PEREVOD, "Nobody understands me."]),
}
for command, took in startup.items():
    print(f"{took:.2f} s  {command}")

fig, ax = plt.subplots(figsize=(8, 3))
bars = ax.barh(list(startup), list(startup.values()), color=[
    COLORS["data"], COLORS["data"], COLORS["muted"], COLORS["muted"], COLORS["enc"]])
ax.bar_label(bars, fmt="{:.2f} s", padding=4, fontsize=9, color=COLORS["muted"])
ax.invert_yaxis()
ax.set_xlabel("seconds, fastest of three runs")
ax.set_title("Where a command spends its time")
plt.show()

write("sentences.txt", """
Memory is the diary we all carry about with us.

Mary has two brothers and a sister.
""")
!perevod < sentences.txt > translations.txt
!paste -d '|' sentences.txt translations.txt

%%writefile perevod/tests/conftest.py
import json

import pytest
import torch
from safetensors.torch import save_file
from tokenizers import Tokenizer, models, normalizers, pre_tokenizers

from perevod import Translator
from perevod.config import GenerationConfig, ModelConfig
from perevod.model import Seq2Seq

# a model small enough to build in a millisecond: twenty words, random weights. It
# translates nonsense, which is fine: these tests check what must hold for any weights
SPECIAL = ["<pad>", "<unk>", "<bos>", "<eos>"]
WORDS = SPECIAL + ["i", "think", "therefore", "am", ".", ",", "tom", "reads", "a", "book",
                   "?", "я", "думаю", "том", "читает", "книгу"]

%%writefile -a perevod/tests/conftest.py


def tiny_tokenizer():
    vocab = {}
    for token_id, word in enumerate(WORDS):
        vocab[word] = token_id
    tokenizer = Tokenizer(models.WordLevel(vocab, unk_token="<unk>"))
    tokenizer.add_special_tokens(SPECIAL)  # so that decoding leaves them out
    tokenizer.normalizer = normalizers.Lowercase()
    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
    return tokenizer

%%writefile -a perevod/tests/conftest.py


@pytest.fixture(scope="session")
def tiny_dir(tmp_path_factory):
    path = tmp_path_factory.mktemp("tiny-model")
    tiny_tokenizer().save(str(path / "tokenizer.json"))
    config = ModelConfig(setup={"norm": "rms", "ffn": "swiglu", "pos": "rope"},
                         vocab_size=len(WORDS), d_model=16, n_heads=2, d_ff=32, n_layers=2,
                         dropout=0.1, max_len=32)
    torch.manual_seed(0)
    save_file(Seq2Seq(config).state_dict(), path / "model.safetensors")
    (path / "config.json").write_text(json.dumps(config.to_dict()))
    generation = GenerationConfig(beam_size=3, max_tokens=10)
    (path / "generation_config.json").write_text(json.dumps(generation.to_dict()))
    return path


@pytest.fixture
def translator(tiny_dir):
    return Translator.from_pretrained(str(tiny_dir), device="cpu")

%%writefile perevod/tests/test_config.py
import pytest

from perevod.config import GenerationConfig, ModelConfig

LLAMA_STYLE = {"norm": "rms", "ffn": "swiglu", "pos": "rope"}


def test_a_typo_in_a_key_is_an_error():
    with pytest.raises(ValueError, match="n_head"):
        ModelConfig.from_dict({"setup": LLAMA_STYLE, "vocab_size": 100, "n_head": 4})


def test_another_design_is_an_error():
    with pytest.raises(ValueError, match="setup"):
        ModelConfig(setup={**LLAMA_STYLE, "norm": "layer"}, vocab_size=100)


def test_the_heads_must_split_the_width():
    with pytest.raises(ValueError, match="heads"):
        ModelConfig(setup=LLAMA_STYLE, vocab_size=100, d_model=250, n_heads=4)


def test_a_beam_needs_at_least_one_hypothesis():
    with pytest.raises(ValueError, match="beam_size"):
        GenerationConfig(beam_size=0)


def test_a_config_survives_a_round_trip():
    config = ModelConfig(setup=LLAMA_STYLE, vocab_size=100)
    assert ModelConfig.from_dict(config.to_dict()) == config

%%writefile perevod/tests/test_model.py
import torch

from perevod.config import ModelConfig
from perevod.model import Seq2Seq


def tiny_model():
    torch.manual_seed(0)
    config = ModelConfig(setup={"norm": "rms", "ffn": "swiglu", "pos": "rope"},
                         vocab_size=20, d_model=16, n_heads=2, d_ff=32, n_layers=2,
                         max_len=32)
    return Seq2Seq(config).eval()


def test_a_position_does_not_see_the_tokens_after_it():
    model = tiny_model()
    src = torch.tensor([[5, 6, 7, 8]])
    tgt = torch.tensor([[2, 9, 10, 11, 12]])
    changed = tgt.clone()
    changed[0, 3:] = torch.tensor([13, 14])
    before, after = model(src, tgt), model(src, changed)
    assert torch.allclose(before[0, :3], after[0, :3], atol=1e-6)
    assert not torch.allclose(before[0, 3:], after[0, 3:])


def test_padding_the_source_changes_nothing():
    model = tiny_model()
    src = torch.tensor([[5, 6, 7]])
    padded = torch.tensor([[5, 6, 7, 0, 0]])
    tgt = torch.tensor([[2, 9, 10]])
    assert torch.allclose(model(src, tgt), model(padded, tgt), atol=1e-5)

%%writefile perevod/tests/test_decoding.py
from perevod import decoding

SENTENCES = ["i think therefore i am .", "tom reads a book", "i am ?", "tom"]


def test_a_beam_of_one_is_greedy_decoding(translator):
    model, tokenizer = translator.model, translator.tokenizer
    greedy = decoding.greedy(model, tokenizer, SENTENCES, max_tokens=10)
    beam = decoding.beam_search(model, tokenizer, SENTENCES, beam_size=1, max_tokens=10)
    assert beam == greedy


def test_a_batch_translates_like_one_sentence_at_a_time(translator):
    alone = []
    for sentence in SENTENCES:
        alone.append(translator.translate(sentence))
    assert translator.translate(SENTENCES) == alone


def test_the_same_input_gives_the_same_translation(translator):
    assert translator.translate(SENTENCES) == translator.translate(SENTENCES)

%%writefile perevod/tests/test_translator.py
import pytest


def test_an_empty_input_has_an_empty_translation(translator):
    assert translator.translate(["", "   ", "tom reads"])[:2] == ["", ""]


def test_a_string_gives_a_string_and_a_list_gives_a_list(translator):
    assert isinstance(translator.translate("tom reads"), str)
    assert isinstance(translator.translate(["tom reads"]), list)


def test_spaces_and_line_breaks_change_nothing(translator):
    expected = translator.translate("tom reads a book")
    assert translator.translate("tom  reads\na book") == expected


def test_a_long_input_is_cut_with_a_warning(translator):
    long_text = " ".join(["tom reads a book ."] * 10)  # 50 tokens, the model has 32 places
    with pytest.warns(UserWarning, match="cut"):
        translation = translator.translate(long_text)
    assert isinstance(translation, str)

%%writefile perevod/tests/test_cli.py
import io
import subprocess
import sys

from perevod import __version__
from perevod.cli import main


def run_python(*args):
    command = [sys.executable, *args]
    result = subprocess.run(command, capture_output=True, text=True, check=True)
    return result.stdout.strip()


def test_the_version():
    assert run_python("-m", "perevod.cli", "--version") == f"perevod {__version__}"


def test_importing_the_package_does_not_load_pytorch():
    assert run_python("-c", "import sys, perevod; print('torch' in sys.modules)") == "False"


def test_each_argument_gives_one_line(tiny_dir, capsys):
    main(["--model", str(tiny_dir), "--device", "cpu", "tom reads", "i think"])
    assert len(capsys.readouterr().out.splitlines()) == 2


def test_one_line_in_one_line_out(tiny_dir, capsys, monkeypatch):
    monkeypatch.setattr("sys.stdin", io.StringIO("tom reads\n\ni think\n"))
    main(["--model", str(tiny_dir), "--device", "cpu"])
    lines = capsys.readouterr().out.splitlines()
    # an empty line stays empty: line n of the output is the translation of line n
    assert len(lines) == 3
    assert lines[1] == ""

!cd perevod && {sys.executable} -m pytest -q -m "not slow" -p no:cacheprovider --color=no

GOLDEN = ["I think, therefore I am.", "Do you dream in color?",
          "Tom is reading a book about the brain.", "My doctor said she was tired.",
          "The researcher said she was exhausted.",
          "Language is a dialect with an army and navy.",
          "I've been learning Russian for three years.",
          "We don't remember days, we remember moments.",
          "Mary has two brothers and a sister.", "Where can I buy a dictionary?",
          "Memory is the diary we all carry about with us.", "Nobody understands me."]
golden = {}
for english, russian in zip(GOLDEN, translator(GOLDEN), strict=True):
    golden[english] = russian
Path("perevod/tests/golden.json").write_text(
    json.dumps(golden, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
display(table(pd.DataFrame({"English": list(golden), "perevod": list(golden.values())})))

%%writefile perevod/tests/test_quality.py
import json
from pathlib import Path

import pandas as pd
import pytest
import sacrebleu
from huggingface_hub import hf_hub_download

from perevod import Translator

# these tests download the published model: `pytest -m "not slow"` skips them
pytestmark = pytest.mark.slow

GOLDEN = Path(__file__).parent / "golden.json"
DATA = "xmarva/tfs-tatoeba-en-ru"
DATA_REVISION = "540be5b93bf77cd0846b02d55f2558a5596887a9"
CHRF_FLOOR = 57.0  # three points under what the published model scores on this sample


@pytest.fixture(scope="module")
def published():
    return Translator.from_pretrained(device="cpu")


def test_the_published_model_translates_as_it_did(published):
    golden = json.loads(GOLDEN.read_text(encoding="utf-8"))
    assert published.translate(list(golden)) == list(golden.values())


def test_chrf_on_a_sample_of_the_test_set(published):
    path = hf_hub_download(DATA, "data/test.parquet", repo_type="dataset",
                           revision=DATA_REVISION)
    test = pd.read_parquet(path).sample(500, random_state=0)
    chrf = sacrebleu.corpus_chrf(published.translate(list(test.en)), [list(test.ru)]).score
    assert chrf >= CHRF_FLOOR

!cd perevod && {sys.executable} -m pytest -q -p no:cacheprovider --color=no

ENCODER_MASK = "        mask = (src != self.config.pad_id)[:, None, None, :]"
MUTANTS = {
    "the decoder sees the next token": [("model.py", ".tril()", ".tril(1)")],
    "the encoder ignores padding": [
        ("model.py", ENCODER_MASK, "        mask = None")],
    "embeddings are not scaled": [("model.py", " * math.sqrt(self.config.d_model)", "")],
    "positions are forgotten": [("model.py", "if angles is not None:", "if False:")],
    "length penalty inverted": [
        ("decoding.py", "total / length_penalty", "total * length_penalty")],
    "dropout left on": [
        ("translator.py", "self.model = model.eval()", "self.model = model"),
        ("decoding.py", "    model.eval()\n", "")],
    "empty input reaches the model": [
        ("translator.py", "[text for text in cleaned if text]", "cleaned"),
        ("translator.py", 'next(done) if text else ""', "next(done)")],
    "no warning for a long input": [("translator.py", "> limit:", "> 10 * limit:")],
    "translations come back in the wrong order": [
        ("decoding.py", "translations[i] = tokenizer.decode(best",
         "translations[k] = tokenizer.decode(best")],
    "input lines keep their newline": [("cli.py", 'line.rstrip("\\n")', "line")],
}

def run_tests(project):
    environment = dict(os.environ, PYTHONPATH=str(project / "src"))
    command = [sys.executable, "-m", "pytest", "-q", "-rf", "--tb=no", "--color=no",
               "-p", "no:cacheprovider"]
    result = subprocess.run(command, cwd=project, env=environment, capture_output=True,
                            text=True)
    failed = []
    for line in result.stdout.splitlines():
        if line.startswith("FAILED"):
            failed.append(line.split()[1].split("::")[1])
    return failed


def mutate(name, edits):
    project = WORK_DIR / "mutants" / name.replace(" ", "-")
    shutil.rmtree(project, ignore_errors=True)
    shutil.copytree("perevod", project)
    for module, old, new in edits:
        path = project / "src" / "perevod" / module
        source = path.read_text(encoding="utf-8")
        assert old in source, (name, old)  # the bug must really be planted
        path.write_text(source.replace(old, new), encoding="utf-8")
    return project