Sofelia TTS training checkpoints (Eliaa final + multispeaker + stage1)

c3f9ad5 verified 5 months ago

11.3 kB

#!/usr/bin/env python3
"""
Kokoro Arabic (Palestinian) Training Data Preparation
======================================================
Builds StyleTTS2 train/val lists from two Palestinian Arabic datasets:
  Speaker 0 "Kore"  : /home/ubuntu/data — Gemini TTS story segments
                      (manifests/segments.jsonl + audio/*.wav + text/segments/*.txt)
  Speaker 1 "Sirin" : hamdallah/Sofelia-Palestinian-TTS HF dataset
                      (metadata.csv in LJSpeech format + wavs/*.wav)
Both datasets are already 24kHz mono 16-bit WAV.
G2P: misaki EspeakG2P(language='ar'), with fixups mapping phonemes that are
not in Kokoro-82M's 178-token vocab onto unused-but-real vocab tokens so the
phonemic contrasts of Arabic survive:
  ħ (U+0127 voiceless pharyngeal fric.) -> ʰ (distinct h-like token)
  ʕ (U+0295 voiced pharyngeal fric.)    -> ʁ (voiced uvular fric., unused by ar)
  ˤ (U+02E4 pharyngealization/emphatic) -> ᵊ (distinct modifier token)
  ̪ (U+032A dental diacritic)           -> removed (not phonemic)
  [ ]                                   -> removed (espeak markup)
Usage:
    python3 scripts/prepare_arabic.py
"""

import json
import random
import sys
from multiprocessing import Pool
from pathlib import Path

import soundfile as sf
from tqdm import tqdm

# ── Paths ────────────────────────────────────────────────────────────────────

REPO = Path(__file__).resolve().parent.parent
TRAINING_DIR = REPO / "training"
KOKORO_CONFIG = TRAINING_DIR / "config.json"

KORE_DIR = Path("/home/ubuntu/data")
KORE_MANIFEST = KORE_DIR / "manifests" / "segments.jsonl"

SOFELIA_DIR = Path(
    "/home/ubuntu/.cache/huggingface/hub/datasets--hamdallah--Sofelia-Palestinian-TTS"
    "/snapshots/316fd018d7855ada768639d40b0deb046be89558/dataset"
)

TRAIN_LIST = TRAINING_DIR / "train_list_ar.txt"
VAL_LIST = TRAINING_DIR / "val_list_ar.txt"
OOD_FILE = TRAINING_DIR / "OOD_texts_ar.txt"
STATS_FILE = TRAINING_DIR / "stats_ar.json"

# ── Filters / split ──────────────────────────────────────────────────────────

MIN_DURATION_S = 2.0
MAX_DURATION_S = 20.0  # caps worst-case batch memory (waveform discriminator)
MIN_PHONEME_LEN = 10
VAL_RATIO = 0.02
N_OOD = 128  # entries reserved for OOD texts (excluded from training entirely)
RANDOM_SEED = 42

# ── Phoneme fixups (see module docstring) ────────────────────────────────────

PHONEME_FIXUPS = {
    "ħ": "ʰ",  # ħ -> ʰ
    "ʕ": "ʁ",  # ʕ -> ʁ
    "ˤ": "ᵊ",  # ˤ -> ᵊ
    "̪": "",        # dental diacritic -> drop
    "[": "",
    "]": "",
}

_g2p = None
_marbuta_cache = {}

import json as _json
import re

_MARBUTA_WORD = re.compile(r"([ؠ-يـ]+ة)(?![ؠ-ي])")

# Arabic punctuation → Latin equivalents that espeak passes through to the
# phoneme string (the Arabic forms are silently dropped, so the model never
# sees pause/question anchors without this).
PUNCT_MAP = str.maketrans({"،": ",", "؛": ";", "؟": "?"})

_LEXICON_PATH = REPO / "training" / "ar_lexicon.json"
_lexicon = None
_lexicon_re = None


def _load_lexicon():
    global _lexicon, _lexicon_re
    if _lexicon is None:
        if _LEXICON_PATH.exists():
            entries = _json.loads(_LEXICON_PATH.read_text())
            _lexicon = {k: v for k, v in entries.items() if not k.startswith("_")}
        else:
            _lexicon = {}
        if _lexicon:
            alts = "|".join(re.escape(w) for w in sorted(_lexicon, key=len, reverse=True))
            _lexicon_re = re.compile(rf"(?<![ؠ-يـ])({alts})(?![ؠ-يـ])")
    return _lexicon, _lexicon_re


