align: anchor each line to the transcript segments it was heard in

src/songtrace/stages/align.py

@@ -4,9 +4,12 @@ from pathlib import Path
 from ..lyrics import alignment_token, normalize_lyrics
 from ..models import on_release
 from ..pipeline import Ctx, PipelineError, Stage
-from .transcribe import classify_lyrics_match, load_transcript, lyrics_similarity, vocal_sections
+from .transcribe import (classify_lyrics_match, line_segments, load_transcript,
+                         lyrics_similarity)
 
 SIMILARITY_THRESHOLD = 0.4
+# how far a chunk's window reaches past its segments, into the gap between them
+PAD_SEC = 1.0
 
 # a stage aligns section by section, so the model is cached across the calls
 # within one run and dropped at the end of it
@@ -50,17 +53,39 @@ def force_align_section(vocals_path: Path, start: float, end: float,
              sum(t.score for t in s) / len(s)) for s in spans]
 
 
-def _assign_lines_to_sections(line_indices: list[int], sections) -> list[list[int]]:
-    """Split matched lyric line indices across vocal sections proportionally to section duration."""
-    total = sum(b - a for a, b in sections) or 1.0
-    n = len(line_indices)
-    groups, elapsed, start = [], 0.0, 0
-    for i, (a, b) in enumerate(sections):
-        elapsed += b - a
-        stop = n if i == len(sections) - 1 else min(n, round(n * elapsed / total))
-        groups.append(line_indices[start:max(stop, start)])
-        start = max(stop, start)
-    return groups
+def _chunks(segments, lines: list[int], anchors: dict[int, tuple[int, int]]
+            ) -> list[tuple[list[int], tuple[float, float]]]:
+    """Group lines heard in the same transcript segments, each group with the
+    window those segments cover.
+
+    Aligning each group only within its own window stops one misheard
+    passage from pushing every later word off its time. A line with no
+    alignable words joins the group before it.
+    """
+    filled, prev = {}, None
+    for li in lines:
+        prev = filled[li] = anchors.get(li, prev)
+    first = next((a for a in filled.values() if a), (0, len(segments) - 1))
+    groups: list[list] = []
+    for li in lines:
+        lo, hi = filled[li] or first
+        if groups and lo <= groups[-1][2]:
+            groups[-1][0].append(li)
+            groups[-1][2] = max(groups[-1][2], hi)
+        else:
+            groups.append([[li], lo, hi])
+    bounds = []
+    for _, lo, hi in groups:
+        times = [t for s in segments[lo:hi + 1] for t in (s.start, s.end)]
+        bounds.append((min(times), max(times)))
+    # pad only into the gap before the neighbouring group's segments
+    out = []
+    for i, ((group, _, _), (a, b)) in enumerate(zip(groups, bounds)):
+        prev_end = bounds[i - 1][1] if i else 0.0
+        next_start = bounds[i + 1][0] if i + 1 < len(bounds) else b + PAD_SEC
+        out.append((group, (max(a - PAD_SEC, min(prev_end, a)),
+                            min(b + PAD_SEC, max(next_start, b)))))
+    return out
 
 
 def _interpolate(rows: list[dict], lo_bound: float, hi_bound: float) -> None:
@@ -112,16 +137,15 @@ def _run(ctx: Ctx) -> None:
         raise PipelineError(
             f"lyrics don't match audio (similarity {sim:.2f}) — "
             "wrong lyrics or wrong recording?")
-    # normalised so an inverted span can never collapse a section onto one instant
-    sections = [(min(a, b), max(a, b)) for a, b in vocal_sections(segments)]
     dropped_lines, missing_spans = classify_lyrics_match(segments, ref)
     matched = [i for i in range(len(ref.lines)) if i not in set(dropped_lines)]
     vocals = ctx.work / "vocals.wav"
     rows = []
-    for sec, line_idxs in zip(sections, _assign_lines_to_sections(matched, sections)):
+    for line_idxs, window in _chunks(segments, matched,
+                                     line_segments(segments, ref)):
         words = [(w, li) for li in line_idxs for w in ref.lines[li]]
         if words:
-            rows.extend(_align_section(vocals, sec, words, ctx.inp.language,
+            rows.extend(_align_section(vocals, window, words, ctx.inp.language,
                                        ctx.cfg.device))
     _enforce_monotonic(rows)
     (ctx.work / "alignment.json").write_text(json.dumps(rows))

src/songtrace/stages/transcribe.py

@@ -6,16 +6,6 @@ from ..pipeline import Ctx, Stage
 from ..whisper_client import TranscriptSegment, transcribe
 
 
-def vocal_sections(segments, min_gap: float = 6.0, pad: float = 1.0):
-    sections = []
-    for s in sorted(segments, key=lambda s: s.start):
-        if sections and s.start - sections[-1][1] < min_gap:
-            sections[-1][1] = max(sections[-1][1], s.end)
-        else:
-            sections.append([s.start, s.end])
-    return [(max(0.0, a - pad), b + pad) for a, b in sections]
-
-
 def _tokens(text: str) -> list[str]:
     return [t for w in text.split() if (t := alignment_token(w))]
 
@@ -39,6 +29,26 @@ def lyrics_similarity(segments, ref: RefLyrics) -> float:
     return matcher.ratio()
 
 
+def line_segments(segments, ref: RefLyrics) -> dict[int, tuple[int, int]]:
+    """For each reference line that was heard, the first and last transcript
+    segment its matched words were heard in.
+
+    Matching is in order, so these ranges never go backwards from one line to
+    the next. Words in a 'replace' block are paired with heard words
+    proportionally, since the block only says the two runs correspond.
+    """
+    matcher, seg_of, line_of = match_lyrics(segments, ref)
+    spans: dict[int, tuple[int, int]] = {}
+    for tag, ai1, ai2, bj1, bj2 in matcher.get_opcodes():
+        if tag not in ("equal", "replace"):
+            continue
+        for j in range(bj1, bj2):
+            si = seg_of[ai1 + (j - bj1) * (ai2 - ai1) // (bj2 - bj1)]
+            lo, hi = spans.get(line_of[j], (si, si))
+            spans[line_of[j]] = (min(lo, si), max(hi, si))
+    return spans
+
+
 def classify_lyrics_match(segments, ref: RefLyrics, pad: float = 1.0
                           ) -> tuple[list[int], list[tuple[float, float]]]:
     """Split reference lines into matched vs. dropped, and flag sung spans no

tests/test_align.py

@@ -123,11 +123,12 @@ def test_non_chronological_transcript(tmp_path, monkeypatch):
 
 
 def test_inverted_section_does_not_collapse(tmp_path, monkeypatch):
-    monkeypatch.setattr(align, "vocal_sections", lambda segments: [(29.0, 5.0)])
-    rows = run_align(tmp_path, monkeypatch)
+    transcript = [{"start": 29.0, "end": 5.0,
+                   "text": "it's one two three four five six seven eight"}]
+    rows = run_align(tmp_path, monkeypatch, transcript=transcript)
     check_invariants(rows, LYRICS)
     assert len({r["start"] for r in rows}) == len(rows)
-    assert rows[0]["start"] >= 5.0 and rows[-1]["end"] <= 29.0
+    assert rows[0]["start"] >= 4.0 and rows[-1]["end"] <= 30.0
 
 
 def test_words_split_across_sections(tmp_path, monkeypatch):
@@ -160,3 +161,47 @@ def test_align_flags_missing_span_for_unmatched_singing(tmp_path, monkeypatch):
     meta = json.loads((tmp_path / "w" / "align_meta.json").read_text())
     assert meta["dropped_lines"] == []
     assert meta["missing_spans"] == [[19.0, 24.0]]
+
+
+def test_contiguous_segments_are_aligned_separately(tmp_path, monkeypatch):
+    # No gap between the segments, so each line must still be held to the
+    # segment it was heard in rather than aligned across both.
+    lyrics = "one two three four\nfive six seven eight"
+    transcript = [{"start": 0.0, "end": 10.0, "text": "one two three four"},
+                  {"start": 10.0, "end": 12.0, "text": "five six seven eight"}]
+    calls = []
+    def spy(vocals_path, start, end, tokens, device):
+        calls.append((start, end, len(tokens)))
+        return fake_align(vocals_path, start, end, tokens, device)
+    monkeypatch.setattr(align, "force_align_section", spy)
+    ctx = make_ctx(tmp_path, lyrics, transcript)
+    align.align_stage().run(ctx)
+    rows = json.loads((ctx.work / "alignment.json").read_text())
+    check_invariants(rows, lyrics)
+    assert calls == [(0.0, 10.0, 4), (10.0, 13.0, 4)]
+    assert all(r["start"] >= 10.0 for r in rows[4:])
+
+
+def test_lines_sharing_a_segment_align_together(tmp_path, monkeypatch):
+    lyrics = "one two\nthree four\nfive six"
+    transcript = [{"start": 2.0, "end": 6.0, "text": "one two three"},
+                  {"start": 6.0, "end": 9.0, "text": "four"},
+                  {"start": 20.0, "end": 23.0, "text": "five six"}]
+    calls = []
+    def spy(vocals_path, start, end, tokens, device):
+        calls.append((start, end, len(tokens)))
+        return fake_align(vocals_path, start, end, tokens, device)
+    monkeypatch.setattr(align, "force_align_section", spy)
+    ctx = make_ctx(tmp_path, lyrics, transcript)
+    align.align_stage().run(ctx)
+    assert calls == [(1.0, 10.0, 4), (19.0, 24.0, 2)]
+
+
+def test_line_without_alignable_words_joins_the_line_before(tmp_path, monkeypatch):
+    lyrics = "one two\n...\nthree four"
+    transcript = [{"start": 1.0, "end": 3.0, "text": "one two"},
+                  {"start": 20.0, "end": 22.0, "text": "three four"}]
+    rows = run_align(tmp_path, monkeypatch, lyrics, transcript)
+    check_invariants(rows, lyrics)
+    assert rows[2]["end"] <= 4.0
+    assert rows[3]["start"] >= 19.0

tests/test_transcribe.py

@@ -1,22 +1,10 @@
 from songtrace.lyrics import normalize_lyrics
 from songtrace.whisper_client import TranscriptSegment
-from songtrace.stages.transcribe import classify_lyrics_match, lyrics_similarity, vocal_sections
+from songtrace.stages.transcribe import classify_lyrics_match, line_segments, lyrics_similarity
 
 def seg(s, e, t=""):
     return TranscriptSegment(start=s, end=e, text=t)
 
-def test_vocal_sections_merges_close_segments():
-    sections = vocal_sections([seg(2, 5), seg(6, 9), seg(30, 33)])
-    assert sections == [(1.0, 10.0), (29.0, 34.0)]
-
-def test_vocal_sections_handles_out_of_order_segments():
-    sections = vocal_sections([seg(30, 33), seg(1, 4)])
-    assert sections == [(0.0, 5.0), (29.0, 34.0)]
-
-def test_vocal_sections_keeps_contained_segment():
-    sections = vocal_sections([seg(1, 10), seg(3, 5)])
-    assert sections == [(0.0, 11.0)]
-
 def test_similarity_high_for_matching_lyrics():
     ref = normalize_lyrics("it's one two three\nfour five six seven eight")
     segs = [seg(0, 2, "It's one two three"), seg(2, 4, "four five six seven eight")]
@@ -51,16 +39,32 @@ def test_classify_lyrics_match_keeps_partially_matched_lines():
     assert dropped == []
 
 def test_classify_lyrics_match_scopes_missing_span_to_adlib_not_whole_section():
-    # matched line and an unmatched ad-lib close enough together that
-    # vocal_sections would merge them into one section; missing_spans must
-    # stay tight around the ad-lib so the review marks the right place.
+    # matched line and an unmatched ad-lib two seconds apart; missing_spans
+    # must stay tight around the ad-lib so the review marks the right place.
     ref = normalize_lyrics("it's one two three\nfour five six seven eight")
     segs = [seg(1, 4, "it's one two three four five six seven eight"),
             seg(6, 8, "completely unrelated ad lib")]
-    sections = vocal_sections(segs)
-    assert sections == [(0.0, 9.0)]      # confirms the two segments do merge
     dropped, missing = classify_lyrics_match(segs, ref)
     assert dropped == []
     assert missing == [(5.0, 9.0)]
     for a, b in missing:
         assert not (a <= 1.0 < b)        # matched line's start (1.0) untouched
+
+
+def test_line_segments_maps_lines_to_the_segments_they_were_heard_in():
+    ref = normalize_lyrics("one two three\nfour five six\nseven eight nine")
+    segs = [seg(1, 4, "one two three four"), seg(5, 8, "five six"),
+            seg(9, 12, "seven eight nine")]
+    assert line_segments(segs, ref) == {0: (0, 0), 1: (0, 1), 2: (2, 2)}
+
+
+def test_line_segments_omits_lines_never_heard():
+    ref = normalize_lyrics("one two three\nfour five six\nseven eight nine")
+    segs = [seg(1, 4, "one two three"), seg(9, 12, "seven eight nine")]
+    assert line_segments(segs, ref) == {0: (0, 0), 2: (1, 1)}
+
+
+def test_line_segments_spreads_a_replace_block_across_its_segments():
+    ref = normalize_lyrics("one two\nthree four")
+    segs = [seg(1, 2, "won too"), seg(3, 4, "tree for")]
+    assert line_segments(segs, ref) == {0: (0, 0), 1: (1, 1)}