decoding.py

"""
decoding.py — inference-time generation: greedy and beam search.

Why this is a separate module: train.py needs greedy decoding for validation
BLEU, evaluate.py needs both. Keeping them here guarantees the two scripts
generate text identically.

The train/inference mismatch
----------------------------
During training the decoder is fed the GROUND TRUTH prefix (teacher forcing).
At inference there is no ground truth, so it is fed its OWN previous outputs.
One bad token early therefore poisons everything after it — the model has never
been trained on its own mistakes. That asymmetry is exactly why beam search
helps: keeping several candidate prefixes alive gives the model a chance to
recover from a locally-attractive but globally-wrong first choice.

Stopping: both decoders stop per-sequence at <eos>. Running for a fixed number
of steps and truncating later would append garbage after the real sentence and
tank BLEU (which is precision-based — every junk token is counted against you).

Efficiency note: both functions re-run the whole decoder stack over the whole
prefix at every step, so generating T tokens costs O(T^2) decoder passes. A KV
cache would make it O(T) by storing each layer's past keys/values. It is left
out on purpose — it adds state-management complexity that has nothing to do
with understanding encoder-decoder attention, and at this scale it is fast
enough.

============================================================================
THIS FILE'S PLACE IN THE CHAIN
============================================================================
Depends on: model.py only (calls model.encode()/model.decode() directly — see
model.py's docstrings on those two methods for why generation bypasses forward()).

Called from:
  - train.py's main(): once per epoch, via translate_corpus(..., method="greedy"),
    to compute validation BLEU for model-selection / early stopping.
  - evaluate.py's main(): via translate_corpus(..., method=<greedy|beam>), to
    produce the final scored hypotheses reported to the user and appended to
    experiments/results.csv.

Internal call chain:
  translate_corpus(model, dataloader, tokenizer, ...)
    for each batch from the dataloader (built by dataset.py's make_dataloader):
      -> greedy_decode(model, batch["src"], ...)      [if method == "greedy"]
         or, once per sentence in the batch:
      -> beam_search_decode(model, one_sentence, ...)  [if method == "beam"]
    -> tokenizer.decode(ids, skip_special_tokens=True) to turn ids back into text
    => returns (hypotheses, references, sources) as plain strings, ready for
       sacrebleu (in train.py/evaluate.py) to score directly.
"""

import torch
import torch.nn.functional as F


@torch.no_grad()
# ^ disables gradient tracking for this whole function — generation is
#   inference-only, so autograd bookkeeping would just waste memory/time.
def greedy_decode(model, src, src_key_padding_mask, bos_id, eos_id, pad_id,
                  max_new_tokens=128):
    """Batched greedy decoding: always take the argmax token.

    src: (B, S) -> returns a list of B token-id lists, <bos>/<eos> stripped.

    Called from translate_corpus() below, once per batch, when method="greedy".
    """
    model.eval()          # disables dropout — deterministic behavior required for reproducible generation
    device = src.device
    B = src.size(0)

    # The encoder runs ONCE. `memory` is then reused at every decoding step —
    # this is the practical payoff of the encoder-decoder split.
    memory = model.encode(src, src_key_padding_mask)   # model.py: TransformerTranslator.encode()

    # Start every sequence in the batch with just <bos> — shape (B, 1).
    ys = torch.full((B, 1), bos_id, dtype=torch.long, device=device)
    # Tracks, per batch element, whether that sequence has already produced <eos>.
    finished = torch.zeros(B, dtype=torch.bool, device=device)

    for _ in range(max_new_tokens):
        # Re-run the WHOLE decoder stack over `ys` as it stands so far (this
        # is the O(T^2) cost the module docstring's "Efficiency note" warns about).
        logits = model.decode(ys, memory,
                              memory_key_padding_mask=src_key_padding_mask,
                              tgt_key_padding_mask=None)  # no pads in a generated prefix
        next_tok = logits[:, -1].argmax(dim=-1)           # (B,) — last position only
        # ^ logits shape is (B, current_len, vocab); [:, -1] takes only the
        #   prediction for the NEXT token (position current_len), argmax
        #   picks the single highest-scoring vocab id — this is "greedy":
        #   no exploration of alternatives, unlike beam_search_decode below.

        # Once a sequence has emitted <eos> we keep appending <pad>, so the
        # tensor stays rectangular but nothing further is really generated.
        next_tok = torch.where(finished, torch.full_like(next_tok, pad_id), next_tok)
        # ^ torch.where(cond, a, b): elementwise "if cond then a else b" —
        #   for already-finished sequences, force the next token to <pad>
        #   regardless of what argmax picked, so they don't keep "generating"
        #   nonsense after their real answer ended.
        ys = torch.cat([ys, next_tok.unsqueeze(1)], dim=1)   # append this step's tokens: (B, len) -> (B, len+1)
        finished |= next_tok.eq(eos_id)                       # mark newly-finished sequences
        if bool(finished.all()):
            break         # early exit once every sequence in the batch has hit <eos>

    # _strip() (below) removes the leading <bos>, trailing <pad>s, and
    # everything from the first <eos> onward, for each sequence independently.
    return [_strip(seq.tolist(), bos_id, eos_id, pad_id) for seq in ys]