def normalize_arabic(text, g2p):
    """Full text normalization: punctuation map + pronunciation lexicon + ة fix."""
    text = text.translate(PUNCT_MAP)
    lex, lex_re = _load_lexicon()
    if lex_re is not None:
        text = lex_re.sub(lambda m: lex[m.group(1)], text)
    return normalize_taa_marbuta(text, g2p)


def normalize_taa_marbuta(text, g2p):
    """Fix espeak's OOV fallback that reads word-final ة as /t/ (e.g. قريبة→qriːbt).
    espeak pronounces in-lexicon ة words correctly (مدرسة→a) but renders unknown
    ones with a final t. For each ة-final word, G2P it alone: if the result ends
    in 't', rewrite the trailing ة to a bare fatha (قريبَ), which espeak reads as
    the correct pausal/dialect final 'a' (qriːba). Cached per word type.
    """

    def repl(m):
        word = m.group(1)
        fixed = _marbuta_cache.get(word)
        if fixed is None:
            try:
                ph, _ = g2p(word)
                broken = ph.strip().rstrip(".").endswith("t")
            except Exception:
                broken = False
            fixed = word[:-1] + "َ" if broken else word  # ة → fatha
            _marbuta_cache[word] = fixed
        return fixed

    return _MARBUTA_WORD.sub(repl, text)


def _init_worker():
    global _g2p
    from misaki import espeak

    _g2p = espeak.EspeakG2P(language="ar")


def _phonemize(item):
    """(idx, text) -> (idx, phonemes or None)"""
    idx, text = item
    try:
        text = normalize_arabic(text, _g2p)
        ph, _ = _g2p(text)
    except Exception:
        return idx, None
    for old, new in PHONEME_FIXUPS.items():
        ph = ph.replace(old, new)
    return idx, ph.strip()


def collect_kore():
    """Speaker 0: Gemini story segments."""
    entries = []
    dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
    with open(KORE_MANIFEST) as f:
        for line in f:
            e = json.loads(line)
            dur = e.get("durationSeconds", 0)
            if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
                dropped["bad_duration"] += 1
                continue
            seg_id = e["segmentId"]
            wav = KORE_DIR / "audio" / f"{seg_id}.wav"
            txt = KORE_DIR / "text" / "segments" / f"{seg_id}.txt"
            if not wav.exists() or not txt.exists():
                dropped["missing_file"] += 1
                continue
            text = txt.read_text().strip()
            if not text:
                dropped["empty_text"] += 1
                continue
            entries.append({"wav": str(wav), "text": text, "speaker": 0, "dur": dur})
    print(f"Kore: kept {len(entries):,}  dropped {dropped}")
    return entries


def collect_sofelia():
    """Speaker 1: Sofelia Palestinian TTS."""
    entries = []
    dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
    with open(SOFELIA_DIR / "metadata.csv") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split("|")
            if len(parts) < 2:
                continue
            wav = SOFELIA_DIR / "wavs" / f"{parts[0]}.wav"
            if not wav.exists():
                dropped["missing_file"] += 1
                continue
            text = parts[1].strip()
            if not text:
                dropped["empty_text"] += 1
                continue
            info = sf.info(str(wav))
            dur = info.frames / info.samplerate
            if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
                dropped["bad_duration"] += 1
                continue
            entries.append({"wav": str(wav), "text": text, "speaker": 1, "dur": dur})
    print(f"Sofelia: kept {len(entries):,}  dropped {dropped}")
    return entries


