laubonghaudoi/legco-speech

1"""Segment full meeting audio into sentence-level clips and write Parquet shards.2 3For each meeting in metadata.csv, this script:4  1. Parses the SRT file to get (start, end, text) segments.5  2. Uses ffmpeg to extract each segment from the opus file (stream copy, no re-encoding).6  3. Batches segments and writes Parquet shards with embedded audio bytes.7 8Usage:9    python -m scripts.segment_audio [--workers N] [--shard-size N] [--out-dir DIR]10"""11 12import argparse13import collections14import csv15import itertools16import re17import subprocess18from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait19from dataclasses import dataclass20from pathlib import Path21 22from rich.console import Console23from rich.progress import (24    BarColumn,25    MofNCompleteColumn,26    Progress,27    SpinnerColumn,28    TextColumn,29    TimeElapsedColumn,30    TimeRemainingColumn,31)32 33REPO_ROOT = Path(__file__).resolve().parent.parent34DEFAULT_SHARD_SIZE = 500035DEFAULT_WORKERS = 436 37console = Console()38 39 40@dataclass41class SrtSegment:42    index: int43    start_seconds: float44    end_seconds: float45    text: str46 47 48def _ts_to_seconds(ts: str) -> float:49    """Convert SRT timestamp (HH:MM:SS,mmm) to seconds."""50    h, m, rest = ts.split(":")51    s, ms = rest.split(",")52    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 100053 54 55def parse_srt(srt_path: Path) -> list[SrtSegment]:56    """Parse an SRT file into a list of timed segments."""57    try:58        content = srt_path.read_text(encoding="utf-8")59    except FileNotFoundError:60        return []61 62    segments: list[SrtSegment] = []63    blocks = re.split(r"\n\s*\n", content.strip())64 65    for block in blocks:66        lines = block.strip().split("\n")67        if len(lines) < 3:68            continue69 70        try:71            idx = int(lines[0].strip())72        except ValueError:73            continue74 75        ts_match = re.match(76            r"(\d{2}:\d{2}:\d{2},\d{3})\s*-->\s*(\d{2}:\d{2}:\d{2},\d{3})",77            lines[1].strip(),78        )79        if not ts_match:80            continue81 82        start = _ts_to_seconds(ts_match.group(1))83        end = _ts_to_seconds(ts_match.group(2))84        text = " ".join(l.strip() for l in lines[2:] if l.strip())85 86        if not text or end <= start:87            continue88 89        segments.append(SrtSegment(index=idx, start_seconds=start, end_seconds=end, text=text))90 91    return segments92 93 94def extract_segment_audio(opus_path: Path, start: float, duration: float) -> bytes | None:95    """Extract a segment from an opus file using ffmpeg stream copy.96 97    Returns the raw OGG/Opus bytes, or None on failure.98    """99    cmd = [100        "ffmpeg",101        "-v", "error",102        "-ss", f"{start:.3f}",103        "-i", str(opus_path),104        "-t", f"{duration:.3f}",105        "-c", "copy",106        "-f", "ogg",107        "pipe:1",108    ]109    try:110        result = subprocess.run(cmd, capture_output=True, timeout=30)111        if result.returncode != 0:112            return None113        if len(result.stdout) < 100:114            return None115        return result.stdout116    except (subprocess.TimeoutExpired, OSError):117        return None118 119 120def process_meeting(row: dict) -> list[dict]:121    """Process a single meeting: parse SRT, extract all audio segments."""122    video_id = row["id"]123    opus_path = REPO_ROOT / row["audio"]124    srt_path = REPO_ROOT / row["subtitles"]125 126    if not opus_path.exists():127        return []128 129    segments = parse_srt(srt_path)130    if not segments:131        return []132 133    results = []134    for seg in segments:135        duration = seg.end_seconds - seg.start_seconds136        audio_bytes = extract_segment_audio(opus_path, seg.start_seconds, duration)137        if audio_bytes is None:138            continue139 140        results.append({141            "video_id": video_id,142            "segment_id": seg.index,143            "audio": {"bytes": audio_bytes, "path": f"{video_id}_{seg.index:05d}.opus"},144            "text": seg.text,145            "start_time": round(seg.start_seconds, 3),146            "end_time": round(seg.end_seconds, 3),147            "duration": round(duration, 3),148        })149 150    return results151 152 153def write_shard(segments: list[dict], shard_idx: int, out_dir: Path) -> Path:154    """Write a list of segment dicts as a Parquet shard with Audio feature."""155    from datasets import Audio, Dataset, Features, Value156 157    features = Features({158        "video_id": Value("string"),159        "segment_id": Value("int32"),160        "audio": Audio(),161        "text": Value("string"),162        "start_time": Value("float64"),163        "end_time": Value("float64"),164        "duration": Value("float64"),165    })166 167    ds = Dataset.from_dict(168        {k: [s[k] for s in segments] for k in segments[0]},169        features=features,170    )171 172    path = out_dir / f"train-{shard_idx:05d}.parquet"173    ds.to_parquet(path)174    return path175 176 177def _flush_buffer(buffer: collections.deque, shard_size: int, shard_idx: int,178                   out_dir: Path, *, force: bool = False) -> int:179    """Write complete shards from buffer. Returns updated shard_idx."""180    while len(buffer) >= shard_size or (force and buffer):181        n = min(shard_size, len(buffer))182        batch = [buffer.popleft() for _ in range(n)]183        shard_path = write_shard(batch, shard_idx, out_dir)184        console.print(185            f"  Wrote shard {shard_idx} ({n} segments) -> {shard_path.name}"186        )187        del batch188        shard_idx += 1189    return shard_idx190 191 192def main() -> None:193    parser = argparse.ArgumentParser(description="Segment audio and build Parquet shards")194    parser.add_argument("--workers", type=int, default=DEFAULT_WORKERS,195                        help="Number of parallel workers (default: %(default)s)")196    parser.add_argument("--shard-size", type=int, default=DEFAULT_SHARD_SIZE,197                        help="Segments per Parquet shard (default: %(default)s)")198    parser.add_argument("--out-dir", type=Path, default=REPO_ROOT / "segmented",199                        help="Output directory for Parquet shards")200    args = parser.parse_args()201 202    args.out_dir.mkdir(parents=True, exist_ok=True)203 204    src = REPO_ROOT / "metadata.csv"205    with open(src, encoding="utf-8", newline="") as f:206        reader = csv.DictReader(f)207        rows = list(reader)208 209    console.print(f"Processing {len(rows)} meetings with {args.workers} workers")210    console.print(f"Shard size: {args.shard_size} segments")211    console.print(f"Output: {args.out_dir}")212 213    buffer: collections.deque[dict] = collections.deque()214    shard_idx = 0215    total_segments = 0216    errors = 0217    meetings_done = 0218    max_in_flight = args.workers * 2219 220    progress = Progress(221        SpinnerColumn(),222        TextColumn("[progress.description]{task.description}"),223        BarColumn(),224        MofNCompleteColumn(),225        TimeElapsedColumn(),226        TimeRemainingColumn(),227        console=console,228    )229 230    with progress:231        task = progress.add_task("Meetings processed", total=len(rows))232        rows_iter = iter(rows)233 234        with ProcessPoolExecutor(max_workers=args.workers) as pool:235            active: dict = {}236 237            # Seed the pool with an initial batch of work238            for row in itertools.islice(rows_iter, max_in_flight):239                f = pool.submit(process_meeting, row)240                active[f] = row["id"]241 242            while active:243                done, _ = wait(active, return_when=FIRST_COMPLETED)244 245                for future in done:246                    video_id = active.pop(future)247                    try:248                        segments = future.result()249                        buffer.extend(segments)250                        total_segments += len(segments)251                        del segments252                    except Exception as e:253                        errors += 1254                        console.print(f"[red]Error processing {video_id}: {e}[/red]")255 256                    # Release the Future's internal result reference257                    future._result = None258 259                    meetings_done += 1260                    progress.advance(task)261 262                    # Submit next meeting to keep the pool fed263                    row = next(rows_iter, None)264                    if row is not None:265                        f = pool.submit(process_meeting, row)266                        active[f] = row["id"]267 268                # Flush complete shards after processing each batch of done futures269                shard_idx = _flush_buffer(270                    buffer, args.shard_size, shard_idx, args.out_dir271                )272 273    # Flush remaining segments274    shard_idx = _flush_buffer(275        buffer, args.shard_size, shard_idx, args.out_dir, force=True276    )277 278    console.print(f"\n[bold green]Done![/bold green]")279    console.print(f"  Total segments: {total_segments}")280    console.print(f"  Total shards:   {shard_idx}")281    console.print(f"  Errors:         {errors}")282 283 284if __name__ == "__main__":285    main()286