adam/transcript_dataset.py

from __future__ import annotations

import json
import re
import shutil
from dataclasses import dataclass
from importlib.util import find_spec
from pathlib import Path
from typing import Any


@dataclass(slots=True)
class TranscriptResult:
    available: bool
    message: str
    samples: list[dict[str, Any]]
    output_folder: str = ""
    backend: str = ""


@dataclass(frozen=True, slots=True)
class TranscriptionBackend:
    id: str
    name: str
    available: bool
    requirement: str


def clean_transcript(text: str) -> str:
    text = re.sub(r"\s+", " ", text.replace("\r", " ").replace("\n", " ")).strip()
    text = re.sub(r"\s+([,.!?;:])", r"\1", text)
    return text


def split_transcript(text: str, *, max_chars: int = 900) -> list[str]:
    cleaned = clean_transcript(text)
    if not cleaned:
        return []
    sentences = re.split(r"(?<=[.!?])\s+", cleaned)
    samples: list[str] = []
    current = ""
    for sentence in sentences:
        if not sentence:
            continue
        if current and len(current) + 1 + len(sentence) > max_chars:
            samples.append(current.strip())
            current = sentence
        else:
            current = f"{current} {sentence}".strip()
    if current:
        samples.append(current.strip())
    return samples


def available_transcription_backends() -> list[TranscriptionBackend]:
    return [
        TranscriptionBackend("whisper", "OpenAI Whisper", find_spec("whisper") is not None, "pip install openai-whisper"),
        TranscriptionBackend("faster_whisper", "faster-whisper", find_spec("faster_whisper") is not None, "pip install faster-whisper"),
    ]


def _select_backend(preferred: str = "auto") -> TranscriptionBackend | None:
    backends = available_transcription_backends()
    if preferred != "auto":
        return next((backend for backend in backends if backend.id == preferred and backend.available), None)
    return next((backend for backend in backends if backend.available), None)


def _transcribe_with_backend(video: Path, backend: TranscriptionBackend) -> str:
    if backend.id == "whisper":
        import whisper

        model = whisper.load_model("base")
        result = model.transcribe(str(video))
        return str(result.get("text", ""))
    if backend.id == "faster_whisper":
        from faster_whisper import WhisperModel

        model = WhisperModel("base", device="auto", compute_type="auto")
        segments, _info = model.transcribe(str(video))
        return " ".join(segment.text for segment in segments)
    raise RuntimeError(f"Unsupported transcription backend: {backend.id}")


def transcript_videos_to_dataset(
    videos: list[str | Path],
    output_folder: str | Path,
    *,
    backend: str = "auto",
    max_chars: int = 900,
    preserve_metadata: bool = True,
) -> TranscriptResult:
    output = Path(output_folder).expanduser().resolve()
    if not videos:
        return TranscriptResult(False, "Choose at least one local video.", [])
    if not shutil.which("ffmpeg"):
        return TranscriptResult(
            False,
            "FFmpeg was not found. ADAM can keep the workflow ready, but audio extraction/transcription needs FFmpeg.",
            [],
        )
    selected_backend = _select_backend(backend)
    if selected_backend is None:
        requirements = ", ".join(item.requirement for item in available_transcription_backends())
        return TranscriptResult(
            False,
            f"No local transcription backend is installed. Install one of: {requirements}.",
            [],
        )
    samples: list[dict[str, Any]] = []
    try:
        for video in videos:
            path = Path(video).expanduser().resolve()
            if not path.is_file():
                return TranscriptResult(False, f"Video not found: {path}", [])
            text = _transcribe_with_backend(path, selected_backend)
            for index, sample in enumerate(split_transcript(text, max_chars=max_chars), 1):
                item = {"text": sample}
                if preserve_metadata:
                    item.update({"source_video": str(path), "sample_index": index})
                samples.append(item)
    except Exception as exc:
        return TranscriptResult(False, f"{selected_backend.name} could not transcribe the selected video(s): {exc}", [], backend=selected_backend.id)
    output.mkdir(parents=True, exist_ok=True)
    txt_path = output / "transcript_samples.txt"
    jsonl_path = output / "transcript_samples.jsonl"
    txt_path.write_text("\n\n".join(item["text"] for item in samples), encoding="utf-8")
    jsonl_path.write_text(
        "\n".join(json.dumps(item, ensure_ascii=False) for item in samples) + ("\n" if samples else ""),
        encoding="utf-8",
    )
    return TranscriptResult(True, f"Exported {len(samples):,} transcript sample(s) with {selected_backend.name}.", samples, str(output), selected_backend.id)