def main():
    vocab = json.load(open(KOKORO_CONFIG))["vocab"]

    entries = collect_kore() + collect_sofelia()
    print(f"Total entries before G2P: {len(entries):,}")

    # ── G2P (parallel) ───────────────────────────────────────────────────
    items = [(i, e["text"]) for i, e in enumerate(entries)]
    with Pool(processes=8, initializer=_init_worker) as pool:
        for idx, ph in tqdm(
            pool.imap_unordered(_phonemize, items, chunksize=64),
            total=len(items),
            desc="G2P",
        ):
            entries[idx]["ipa"] = ph

    # ── Validate against Kokoro vocab ────────────────────────────────────
    valid = []
    n_g2p_fail = n_short = n_unknown = 0
    unknown_chars = {}
    for e in entries:
        ph = e.get("ipa")
        if not ph:
            n_g2p_fail += 1
            continue
        if len(ph) < MIN_PHONEME_LEN:
            n_short += 1
            continue
        bad = [ch for ch in ph if ch not in vocab]
        if bad:
            n_unknown += 1
            for ch in bad:
                unknown_chars[ch] = unknown_chars.get(ch, 0) + 1
            continue
        valid.append(e)

    print(f"Valid after G2P: {len(valid):,}")
    print(f"  G2P failures   : {n_g2p_fail}")
    print(f"  Too short      : {n_short}")
    print(f"  Unknown phoneme: {n_unknown}")
    if unknown_chars:
        for ch, c in sorted(unknown_chars.items(), key=lambda x: -x[1])[:20]:
            print(f"    {ch!r} U+{ord(ch):04X}: {c}")

    if not valid:
        print("ERROR: no valid entries")
        sys.exit(1)

    # ── Split: OOD / val / train ─────────────────────────────────────────
    rng = random.Random(RANDOM_SEED)
    rng.shuffle(valid)
    ood_entries = valid[:N_OOD]
    rest = valid[N_OOD:]
    n_val = max(1, int(len(rest) * VAL_RATIO))
    val_entries = rest[:n_val]
    train_entries = rest[n_val:]

    def write_list(path, ents):
        with open(path, "w") as f:
            for e in ents:
                f.write(f"{e['wav']}|{e['ipa']}|{e['speaker']}\n")

    write_list(TRAIN_LIST, train_entries)
    write_list(VAL_LIST, val_entries)

    with open(OOD_FILE, "w") as f:
        f.write("\n".join(e["ipa"] for e in ood_entries) + "\n")

    # ── Stats ────────────────────────────────────────────────────────────
    def hours(ents):
        return round(sum(e["dur"] for e in ents) / 3600, 2)

    stats = {
        "train": {"files": len(train_entries), "hours": hours(train_entries)},
        "val": {"files": len(val_entries), "hours": hours(val_entries)},
        "ood_texts": len(ood_entries),
        "speakers": {
            "0_kore": {
                "files": sum(1 for e in valid if e["speaker"] == 0),
                "hours": hours([e for e in valid if e["speaker"] == 0]),
            },
            "1_sirin": {
                "files": sum(1 for e in valid if e["speaker"] == 1),
                "hours": hours([e for e in valid if e["speaker"] == 1]),
            },
        },
    }
    with open(STATS_FILE, "w") as f:
        json.dump(stats, f, indent=2)

    print(json.dumps(stats, indent=2))
    print(f"\nWrote {TRAIN_LIST}\n      {VAL_LIST}\n      {OOD_FILE}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""
	Kokoro Arabic (Palestinian) Training Data Preparation
	======================================================
	Builds StyleTTS2 train/val lists from two Palestinian Arabic datasets:

	Speaker 0 "Kore" : /home/ubuntu/data — Gemini TTS story segments
	(manifests/segments.jsonl + audio/.wav + text/segments/.txt)
	Speaker 1 "Sirin" : hamdallah/Sofelia-Palestinian-TTS HF dataset
	(metadata.csv in LJSpeech format + wavs/*.wav)

	Both datasets are already 24kHz mono 16-bit WAV.

	G2P: misaki EspeakG2P(language='ar'), with fixups mapping phonemes that are
	not in Kokoro-82M's 178-token vocab onto unused-but-real vocab tokens so the
	phonemic contrasts of Arabic survive:

	ħ (U+0127 voiceless pharyngeal fric.) -> ʰ (distinct h-like token)
	ʕ (U+0295 voiced pharyngeal fric.) -> ʁ (voiced uvular fric., unused by ar)
	ˤ (U+02E4 pharyngealization/emphatic) -> ᵊ (distinct modifier token)
	̪ (U+032A dental diacritic) -> removed (not phonemic)
	[ ] -> removed (espeak markup)

	Usage:
	python3 scripts/prepare_arabic.py
	"""

	import json
	import random
	import sys
	from multiprocessing import Pool
	from pathlib import Path

	import soundfile as sf
	from tqdm import tqdm

	# ── Paths ────────────────────────────────────────────────────────────────────

	REPO = Path(__file__).resolve().parent.parent
	TRAINING_DIR = REPO / "training"
	KOKORO_CONFIG = TRAINING_DIR / "config.json"

	KORE_DIR = Path("/home/ubuntu/data")
	KORE_MANIFEST = KORE_DIR / "manifests" / "segments.jsonl"

	SOFELIA_DIR = Path(
	"/home/ubuntu/.cache/huggingface/hub/datasets--hamdallah--Sofelia-Palestinian-TTS"
	"/snapshots/316fd018d7855ada768639d40b0deb046be89558/dataset"
	)

	TRAIN_LIST = TRAINING_DIR / "train_list_ar.txt"
	VAL_LIST = TRAINING_DIR / "val_list_ar.txt"
	OOD_FILE = TRAINING_DIR / "OOD_texts_ar.txt"
	STATS_FILE = TRAINING_DIR / "stats_ar.json"

	# ── Filters / split ──────────────────────────────────────────────────────────

	MIN_DURATION_S = 2.0
	MAX_DURATION_S = 20.0 # caps worst-case batch memory (waveform discriminator)
	MIN_PHONEME_LEN = 10
	VAL_RATIO = 0.02
	N_OOD = 128 # entries reserved for OOD texts (excluded from training entirely)
	RANDOM_SEED = 42

	# ── Phoneme fixups (see module docstring) ────────────────────────────────────

	PHONEME_FIXUPS = {
	"ħ": "ʰ", # ħ -> ʰ
	"ʕ": "ʁ", # ʕ -> ʁ
	"ˤ": "ᵊ", # ˤ -> ᵊ
	"̪": "", # dental diacritic -> drop
	"[": "",
	"]": "",
	}

	_g2p = None
	_marbuta_cache = {}

	import json as _json
	import re

	_MARBUTA_WORD = re.compile(r"([ؠ-يـ]+ة)(?![ؠ-ي])")

	# Arabic punctuation → Latin equivalents that espeak passes through to the
	# phoneme string (the Arabic forms are silently dropped, so the model never
	# sees pause/question anchors without this).
	PUNCT_MAP = str.maketrans({"،": ",", "؛": ";", "؟": "?"})

	_LEXICON_PATH = REPO / "training" / "ar_lexicon.json"
	_lexicon = None
	_lexicon_re = None


	def _load_lexicon():
	global _lexicon, _lexicon_re
	if _lexicon is None:
	if _LEXICON_PATH.exists():
	entries = _json.loads(_LEXICON_PATH.read_text())
	_lexicon = {k: v for k, v in entries.items() if not k.startswith("_")}
	else:
	_lexicon = {}
	if _lexicon:
	alts = "\|".join(re.escape(w) for w in sorted(_lexicon, key=len, reverse=True))
	_lexicon_re = re.compile(rf"(?<![ؠ-يـ])({alts})(?![ؠ-يـ])")
	return _lexicon, _lexicon_re


	def normalize_arabic(text, g2p):
	"""Full text normalization: punctuation map + pronunciation lexicon + ة fix."""
	text = text.translate(PUNCT_MAP)
	lex, lex_re = _load_lexicon()
	if lex_re is not None:
	text = lex_re.sub(lambda m: lex[m.group(1)], text)
	return normalize_taa_marbuta(text, g2p)


	def normalize_taa_marbuta(text, g2p):
	"""Fix espeak's OOV fallback that reads word-final ة as /t/ (e.g. قريبة→qriːbt).

	espeak pronounces in-lexicon ة words correctly (مدرسة→a) but renders unknown
	ones with a final t. For each ة-final word, G2P it alone: if the result ends
	in 't', rewrite the trailing ة to a bare fatha (قريبَ), which espeak reads as
	the correct pausal/dialect final 'a' (qriːba). Cached per word type.
	"""

	def repl(m):
	word = m.group(1)
	fixed = _marbuta_cache.get(word)
	if fixed is None:
	try:
	ph, _ = g2p(word)
	broken = ph.strip().rstrip(".").endswith("t")
	except Exception:
	broken = False
	fixed = word[:-1] + "َ" if broken else word # ة → fatha
	_marbuta_cache[word] = fixed
	return fixed

	return _MARBUTA_WORD.sub(repl, text)


	def _init_worker():
	global _g2p
	from misaki import espeak

	_g2p = espeak.EspeakG2P(language="ar")


	def _phonemize(item):
	"""(idx, text) -> (idx, phonemes or None)"""
	idx, text = item
	try:
	text = normalize_arabic(text, _g2p)
	ph, _ = _g2p(text)
	except Exception:
	return idx, None
	for old, new in PHONEME_FIXUPS.items():
	ph = ph.replace(old, new)
	return idx, ph.strip()


	def collect_kore():
	"""Speaker 0: Gemini story segments."""
	entries = []
	dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
	with open(KORE_MANIFEST) as f:
	for line in f:
	e = json.loads(line)
	dur = e.get("durationSeconds", 0)
	if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
	dropped["bad_duration"] += 1
	continue
	seg_id = e["segmentId"]
	wav = KORE_DIR / "audio" / f"{seg_id}.wav"
	txt = KORE_DIR / "text" / "segments" / f"{seg_id}.txt"
	if not wav.exists() or not txt.exists():
	dropped["missing_file"] += 1
	continue
	text = txt.read_text().strip()
	if not text:
	dropped["empty_text"] += 1
	continue
	entries.append({"wav": str(wav), "text": text, "speaker": 0, "dur": dur})
	print(f"Kore: kept {len(entries):,} dropped {dropped}")
	return entries


	def collect_sofelia():
	"""Speaker 1: Sofelia Palestinian TTS."""
	entries = []
	dropped = {"bad_duration": 0, "missing_file": 0, "empty_text": 0}
	with open(SOFELIA_DIR / "metadata.csv") as f:
	for line in f:
	line = line.strip()
	if not line:
	continue
	parts = line.split("\|")
	if len(parts) < 2:
	continue
	wav = SOFELIA_DIR / "wavs" / f"{parts[0]}.wav"
	if not wav.exists():
	dropped["missing_file"] += 1
	continue
	text = parts[1].strip()
	if not text:
	dropped["empty_text"] += 1
	continue
	info = sf.info(str(wav))
	dur = info.frames / info.samplerate
	if not (MIN_DURATION_S <= dur <= MAX_DURATION_S):
	dropped["bad_duration"] += 1
	continue
	entries.append({"wav": str(wav), "text": text, "speaker": 1, "dur": dur})
	print(f"Sofelia: kept {len(entries):,} dropped {dropped}")
	return entries


	def main():
	vocab = json.load(open(KOKORO_CONFIG))["vocab"]

	entries = collect_kore() + collect_sofelia()
	print(f"Total entries before G2P: {len(entries):,}")

	# ── G2P (parallel) ───────────────────────────────────────────────────
	items = [(i, e["text"]) for i, e in enumerate(entries)]
	with Pool(processes=8, initializer=_init_worker) as pool:
	for idx, ph in tqdm(
	pool.imap_unordered(_phonemize, items, chunksize=64),
	total=len(items),
	desc="G2P",
	):
	entries[idx]["ipa"] = ph

	# ── Validate against Kokoro vocab ────────────────────────────────────
	valid = []
	n_g2p_fail = n_short = n_unknown = 0
	unknown_chars = {}
	for e in entries:
	ph = e.get("ipa")
	if not ph:
	n_g2p_fail += 1
	continue
	if len(ph) < MIN_PHONEME_LEN:
	n_short += 1
	continue
	bad = [ch for ch in ph if ch not in vocab]
	if bad:
	n_unknown += 1
	for ch in bad:
	unknown_chars[ch] = unknown_chars.get(ch, 0) + 1
	continue
	valid.append(e)

	print(f"Valid after G2P: {len(valid):,}")
	print(f" G2P failures : {n_g2p_fail}")
	print(f" Too short : {n_short}")
	print(f" Unknown phoneme: {n_unknown}")
	if unknown_chars:
	for ch, c in sorted(unknown_chars.items(), key=lambda x: -x[1])[:20]:
	print(f" {ch!r} U+{ord(ch):04X}: {c}")

	if not valid:
	print("ERROR: no valid entries")
	sys.exit(1)

	# ── Split: OOD / val / train ─────────────────────────────────────────
	rng = random.Random(RANDOM_SEED)
	rng.shuffle(valid)
	ood_entries = valid[:N_OOD]
	rest = valid[N_OOD:]
	n_val = max(1, int(len(rest) * VAL_RATIO))
	val_entries = rest[:n_val]
	train_entries = rest[n_val:]

	def write_list(path, ents):
	with open(path, "w") as f:
	for e in ents:
	f.write(f"{e['wav']}\|{e['ipa']}\|{e['speaker']}\n")

	write_list(TRAIN_LIST, train_entries)
	write_list(VAL_LIST, val_entries)

	with open(OOD_FILE, "w") as f:
	f.write("\n".join(e["ipa"] for e in ood_entries) + "\n")

	# ── Stats ────────────────────────────────────────────────────────────
	def hours(ents):
	return round(sum(e["dur"] for e in ents) / 3600, 2)

	stats = {
	"train": {"files": len(train_entries), "hours": hours(train_entries)},
	"val": {"files": len(val_entries), "hours": hours(val_entries)},
	"ood_texts": len(ood_entries),
	"speakers": {
	"0_kore": {
	"files": sum(1 for e in valid if e["speaker"] == 0),
	"hours": hours([e for e in valid if e["speaker"] == 0]),
	},
	"1_sirin": {
	"files": sum(1 for e in valid if e["speaker"] == 1),
	"hours": hours([e for e in valid if e["speaker"] == 1]),
	},
	},
	}
	with open(STATS_FILE, "w") as f:
	json.dump(stats, f, indent=2)

	print(json.dumps(stats, indent=2))
	print(f"\nWrote {TRAIN_LIST}\n {VAL_LIST}\n {OOD_FILE}")


	if __name__ == "__main__":
	main()