code/src/lrs_alignment.py

"""Parse LRS2/LRS3 per-word alignment text files and count word frequencies.

Format (pretrain splits of both datasets — this is what makes word-level
clipping possible at all; the alignment is provided, not computed by us):

    Text:  ABOUT TWO THIRDS OF THE COST
    Conf:  1

    WORD START END ASDSCORE
    ABOUT 0.00 0.34 1.00000
    TWO 0.34 0.54 1.00000
    ...

Only files with a WORD/START/END table carry per-word timing. LRS2's "main"
split and LRS3's trainval/test splits typically have just a Text: line and
are skipped — there's no way to isolate one word from a whole-sentence clip
without timing, so those entries simply yield no word rows.
"""
from __future__ import annotations

import re
from collections import Counter
from typing import Iterable, List, Optional, Tuple

from src.lrs_source import LRSSource

_WORD_ROW = re.compile(r"^([A-Za-z']+)\s+([\d.]+)\s+([\d.]+)(?:\s+([\d.]+))?\s*$")

# Pure function words — deliberately minimal. Content words like ABOUT or
# PEOPLE are good lip-reading classes (LRW's own 500-word vocabulary keeps
# them), so we only strip words carrying no visual/lexical content.
DEFAULT_STOPWORDS = {
    "THE", "A", "AN", "TO", "OF", "IS", "IN", "AND", "THAT", "IT", "ARE",
    "WAS", "FOR", "ON", "AS", "WITH", "BE", "AT", "HAVE", "THIS", "FROM",
    "OR", "BY", "BUT", "WE", "CAN", "YOU", "I", "HE", "SHE", "THEY", "THEM",
    "HIS", "HER", "ITS", "OUR", "YOUR", "THEIR", "NOT", "DO", "DOES", "DID",
    "WILL", "WOULD", "SHOULD", "COULD", "HAS", "HAD",
}


def parse_word_timings(text: str) -> Optional[List[Tuple[str, float, float, float]]]:
    """Return [(WORD, start_sec, end_sec, conf)] or None if untimed."""
    lines = text.splitlines()
    header_idx = None
    for i, line in enumerate(lines):
        upper = line.strip().upper()
        if upper.startswith("WORD") and "START" in upper:
            header_idx = i
            break
    if header_idx is None:
        return None
    rows = []
    for line in lines[header_idx + 1:]:
        line = line.strip()
        if not line:
            continue
        m = _WORD_ROW.match(line)
        if not m:
            continue
        word = m.group(1).upper()
        start, end = float(m.group(2)), float(m.group(3))
        conf = float(m.group(4)) if m.group(4) is not None else 1.0
        rows.append((word, start, end, conf))
    return rows or None


def iter_word_rows(sources: Iterable[LRSSource], min_duration: float = 0.0,
                   max_duration: float = 10.0, min_conf: float = 0.0):
    """Yield (source, entry, word, start, end) for every qualifying word
    occurrence across all given sources. Reads only .txt members — cheap
    even across LRS3's 100 shards."""
    for source in sources:
        with source:
            for entry in source.entries():
                text = source.read_text(entry.txt_member)
                rows = parse_word_timings(text)
                if not rows:
                    continue
                for word, start, end, conf in rows:
                    dur = end - start
                    if dur < min_duration or dur > max_duration:
                        continue
                    if conf < min_conf:
                        continue
                    yield source, entry, word, start, end


def count_word_frequencies(sources: Iterable[LRSSource], min_duration: float = 0.0,
                           max_duration: float = 10.0, min_conf: float = 0.0) -> Counter:
    counts: Counter = Counter()
    for _, _, word, _, _ in iter_word_rows(sources, min_duration,
                                          max_duration, min_conf):
        counts[word] += 1
    return counts