@torch.no_grad()
def beam_search_decode(model, src, src_key_padding_mask, bos_id, eos_id, pad_id,
                       beam_size=5, max_new_tokens=128, length_penalty=0.6):
    """Beam search, one source sentence at a time (src must be (1, S)).

    Kept to batch-size 1 for readability: the beams themselves are batched, so
    it is still a single forward pass per step over `beam_size` rows.

    Length penalty (GNMT): score / ((5 + len) / 6)^alpha. Raw log-probabilities
    are sums of negative numbers, so without normalization the shortest
    hypothesis almost always wins and the model produces truncated translations.

    Called from translate_corpus() below, once PER SENTENCE (a Python list
    comprehension over the batch), when method="beam" — note this is why beam
    search is noticeably slower than greedy_decode's batched version.
    """
    model.eval()
    device = src.device
    assert src.size(0) == 1, "beam_search_decode expects one sentence at a time"

    memory = model.encode(src, src_key_padding_mask)            # (1, S, d) — encode the ONE sentence
    # Expand memory to `beam_size` copies so every beam can be processed as
    # one batched decode() call. .expand() is a view (no data copy) until
    # .contiguous() forces actual memory layout, which nn.MultiheadAttention needs.
    memory = memory.expand(beam_size, -1, -1).contiguous()      # (k, S, d)
    mem_mask = (src_key_padding_mask.expand(beam_size, -1).contiguous()
                if src_key_padding_mask is not None else None)

    # beams: (k, current_len) — k candidate token sequences, all starting at <bos>.
    beams = torch.full((beam_size, 1), bos_id, dtype=torch.long, device=device)
    # beam_scores: cumulative log-probability of each beam so far.
    beam_scores = torch.full((beam_size,), float("-inf"), device=device)
    beam_scores[0] = 0.0     # only beam 0 is live at step 0; the others are
                             # clones of it, and without this they would all
                             # expand to the same tokens and waste the beam.
    finished = []            # list of (normalized_score, token_id_list)

    for step in range(max_new_tokens):
        logits = model.decode(beams, memory, memory_key_padding_mask=mem_mask)
        # log_softmax turns raw logits into log-probabilities — needed
        # because beam scores accumulate ADDITIVELY across steps (summing
        # log-probs == multiplying probabilities, in log space).
        log_probs = F.log_softmax(logits[:, -1], dim=-1)        # (k, V)

        # Total score of each candidate continuation = beam score + token logprob
        # Broadcasting: (k, 1) + (k, V) -> (k, V), every beam's running score
        # added to every possible next-token log-prob for that beam.
        cand = beam_scores.unsqueeze(1) + log_probs             # (k, V)
        flat = cand.view(-1)                                     # flatten to (k*V,) to do ONE global top-k
        top_scores, top_idx = flat.topk(beam_size)
        # ^ picks the best `beam_size` (beam, token) COMBINATIONS across ALL
        #   k beams at once — this is what lets beam search prune to a fixed
        #   width while still considering every beam's every possible next token.

        # Undo the flattening: recover which original beam and which token
        # each of the top_idx entries came from.
        beam_idx = torch.div(top_idx, log_probs.size(-1), rounding_mode="floor")
        # ^ integer division by vocab size V: which beam (row) this came from
        token_idx = top_idx % log_probs.size(-1)
        # ^ remainder: which token (column) within that beam's row

        # Build the new set of k beams: take each surviving beam's PAST
        # tokens (indexed by beam_idx, which may repeat or drop the original
        # 0..k-1 ordering) and append its newly chosen token.
        beams = torch.cat([beams[beam_idx], token_idx.unsqueeze(1)], dim=1)
        beam_scores = top_scores

        # Retire any beam that just produced <eos>, and mark its slot dead so
        # it is not extended again.
        for b in range(beam_size):
            if token_idx[b].item() == eos_id:
                lp = ((5 + beams.size(1)) / 6) ** length_penalty   # GNMT length penalty (see docstring)
                finished.append((beam_scores[b].item() / lp,
                                 _strip(beams[b].tolist(), bos_id, eos_id, pad_id)))
                beam_scores[b] = float("-inf")
                # ^ -inf score means this slot can never win future topk
                #   comparisons, effectively freezing it out without shrinking
                #   the tensor (keeping shapes fixed at (beam_size, ...) is
                #   simpler than dynamically resizing).

        if len(finished) >= beam_size or bool(torch.isinf(beam_scores).all()):
            break   # stop once we have enough finished candidates, or every beam has died

    if not finished:
        # Hit the length limit with nothing finished — take the best live beam.
        best = int(beam_scores.argmax())
        lp = ((5 + beams.size(1)) / 6) ** length_penalty
        finished.append((beam_scores[best].item() / lp,
                         _strip(beams[best].tolist(), bos_id, eos_id, pad_id)))

    finished.sort(key=lambda x: x[0], reverse=True)   # highest length-normalized score first
    return finished[0][1]                              # return only the single best hypothesis's token ids


