Spaces:
Running
Running
File size: 2,634 Bytes
0b6ab33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """Format synthesized data into HuggingFace Dataset for training."""
import json
import re
from pathlib import Path
from datasets import Dataset, Audio, DatasetDict
def extract_tags(tagged_text: str) -> list[dict]:
"""Extract tags and their positions from tagged text."""
tags = []
for match in re.finditer(r'\[([^\]]+)\]', tagged_text):
tags.append({
"tag": match.group(1),
"start_char": match.start(),
"end_char": match.end(),
})
return tags
def strip_tags(tagged_text: str) -> str:
"""Remove all [tags] from text, leaving plain transcription."""
return re.sub(r'\[[^\]]+\]\s*', '', tagged_text).strip()
def has_emphasis(text: str) -> bool:
"""Check if text contains CAPS emphasis (2+ consecutive uppercase letters forming a word)."""
return bool(re.search(r'\b[A-Z]{2,}\b', text))
def format_dataset(
manifest_path: str = "data/audio/manifest.json",
output_dir: str = "data/processed",
) -> DatasetDict:
"""Convert manifest + audio files into a HuggingFace Dataset."""
with open(manifest_path) as f:
manifest = json.load(f)
records = [m for m in manifest if m.get("synthesis_success", True)]
rows = []
for record in records:
tagged_text = record["tagged_text"]
plain_text = record.get("plain_text", strip_tags(tagged_text))
tags = extract_tags(tagged_text)
rows.append({
"audio": record["audio_path"],
"tagged_text": tagged_text,
"plain_text": plain_text,
"tags": json.dumps(tags),
"tag_count": len(tags),
"has_emphasis": has_emphasis(tagged_text),
"slice_type": record.get("slice_type", "unknown"),
"domain": record.get("domain", "unknown"),
"voice_id": record.get("voice_id", "unknown"),
})
ds = Dataset.from_dict({k: [r[k] for r in rows] for k in rows[0]})
ds = ds.cast_column("audio", Audio(sampling_rate=16000))
# Train/val/test split: 80/10/10
ds_split = ds.train_test_split(test_size=0.2, seed=42)
val_test = ds_split["test"].train_test_split(test_size=0.5, seed=42)
final = DatasetDict({
"train": ds_split["train"],
"validation": val_test["train"],
"test": val_test["test"],
})
Path(output_dir).mkdir(parents=True, exist_ok=True)
final.save_to_disk(output_dir)
print(f"Dataset saved to {output_dir}")
print(f" Train: {len(final['train'])}, Val: {len(final['validation'])}, Test: {len(final['test'])}")
return final
if __name__ == "__main__":
format_dataset()
|