transcribe/chunks.py

"""Plan bounded audio uploads and stitch their overlapping word transcripts."""

import argparse
import array
import json
import math
from pathlib import Path
import sys
import wave

from subtitles import convert, normalized_words
from transcript import render_text

CONTEXT = 3.0  # Audio on each side of a seam: six seconds shared by two uploads.
ENCODING_MARGIN = 0.1  # MP3 encoder delay and frame padding at 16 kHz.
PAUSE_SEARCH = 30.0
MIN_SILENCE = 0.3
SILENCE_PEAK = 32768 * 10 ** (-35 / 20)


def pause_before(audio, earliest, latest):
    """Find the last quiet interval in a small PCM window before the limit."""
    rate = audio.getframerate()
    first, last = round(earliest * rate), round(latest * rate)
    audio.setpos(first)
    samples = array.array("h", audio.readframes(last - first))
    if sys.byteorder != "little":
        samples.byteswap()
    frame = round(0.02 * rate)
    quiet_start = None
    boundary = latest
    for offset in range(0, len(samples) + frame, frame):
        block = samples[offset:offset + frame]
        quiet = bool(block) and max(map(abs, block)) <= SILENCE_PEAK
        if quiet and quiet_start is None:
            quiet_start = offset
        elif not quiet and quiet_start is not None:
            quiet_end = min(offset, len(samples))
            if (quiet_end - quiet_start) / rate >= MIN_SILENCE:
                boundary = (first + (quiet_start + quiet_end) / 2) / rate
            quiet_start = None
    return boundary


def plan(path, max_seconds=1800):
    if not math.isfinite(max_seconds) or not 60 <= max_seconds <= 1800:
        raise ValueError("Chunk length must be between 1 and 30 minutes")
    with wave.open(str(path), "rb") as audio:
        if (audio.getnchannels(), audio.getsampwidth(), audio.getframerate()) != (1, 2, 16000):
            raise ValueError("Chunk planning requires mono 16 kHz 16-bit PCM WAV")
        duration = audio.getnframes() / audio.getframerate()
        if not duration:
            raise ValueError("Audio file contains no samples")
        chunks = []
        core_start = 0.0
        while core_start < duration:
            start = max(0.0, core_start - CONTEXT)
            if duration - start <= max_seconds - ENCODING_MARGIN:
                core_end = duration
            else:
                latest = start + max_seconds - ENCODING_MARGIN - CONTEXT
                core_end = pause_before(audio, max(core_start, latest - PAUSE_SEARCH), latest)
            end = min(duration, core_end + CONTEXT)
            chunks.append({"start": start, "end": end, "seam": core_start})
            core_start = core_end
    return chunks


def midpoint(word):
    return (word["start"] + word["end"]) / 2


def token(word):
    return "".join(c for c in word["word"].casefold() if c.isalnum())


def overlap_matches(left, right, seam):
    """Find matching words at matching times in the shared audio."""
    radius = CONTEXT - 0.5
    left_near = [(i, word, token(word)) for i, word in enumerate(left)
                 if abs(midpoint(word) - seam) <= radius]
    right_near = [(i, word, token(word)) for i, word in enumerate(right)
                  if abs(midpoint(word) - seam) <= radius]
    matches = {}
    for i, a, a_token in left_near:
        for j, b, b_token in right_near:
            difference = abs(midpoint(a) - midpoint(b))
            if a_token and a_token == b_token and difference <= 1.0:
                matches[i, j] = (difference, abs((midpoint(a) + midpoint(b)) / 2 - seam))
    return matches, bool(left_near or right_near)


def stitch(left, right, seam):
    """Switch transcripts at a shared word, away from either audio cut.

    Match text AND time to avoid deleting intentional repetitions. Prefer an
    anchor with a matching neighbour; accept a single shared word when speech
    is sparse. Timestamp ownership is the fallback if hypotheses disagree.
    """
    matches, speech_near = overlap_matches(left, right, seam)
    if matches:
        def score(pair):
            i, j = pair
            run = [matches[pair][0]]
            for direction in (-1, 1):
                step = direction
                while (i + step, j + step) in matches:
                    run.append(matches[i + step, j + step][0])
                    step += direction
            # A whole matching phrase disambiguates "yes yes yes" better than
            # a single word that just happens to be closer to the seam.
            return (-len(run), round(sum(run) / len(run), 3), matches[pair][1])

        i, j = min(matches, key=score)
        return left[:i + 1] + right[j + 1:]

    if speech_near:
        print(
            f"Warning: no shared words near {seam:.2f}s; joining by timestamps. "
            "Review the transcript at this boundary.", file=sys.stderr,
        )
    return ([word for word in left if midpoint(word) < seam] +
            [word for word in right if midpoint(word) >= seam])


