core/alignment.py

import difflib
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

from core.text_utils import clean_token


@dataclass
class AlignItem:
    ref_word: Optional[str]
    hyp_word: Optional[str]
    status: str
    ref_start: Optional[float]
    ref_end: Optional[float]
    hyp_start: Optional[float]
    hyp_end: Optional[float]
    suggestion: str
    severity: Optional[float] = None
    confidence: Optional[float] = None
    decision: Optional[str] = None
    evidence: Optional[str] = None
    issue_type: Optional[str] = None


def align_words(
    ref_words: List[Dict[str, Any]],
    hyp_words: List[Dict[str, Any]],
    suggest_for_word,
) -> List[AlignItem]:
    ref_items = [(w, clean_token(w["word"])) for w in ref_words]
    hyp_items = [(w, clean_token(w["word"])) for w in hyp_words]
    ref_items = [(w, t) for (w, t) in ref_items if t]
    hyp_items = [(w, t) for (w, t) in hyp_items if t]
    ref_tokens = [t for _, t in ref_items]
    hyp_tokens = [t for _, t in hyp_items]

    def dp_align(ref_slice: List[Tuple[Dict[str, Any], str]], hyp_slice: List[Tuple[Dict[str, Any], str]]):
        if not ref_slice and not hyp_slice:
            return []
        if not ref_slice:
            return [
                AlignItem(
                    ref_word=None,
                    hyp_word=hyp["word"],
                    status="insert",
                    ref_start=None,
                    ref_end=None,
                    hyp_start=hyp["start"],
                    hyp_end=hyp["end"],
                    suggestion=suggest_for_word(hyp["word"], "insert"),
                    issue_type="content",
                )
                for hyp, _ in hyp_slice
            ]
        if not hyp_slice:
            return [
                AlignItem(
                    ref_word=ref["word"],
                    hyp_word=None,
                    status="delete",
                    ref_start=ref["start"],
                    ref_end=ref["end"],
                    hyp_start=None,
                    hyp_end=None,
                    suggestion=suggest_for_word(ref["word"], "delete"),
                    issue_type="content",
                )
                for ref, _ in ref_slice
            ]

        ref_toks = [t for _, t in ref_slice]
        hyp_toks = [t for _, t in hyp_slice]
        n, m = len(ref_toks), len(hyp_toks)
        dp = [[0] * (m + 1) for _ in range(n + 1)]
        back = [[""] * (m + 1) for _ in range(n + 1)]
        for i in range(1, n + 1):
            dp[i][0] = i
            back[i][0] = "del"
        for j in range(1, m + 1):
            dp[0][j] = j
            back[0][j] = "ins"

        for i in range(1, n + 1):
            for j in range(1, m + 1):
                cost = 0 if ref_toks[i - 1] == hyp_toks[j - 1] else 1
                subs = dp[i - 1][j - 1] + cost
                dele = dp[i - 1][j] + 1
                ins = dp[i][j - 1] + 1
                best = min(subs, dele, ins)
                dp[i][j] = best
                if best == subs:
                    back[i][j] = "eq" if cost == 0 else "sub"
                elif best == dele:
                    back[i][j] = "del"
                else:
                    back[i][j] = "ins"

        items: List[AlignItem] = []
        i, j = n, m
        while i > 0 or j > 0:
            op = back[i][j]
            if op in ("eq", "sub"):
                ref = ref_slice[i - 1][0]
                hyp = hyp_slice[j - 1][0]
                status = "equal" if op == "eq" else "replace"
                items.append(
                    AlignItem(
                        ref_word=ref["word"],
                        hyp_word=hyp["word"],
                        status=status,
                        ref_start=ref["start"],
                        ref_end=ref["end"],
                        hyp_start=hyp["start"],
                        hyp_end=hyp["end"],
                        suggestion=suggest_for_word(ref["word"], status),
                        issue_type="content" if status != "equal" else None,
                    )
                )
                i -= 1
                j -= 1
            elif op == "del":
                ref = ref_slice[i - 1][0]
                items.append(
                    AlignItem(
                        ref_word=ref["word"],
                        hyp_word=None,
                        status="delete",
                        ref_start=ref["start"],
                        ref_end=ref["end"],
                        hyp_start=None,
                        hyp_end=None,
                        suggestion=suggest_for_word(ref["word"], "delete"),
                        issue_type="content",
                    )
                )
                i -= 1
            else:
                hyp = hyp_slice[j - 1][0]
                items.append(
                    AlignItem(
                        ref_word=None,
                        hyp_word=hyp["word"],
                        status="insert",
                        ref_start=None,
                        ref_end=None,
                        hyp_start=hyp["start"],
                        hyp_end=hyp["end"],
                        suggestion=suggest_for_word(hyp["word"], "insert"),
                        issue_type="content",
                    )
                )
                j -= 1

        items.reverse()
        return items

    if not ref_tokens and not hyp_tokens:
        return []

    matcher = difflib.SequenceMatcher(a=ref_tokens, b=hyp_tokens, autojunk=False)
    blocks = [b for b in matcher.get_matching_blocks() if b.size > 0]
    if not blocks:
        return dp_align(ref_items, hyp_items)

    aligned: List[AlignItem] = []
    ref_i = 0
    hyp_i = 0
    for block in blocks:
        if ref_i < block.a or hyp_i < block.b:
            aligned.extend(dp_align(ref_items[ref_i:block.a], hyp_items[hyp_i:block.b]))
        for k in range(block.size):
            ref = ref_items[block.a + k][0]
            hyp = hyp_items[block.b + k][0]
            aligned.append(
                AlignItem(
                    ref_word=ref["word"],
                    hyp_word=hyp["word"],
                    status="equal",
                    ref_start=ref["start"],
                    ref_end=ref["end"],
                    hyp_start=hyp["start"],
                    hyp_end=hyp["end"],
                    suggestion="",
                )
            )
        ref_i = block.a + block.size
        hyp_i = block.b + block.size

    if ref_i < len(ref_items) or hyp_i < len(hyp_items):
        aligned.extend(dp_align(ref_items[ref_i:], hyp_items[hyp_i:]))

    return aligned


def compute_wer(items: List[AlignItem]) -> float:
    subs = sum(1 for x in items if x.status == "replace")
    ins = sum(1 for x in items if x.status == "insert")
    dele = sum(1 for x in items if x.status == "delete")
    ref_len = sum(1 for x in items if x.ref_word is not None)
    if ref_len == 0:
        return 1.0
    return (subs + ins + dele) / ref_len