scripts/compose_cosyvoice_inputs.py

"""Stage 5: compose CosyVoice 3 inputs (3 variants per segment).

Reads:
- data/hindi_emphasis_test/translations/seg_NNN.json   (Stage 4 output)
- data/hindi_emphasis_test/phase1/seg_NNN.json         (Phase 1 prosody data)

Writes 3 variant JSONs per segment to data/hindi_emphasis_test/cosyvoice_inputs/:
- seg_NNN_variant_a.json  -- <strong> only + minimal tone instruct
- seg_NNN_variant_b.json  -- A + [breath] at sentence boundaries
- seg_NNN_variant_c.json  -- <strong> + instruction names emphasized words
                              (reinforcement — our expected winner based on
                               empirical test 'D' in chat)

Each JSON contains tts_text, instruct_text, prompt_wav, ready to feed into
cosyvoice.inference_instruct2() in the Colab notebook.

Instruction length budget: <= 10 words, SINGLE comma-separated clause.
No internal periods, no quotes, no em-dashes. This is empirically what
CosyVoice 3 parses reliably without reading the instruction aloud.
"""

from __future__ import annotations

import json
import re
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
TRANSLATIONS_DIR = ROOT / "data" / "hindi_emphasis_test" / "translations"
PHASE1_DIR = ROOT / "data" / "hindi_emphasis_test" / "phase1"
SEGMENTS_DIR = ROOT / "data" / "hindi_emphasis_test" / "segments"
OUT_DIR = ROOT / "data" / "hindi_emphasis_test" / "cosyvoice_inputs"

# Overall tone adjective (shortest possible — one emotion word carries more
# impact than a verbose "warm cheerful energetic announcer" phrase).
EMOTION_ADJ = "cheerful"
END_OF_PROMPT = "<|endofprompt|>"


# ── Prosody descriptors (single-word, one-shot) ────────────────────────────

def pace_adj(n_words: int, duration: float) -> str:
    """Return a single word: slow | moderate | brisk."""
    if duration <= 0:
        return "moderate"
    rate = n_words / duration
    if rate < 2.0:
        return "slow"
    if rate < 3.5:
        return "moderate"
    return "brisk"


# ── tts_text builders ──────────────────────────────────────────────────────

