BLEU 指標

from __future__ import annotations

import math
from collections import Counter
from collections.abc import Iterable, Sequence


def ngram_counts(tokens: Sequence[str], n: int) -> Counter[tuple[str, ...]]:
    """トークン列から n-gram の頻度を数える。"""
    return Counter(tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1))


def modified_precision(
    candidate: Sequence[str],
    references: Iterable[Sequence[str]],
    n: int,
) -> tuple[int, int]:
    """候補文の n-gram が参照文に出現した件数と候補文の n-gram 総数を返す。"""
    cand_counts = ngram_counts(candidate, n)
    max_ref: Counter[tuple[str, ...]] = Counter()
    for ref in references:
        max_ref |= ngram_counts(ref, n)
    overlap = {ng: min(count, max_ref[ng]) for ng, count in cand_counts.items()}
    return sum(overlap.values()), max(1, sum(cand_counts.values()))


def brevity_penalty(candidate_len: int, reference_lens: Iterable[int]) -> float:
    """候補文が短すぎる場合に課す brevity penalty を計算する。"""
    if candidate_len == 0:
        return 0.0
    closest_ref_len = min(reference_lens, key=lambda r: (abs(r - candidate_len), r))
    ratio = candidate_len / closest_ref_len
    if ratio > 1:
        return 1.0
    return math.exp(1 - 1 / ratio)


def bleu(candidate: str, references: Sequence[str], max_n: int = 4) -> float:
    """候補文と参照文の集合から BLEU スコアを計算する。"""
    candidate_tokens = candidate.split()
    reference_tokens = [ref.split() for ref in references]
    precisions: list[float] = []
    for n_value in range(1, max_n + 1):
        overlap, total = modified_precision(candidate_tokens, reference_tokens, n_value)
        precisions.append(overlap / total)
    if min(precisions) == 0:
        return 0.0
    geometric_mean = math.exp(sum(math.log(p) for p in precisions) / max_n)
    penalty = brevity_penalty(len(candidate_tokens), (len(ref) for ref in reference_tokens))
    return penalty * geometric_mean


if __name__ == "__main__":
    candidate_sentence = "the cat is on the mat"
    reference_sentences = [
        "there is a cat on the mat",
        "the cat sits on the mat",
    ]
    score = bleu(candidate_sentence, reference_sentences)
    print(f"BLEU = {score:.3f}")