def speaker_mapping(left, right, seam):
    """Reuse a speaker ID only with unambiguous shared-word evidence.

    Require two distinct word pairs, with no competing speaker correspondence
    in either direction. Numeric provider IDs alone carry no identity evidence.
    """
    matches, _ = overlap_matches(left, right, seam)
    left_words, right_words = {}, {}
    candidates, reverse = {}, {}
    for i, j in matches:
        left_words.setdefault(i, set()).add(j)
        right_words.setdefault(j, set()).add(i)
        old, new = left[i].get("speaker"), right[j].get("speaker")
        if old is not None and new is not None:
            candidates.setdefault(new, set()).add(old)
            reverse.setdefault(old, set()).add(new)
    support = {}
    for i, j in matches:
        old, new = left[i].get("speaker"), right[j].get("speaker")
        if (old is not None and new is not None and
                len(left_words[i]) == len(right_words[j]) == 1):
            support[new, old] = support.get((new, old), 0) + 1
    return {new: old for (new, old), count in support.items()
            if count >= 2 and len(candidates[new]) == len(reverse[old]) == 1}


def merge(chunks, responses, scope_speakers=False):
    """Stitch words, optionally leaving speaker identity to an audio pass."""
    if len(chunks) != len(responses) or not chunks:
        raise ValueError("Expected one transcription response per chunk")
    merged = []
    previous = []
    next_speaker = 0
    for index, (chunk, response) in enumerate(zip(chunks, responses)):
        try:
            words = normalized_words(response, "Chunk stitching")
            for word in words:
                if word["end"] > chunk["end"] - chunk["start"] + 1.0:
                    raise ValueError("Word timestamp exceeds the chunk duration")
                word["start"] += chunk["start"]
                word["end"] += chunk["start"]
            mapping = {} if scope_speakers else speaker_mapping(previous, words, chunk["seam"])
            unmatched = False
            for word in words:
                speaker = word.get("speaker")
                if speaker is not None:
                    if speaker not in mapping:
                        mapping[speaker] = next_speaker
                        next_speaker += 1
                        unmatched = True
                    word["speaker"] = mapping[speaker]
            if index and unmatched and not scope_speakers:
                print(
                    f"Warning: speaker continuity is uncertain near {chunk['seam']:.2f}s; "
                    "unmatched speakers received new transcript IDs.", file=sys.stderr,
                )
        except ValueError as error:
            raise ValueError(f"Chunk {index + 1}: {error}") from error
        merged = stitch(merged, words, chunk["seam"]) if index else words
        previous = words
    # Overlap trimming may discard a speaker's only words. Number only the
    # speakers present in the final transcript, in order of first appearance.
    visible_speakers = {}
    for word in merged:
        speaker = word.get("speaker")
        if speaker is not None:
            word["speaker"] = visible_speakers.setdefault(speaker, len(visible_speakers))
    return {"text": " ".join(word["word"] for word in merged), "words": merged}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    planner = commands.add_parser("plan")
    planner.add_argument("audio")
    planner.add_argument("minutes", type=int)
    merger = commands.add_parser("merge")
    merger.add_argument("manifest", type=Path)
    merger.add_argument("--srt", action="store_true")
    merger.add_argument("--json", action="store_true", help="Emit stitched words for further processing")
    merger.add_argument("--scope-speakers", action="store_true",
                        help="Keep request-local speakers separate for subsequent audio comparison")
    args = parser.parse_args()
    try:
        if args.command == "plan":
            print(json.dumps(plan(args.audio, args.minutes * 60)))
        else:
            chunks = json.loads(args.manifest.read_text(encoding="utf-8"))
            responses = [json.loads(
                (args.manifest.parent / f"response-{i}.json").read_text(encoding="utf-8")
            ) for i in range(len(chunks))]
            if args.json and args.srt or args.scope_speakers and not args.json:
                raise ValueError("Use --scope-speakers with --json, and --json without --srt")
            response = merge(chunks, responses, scope_speakers=args.scope_speakers)
            if args.json:
                print(json.dumps(response))
            else:
                sys.stdout.write(convert(response) if args.srt else render_text(response))
    except (ValueError, OSError, wave.Error) as error:
        parser.exit(1, f"Audio chunking failed: {error}\n")


if __name__ == "__main__":
    main()