def _strip(ids, bos_id, eos_id, pad_id):
    """Drop the leading <bos> and everything from the first <eos> onward.

    Called from both greedy_decode() and beam_search_decode() to turn raw
    generated-id sequences (which still contain <bos>/<eos>/<pad>) into the
    "real" translated token ids that tokenizer.decode() should convert to text.
    """
    out = []
    for i, t in enumerate(ids):
        if i == 0 and t == bos_id:
            continue        # skip the leading <bos> specifically at position 0
        if t == eos_id:
            break            # stop entirely — anything after <eos> is not part of the sentence
        if t == pad_id:
            continue        # defensive: skip any stray <pad> (shouldn't normally appear before <eos>)
        out.append(t)
    return out


@torch.no_grad()
def translate_corpus(model, dataloader, tokenizer, bos_id, eos_id, pad_id,
                     device, method="greedy", beam_size=5, max_new_tokens=128,
                     length_penalty=0.6, progress_every=0):
    """Decode a whole DataLoader.

    Returns (hypotheses, references, sources) as lists of DETOKENIZED strings —
    sacrebleu scores plain text, never subword ids. tokenizer.decode with
    skip_special_tokens=True plus the byte-level decoder gives back normal
    spacing automatically.

    This is the ONE function train.py and evaluate.py both call — see this
    file's module docstring for the exact call sites. `dataloader` is
    whatever dataset.py's make_dataloader() produced; `tokenizer` is the
    same `tokenizers.Tokenizer` object dataset.py's load_tokenizer() returned.
    """
    hyps, refs, srcs = [], [], []
    for i, batch in enumerate(dataloader):
        # batch is one of dataset.py's collate() output dicts — see
        # dataset.py's make_collate_fn() for exactly what keys exist.
        src = batch["src"].to(device)
        src_mask = batch["src_key_padding_mask"].to(device)

        if method == "greedy":
            out_ids = greedy_decode(model, src, src_mask, bos_id, eos_id, pad_id,
                                    max_new_tokens=max_new_tokens)
            # ^ returns a LIST of token-id lists, one per sentence in the batch, already
        elif method == "beam":
            # No batched beam-search here: loop over each sentence in this
            # batch individually (src[j:j+1] keeps the batch dim as size 1,
            # which is what beam_search_decode's assert requires).
            out_ids = [
                beam_search_decode(model, src[j:j + 1], src_mask[j:j + 1],
                                   bos_id, eos_id, pad_id, beam_size=beam_size,
                                   max_new_tokens=max_new_tokens,
                                   length_penalty=length_penalty)
                for j in range(src.size(0))
            ]
        else:
            raise ValueError(f"unknown decoding method: {method}")

        # tokenizer.decode(): the exact inverse of the encode_batch() call in
        # dataset.py's TranslationDataset.__init__ — turns integer ids back
        # into normal, correctly-spaced text (see tokenizer.py's ByteLevel
        # decoder comment for why this "just works").
        hyps.extend(tokenizer.decode(ids, skip_special_tokens=True).strip()
                    for ids in out_ids)
        refs.extend(batch["tgt_text"])    # ground-truth English strings, straight from dataset.py
        srcs.extend(batch["src_text"])    # original Russian strings, for printing/inspection

        if progress_every and (i + 1) % progress_every == 0:
            print(f"    decoded {len(hyps)} sentences...", flush=True)

    return hyps, refs, srcs