extract_demo_samples.py

#!/usr/bin/env python3
"""
Extract demo samples from the Cantonese dataset for inference testing
Selects samples with varying durations to demonstrate model capabilities
"""

import pyarrow as pa
import shutil
import os
import json

# Configuration
DATASET_PATH = 'data/cantonese_data_pinyin/raw.arrow'
OUTPUT_DIR = 'demo_samples'
NUM_SAMPLES = 5

# Read the dataset
print("Reading dataset...")
reader = pa.ipc.open_file(DATASET_PATH)
table = reader.read_all()

print(f"Total samples in dataset: {len(table)}")

# Create output directory
os.makedirs(OUTPUT_DIR, exist_ok=True)
print(f"\nCreated output directory: {OUTPUT_DIR}/")

# Select samples with different durations
# We want: short, medium-short, medium, medium-long, long
durations = [(table['duration'][i].as_py(), i) for i in range(len(table))]
durations.sort()

# Pick samples from different duration ranges
indices = [
    durations[len(durations) // 5][1],      # 20th percentile (short)
    durations[2 * len(durations) // 5][1],  # 40th percentile (medium-short)
    durations[len(durations) // 2][1],      # 50th percentile (medium)
    durations[3 * len(durations) // 5][1],  # 60th percentile (medium-long)
    durations[4 * len(durations) // 5][1],  # 80th percentile (long)
]

# Extract samples
samples_metadata = []

for idx, dataset_idx in enumerate(indices):
    audio_path = table['audio_path'][dataset_idx].as_py()
    text = table['text'][dataset_idx].as_py()
    duration = table['duration'][dataset_idx].as_py()

    # Check if source file exists
    if not os.path.exists(audio_path):
        print(f"⚠️  Warning: Source file not found: {audio_path}")
        continue

    # Copy audio file with new name
    dest_filename = f"demo_{idx+1}.wav"
    dest_path = os.path.join(OUTPUT_DIR, dest_filename)
    shutil.copy2(audio_path, dest_path)

    # Store metadata
    sample_info = {
        'id': idx + 1,
        'filename': dest_filename,
        'audio_path': dest_path,
        'text': text.strip(),
        'duration': round(duration, 2),
        'original_path': audio_path
    }
    samples_metadata.append(sample_info)

    print(f"\n✓ Extracted sample {idx+1}:")
    print(f"  File: {dest_filename}")
    print(f"  Duration: {duration:.2f}s")
    print(f"  Text: {text.strip()}")

# Save metadata to JSON
metadata_path = os.path.join(OUTPUT_DIR, 'samples_metadata.json')
with open(metadata_path, 'w', encoding='utf-8') as f:
    json.dump(samples_metadata, f, ensure_ascii=False, indent=2)

print(f"\n✓ Metadata saved to: {metadata_path}")

# Create a simple text file with reference info
readme_path = os.path.join(OUTPUT_DIR, 'README.txt')
with open(readme_path, 'w', encoding='utf-8') as f:
    f.write("Demo Samples for F5-TTS Inference\n")
    f.write("=" * 50 + "\n\n")
    f.write(f"Extracted from: Cantonese dataset (pinyin format)\n")
    f.write(f"Total samples: {len(samples_metadata)}\n")
    f.write(f"Text format: Pinyin with tone numbers\n\n")
    f.write("Sample Details:\n")
    f.write("-" * 50 + "\n\n")

    for sample in samples_metadata:
        f.write(f"Sample {sample['id']}: {sample['filename']}\n")
        f.write(f"  Duration: {sample['duration']}s\n")
        f.write(f"  Text: {sample['text']}\n\n")

print(f"✓ README saved to: {readme_path}")

print("\n" + "=" * 70)
print("Demo samples extraction completed! 🎉")
print("=" * 70)
print(f"\nExtracted {len(samples_metadata)} samples to: {OUTPUT_DIR}/")
print("\nNext steps:")
print("1. Review the samples in demo_samples/")
print("2. Use these samples for inference testing")
print("3. Check samples_metadata.json for details")