def _clean_exotic_punctuation(text: str) -> str:
    """Remove em-dash / en-dash / double-dash that CosyVoice 3's text frontend
    handles unreliably (causes stutters and partial instruction leak)."""
    text = text.replace(" — ", ", ").replace("—", ",")
    text = text.replace(" – ", ", ").replace("–", ",")
    text = text.replace(" -- ", ", ").replace("--", ",")
    text = re.sub(r",\s*,", ",", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def build_tts_text_with_strong(english_words: list[str], emphasized_indices: set[int]) -> str:
    """Wrap each emphasized word in <strong>...</strong> (punctuation outside).
    Also strips em-dashes / en-dashes that CosyVoice 3 mis-handles."""
    out = []
    for i, w in enumerate(english_words):
        if i in emphasized_indices:
            trailing = ""
            stripped = w
            while stripped and stripped[-1] in ".,;:!?":
                trailing = stripped[-1] + trailing
                stripped = stripped[:-1]
            out.append(f"<strong>{stripped}</strong>{trailing}")
        else:
            out.append(w)
    return _clean_exotic_punctuation(" ".join(out))


def add_breaths_at_sentence_boundaries(tts_text: str) -> str:
    """Insert [breath] after sentence-ending punctuation. Uses SQUARE brackets
    per the CosyVoice 3 tokenizer (angle brackets are only for wrapper tokens
    like <strong> and system tokens like <|endofprompt|>)."""
    return re.sub(r"([.!?])\s+", r"\1 [breath] ", tts_text)


# ── instruct_text builders ─────────────────────────────────────────────────
# Budget: <= 10 words, single comma-separated clause, no periods, no quotes.

def _scrub_word_for_instruct(w: str) -> str:
    """Strip punctuation and quotes from an emphasized word so it can be
    named safely in the instruct_text (no quotes, no em-dashes, no periods)."""
    w = w.rstrip(".,;:!?-—").strip("'\"")
    # Drop any apostrophes in the middle too (India's -> Indias)
    return w.replace("'", "").replace('"', "")


def variant_a_instruct(phase1: dict) -> str:
    """Minimal: tone + pace. Single clause. ~3-4 words."""
    pace = pace_adj(phase1["n_words"], phase1["duration_seconds"])
    return f"{EMOTION_ADJ} {pace} tone"


def variant_b_instruct(phase1: dict) -> str:
    """Same as A — the tts_text difference (with [breath]) is what B tests."""
    return variant_a_instruct(phase1)


def variant_c_instruct(phase1: dict, emphasized_en_words: list[str]) -> str:
    """Bare emphasis-naming only (no tone prefix, no comma).

    Empirically, 'cheerful brisk tone, emphasize X' leaks in CosyVoice 3 —
    the comma + multi-concept structure confuses the separator parser.
    The bare pattern 'emphasize X Y' (our successful diagnostic 'D') works.
    If no emphasized words, fall back to the A/B tone-only instruct.
    """
    if not emphasized_en_words:
        return variant_a_instruct(phase1)
    clean = [_scrub_word_for_instruct(w) for w in emphasized_en_words]
    seen, unique = set(), []
    for w in clean:
        if w.lower() not in seen:
            seen.add(w.lower())
            unique.append(w)
    if len(unique) > 3:
        unique = unique[:3]
    return f"emphasize {' '.join(unique)}"


# ── Main composition loop ──────────────────────────────────────────────────

def compose_for_segment(seg_id: str) -> list[dict]:
    trans = json.loads((TRANSLATIONS_DIR / f"{seg_id}.json").read_text(encoding="utf-8"))
    phase1 = json.loads((PHASE1_DIR / f"{seg_id}.json").read_text(encoding="utf-8"))

    english_words = trans["english_words"]
    emphasized_indices = set()
    emphasized_en_words = []
    for a in trans["emphasis_alignments"]:
        for idx in a["en_word_indices"]:
            emphasized_indices.add(idx)
            emphasized_en_words.append(english_words[idx])

    # Variant A: <strong> + short tone instruct
    tts_a = build_tts_text_with_strong(english_words, emphasized_indices)
    instruct_a = variant_a_instruct(phase1)

    # Variant B: A + [breath] at sentence boundaries
    tts_b = add_breaths_at_sentence_boundaries(tts_a)
    instruct_b = variant_b_instruct(phase1)

    # Variant C: A + explicit emphasis-word naming (our expected winner)
    tts_c = tts_a
    instruct_c = variant_c_instruct(phase1, emphasized_en_words)

    prompt_wav = str((SEGMENTS_DIR / f"{seg_id}.wav").as_posix())

    variants = []
    for letter, tts, instr, desc in [
        ("a", tts_a, instruct_a, "<strong> only + minimal tone instruct"),
        ("b", tts_b, instruct_b, "A + [breath] at sentence boundaries"),
        ("c", tts_c, instruct_c, "<strong> + instruction names emphasized words"),
    ]:
        # Guard: strip any trailing punctuation before <|endofprompt|>
        instr_clean = instr.rstrip(".,;:!?- ").strip()
        variants.append({
            "segment": f"{seg_id}.wav",
            "variant": letter,
            "variant_description": desc,
            "tts_text": tts,
            "instruct_text": instr_clean + END_OF_PROMPT,
            "prompt_wav": prompt_wav,
            "metadata": {
                "english_text": trans["english_text"],
                "n_emphasized": len(emphasized_indices),
                "emphasized_words": emphasized_en_words,
                "speaker_context": trans["speaker_context"],
                "instruct_word_count": len(instr_clean.split()),
            },
        })
    return variants


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    seg_files = sorted(TRANSLATIONS_DIR.glob("seg_*.json"))
    print(f"Composing variants for {len(seg_files)} segments...")
    n_written = 0
    for seg_file in seg_files:
        seg_id = seg_file.stem
        variants = compose_for_segment(seg_id)
        for v in variants:
            out_path = OUT_DIR / f"{seg_id}_variant_{v['variant']}.json"
            out_path.write_text(json.dumps(v, indent=2, ensure_ascii=False), encoding="utf-8")
            n_written += 1

        a, b, c = variants
        print(f"\n--- {seg_id} ---")
        print(f"  A ({a['metadata']['instruct_word_count']}w): tts={a['tts_text']}")
        print(f"           ins={a['instruct_text']}")
        print(f"  B ({b['metadata']['instruct_word_count']}w): tts={b['tts_text']}")
        print(f"           ins={b['instruct_text']}")
        print(f"  C ({c['metadata']['instruct_word_count']}w): tts={c['tts_text']}")
        print(f"           ins={c['instruct_text']}")

    print(f"\nWrote {n_written} variant files to {OUT_DIR}")


if __name__ == "__main__":
    main()