subtitlemaker 0.1.0__py3-none-any.whl

subtitlemaker/segmenter.py

@@ -0,0 +1,304 @@
+"""Turn Whisper segments into formatted subtitle cues.
+
+Pure logic: no printing, no file access, no model.
+"""
+
+import re
+
+
+def clean_text(text):
+    """Clean up Whisper text."""
+
+    text = text.strip()
+
+    # Collapse multiple whitespace characters
+    text = re.sub(r"\s+", " ", text)
+
+    # Remove spaces before punctuation
+    text = re.sub(r"\s+([,.!?;:])", r"\1", text)
+
+    # Clean brackets
+    text = re.sub(r"([(\[])\s+", r"\1", text)
+    text = re.sub(r"\s+([)\]])", r"\1", text)
+
+    return text.strip()
+
+
+def split_into_lines(text, style):
+    """
+    Split subtitle into a maximum of `style.max_lines` lines.
+
+    Attempts to:
+    - keep lines below `style.max_line_length`
+    - balance line lengths
+    - prefer punctuation boundaries
+    """
+
+    words = text.split()
+
+    if not words:
+        return []
+
+    # Fits on one line
+    if len(text) <= style.max_line_length:
+        return [text]
+
+    # Find the best two-line split
+    best_split = None
+    best_score = float("inf")
+
+    for i in range(1, len(words)):
+        left = " ".join(words[:i])
+        right = " ".join(words[i:])
+
+        if len(left) > style.max_line_length:
+            continue
+
+        if len(right) > style.max_line_length:
+            continue
+
+        # Prefer balanced lines
+        score = abs(len(left) - len(right))
+
+        # Prefer punctuation at the end of line 1
+        if left.endswith((".", ",", "!", "?", ";", ":")):
+            score -= 15
+
+        if score < best_score:
+            best_score = score
+            best_split = (left, right)
+
+    if best_split:
+        return list(best_split)
+
+    # Fallback: split approximately in half.
+    midpoint = len(words) // 2
+
+    return [
+        " ".join(words[:midpoint]),
+        " ".join(words[midpoint:]),
+    ]
+
+
+# A cue is never shown for less than this, whatever the neighbours are doing.
+_MIN_VISIBLE = 0.1
+
+
+def _fits(text, style):
+    """True if `text` can be laid out within the line and length limits."""
+
+    if len(text) <= style.max_line_length:
+        return True
+
+    if style.max_lines < 2:
+        return False
+
+    words = text.split()
+
+    return any(
+        len(" ".join(words[:i])) <= style.max_line_length
+        and len(" ".join(words[i:])) <= style.max_line_length
+        for i in range(1, len(words))
+    )
+
+
+def _join(tokens):
+    return clean_text("".join(token[0] for token in tokens))
+
+
+def _best_break(tokens):
+    """
+    Where to end the cue when `tokens` plus one more word no longer fit.
+
+    Prefers the end of a sentence, then a clause (comma, semicolon, colon)
+    in the second half of the text. Returns the number of tokens to keep;
+    `len(tokens)` means no good break was found.
+    """
+
+    sentence = clause = None
+    total = len(_join(tokens))
+
+    for i in range(1, len(tokens)):
+
+        word = tokens[i - 1][0].strip()
+
+        if word.endswith((".", "!", "?")):
+            sentence = i
+
+        elif word.endswith((",", ";", ":")) and len(_join(tokens[:i])) * 2 >= total:
+            clause = i
+
+    if sentence is not None:
+        return sentence
+
+    return clause if clause is not None else len(tokens)
+
+
+def _chunk(tokens, style):
+    """
+    Split a segment's tokens into runs that each fit one cue.
+
+    A token is `(text, start, end)` where the text carries its own leading
+    space and the times may be None. A run only grows while it can still be
+    laid out in valid lines, so every chunk (except a single over-long word)
+    splits cleanly in `split_into_lines`.
+    """
+
+    chunks = []
+    current = []
+
+    for token in tokens:
+
+        if not current or _fits(_join(current + [token]), style):
+            current.append(token)
+            continue
+
+        cut = _best_break(current)
+
+        # The leftover plus the new word must still fit; if not, break here.
+        if not _fits(_join(current[cut:] + [token]), style):
+            cut = len(current)
+
+        chunks.append(current[:cut])
+        current = current[cut:] + [token]
+
+    if current:
+        chunks.append(current)
+
+    return chunks
+
+
+def _segment_tokens(segment, text):
+    """
+    The words of a segment as `(text, start, end)` tokens.
+
+    Uses Whisper's word timestamps when it provided them and they agree with
+    the segment text. Otherwise the tokens carry no timing and the caller
+    spreads the segment's duration over them.
+    """
+
+    words = getattr(segment, "words", None) or []
+
+    tokens = [
+        (word.word, word.start, word.end)
+        for word in words
+        if word.word.strip()
+    ]
+
+    usable = (
+        tokens
+        and all(start is not None and end is not None for _, start, end in tokens)
+        and "".join(word for word, _, _ in tokens).split() == text.split()
+    )
+
+    if usable:
+        return tokens
+
+    return [(" " + word, None, None) for word in text.split()]
+
+
+def _segment_cues(segment, style):
+    """Cues for one Whisper segment, with their speech timing and no padding."""
+
+    text = clean_text(segment.text)
+
+    if not text:
+        return []
+
+    chunks = _chunk(_segment_tokens(segment, text), style)
+
+    if all(token[1] is not None for chunk in chunks for token in chunk):
+        return [
+            {
+                "start": float(chunk[0][1]),
+                "end": max(float(chunk[-1][2]), float(chunk[0][1])),
+                "text": _join(chunk),
+            }
+            for chunk in chunks
+        ]
+
+    # No word timing: share the segment's duration out by text length.
+    start = float(segment.start)
+    end = float(segment.end)
+
+    total_chars = sum(len(_join(chunk)) for chunk in chunks)
+
+    cues = []
+    cursor = start
+
+    for chunk in chunks:
+
+        chunk_text = _join(chunk)
+        chunk_end = cursor + (end - start) * len(chunk_text) / total_chars
+
+        cues.append({"start": cursor, "end": chunk_end, "text": chunk_text})
+
+        cursor = chunk_end
+
+    return cues
+
+
+def _apply_timing(cues, style):
+    """
+    Make cue timing readable without moving speech around.
+
+    In order, for every cue:
+      1. It starts when its speech starts, unless the previous cue is still
+         up (only possible when source segments overlap).
+      2. It is held at least `min_duration` and long enough to read at
+         `max_cps`, but never longer than `max_duration`.
+      3. It must end `min_gap` before the next cue's speech starts. This wins
+         over rule 2: readability padding never delays or overlaps the next
+         cue, so a very short gap can leave a cue under `min_duration`.
+    """
+
+    timed = []
+    previous_end = None
+
+    for i, cue in enumerate(cues):
+
+        start = cue["start"]
+
+        if previous_end is not None:
+            start = max(start, previous_end + style.min_gap)
+
+        plain = cue["text"].replace("\n", " ")
+
+        wanted = max(
+            cue["end"],
+            start + style.min_duration,
+            start + len(plain) / style.max_cps,
+        )
+
+        end = min(wanted, start + style.max_duration)
+
+        if i + 1 < len(cues):
+            end = min(end, cues[i + 1]["start"] - style.min_gap)
+
+        end = max(end, start + _MIN_VISIBLE)
+
+        timed.append({"start": start, "end": end, "text": cue["text"]})
+
+        previous_end = end
+
+    return timed
+
+
+def build_subtitles(segments, style):
+    """Convert Whisper segments into formatted subtitles."""
+
+    cues = []
+
+    for segment in segments:
+
+        for cue in _segment_cues(segment, style):
+
+            lines = split_into_lines(cue["text"], style)
+
+            if not lines:
+                continue
+
+            cue["text"] = "\n".join(lines)
+            cues.append(cue)
+
+    return _apply_timing(cues, style)

subtitlemaker/transcribe.py

@@ -0,0 +1,253 @@
+"""Whisper model access and language handling."""
+
+import logging
+import re
+from dataclasses import replace
+from pathlib import Path
+
+import ctranslate2
+import numpy as np
+from faster_whisper import WhisperModel, available_models
+
+from ._cuda import preload_cuda_libs
+
+try:
+    from faster_whisper.tokenizer import _LANGUAGE_CODES
+except ImportError:
+    _LANGUAGE_CODES = None
+
+log = logging.getLogger(__name__)
+
+# user/model, as used on the Hugging Face Hub
+_REPO_ID = re.compile(r"[A-Za-z0-9][\w.-]*/[\w.-]+")
+
+# Defaults used when the model or compute type is left on "auto".
+DEFAULT_MODEL = {"cuda": "large-v3", "cpu": "small"}
+DEFAULT_COMPUTE_TYPE = {"cuda": "float16", "cpu": "int8"}
+
+CUDA_HINT = (
+    "To use the GPU, install the CUDA 12 libraries with "
+    "'pip install subtitlemaker[cuda]' (or use --device cpu)."
+)
+
+
+class ModelError(RuntimeError):
+    """A model could not be selected or loaded. The message is user-facing."""
+
+
+def cuda_available():
+    """True if CTranslate2 can see at least one CUDA device."""
+
+    try:
+        return ctranslate2.get_cuda_device_count() > 0
+
+    except RuntimeError:
+        return False
+
+
+def resolve_model_config(config, cuda=None):
+    """
+    Turn "auto" settings into concrete ones.
+
+    - device:       cuda if a GPU is visible, otherwise cpu (with a warning)
+    - model:        large-v3 on cuda, small on cpu
+    - compute type: float16 on cuda, int8 on cpu
+
+    `cuda` overrides GPU detection (used by tests). Raises `ModelError`
+    for choices that cannot work: --device cuda without a GPU, an unknown
+    model name, or a compute type the device does not support.
+    """
+
+    if cuda is None:
+        cuda = cuda_available()
+
+    if config.device == "auto":
+
+        if cuda:
+            device = "cuda"
+
+        else:
+            device = "cpu"
+            log.warning(
+                "no CUDA device found, using the CPU. "
+                "Transcription will be much slower."
+            )
+
+    elif config.device == "cuda" and not cuda:
+        raise ModelError(
+            "--device cuda was requested but no CUDA device was found. "
+            "Use --device auto or --device cpu."
+        )
+
+    else:
+        device = config.device
+
+    name = config.name
+
+    if name == "auto":
+        name = DEFAULT_MODEL[device]
+
+    elif name not in available_models() and not Path(name).exists():
+
+        if _REPO_ID.fullmatch(name):
+            pass  # a Hugging Face repo id, downloaded on first use
+
+        elif "/" in name or name.startswith("."):
+            raise ModelError(f"model path not found: {name}")
+
+        else:
+            raise ModelError(
+                f"unknown model '{name}'. Use one of: "
+                f"{', '.join(available_models())}; a Hugging Face repo "
+                f"id such as 'user/model'; or a path to a converted model."
+            )
+
+    compute_type = config.compute_type
+
+    if compute_type == "auto":
+        compute_type = DEFAULT_COMPUTE_TYPE[device]
+
+    else:
+        supported = ctranslate2.get_supported_compute_types(device)
+
+        if compute_type not in supported:
+            raise ModelError(
+                f"compute type '{compute_type}' is not supported on "
+                f"{device}. Supported: {', '.join(sorted(supported))}."
+            )
+
+    return replace(
+        config,
+        name=name,
+        device=device,
+        compute_type=compute_type,
+    )
+
+
+def load_model(config, allow_cpu_fallback=False):
+    """
+    Load the model described by a resolved `ModelConfig`.
+
+    Returns `(model, config_used)`. With `allow_cpu_fallback`, a CUDA
+    failure (typically missing CUDA libraries) falls back to the CPU with
+    a warning instead of failing; the returned config shows what was used.
+    """
+
+    if config.device == "cuda":
+        preload_cuda_libs()
+
+    try:
+        return _create(config), config
+
+    except RuntimeError as e:
+
+        if config.device != "cuda":
+            raise ModelError(f"couldn't load model '{config.name}': {e}") from e
+
+        if not allow_cpu_fallback:
+            raise ModelError(
+                f"couldn't start CUDA: {e}\n{CUDA_HINT}"
+            ) from e
+
+        log.warning("CUDA could not be used (%s)", e)
+        log.warning(CUDA_HINT)
+
+        # The CUDA defaults no longer fit: pick CPU ones unless the user
+        # chose the model or compute type themselves.
+        fallback = replace(
+            config,
+            device="cpu",
+            name=(
+                DEFAULT_MODEL["cpu"]
+                if config.name == DEFAULT_MODEL["cuda"]
+                else config.name
+            ),
+            compute_type=DEFAULT_COMPUTE_TYPE["cpu"],
+        )
+
+        log.warning(
+            "falling back to the CPU with model '%s'.", fallback.name
+        )
+
+        try:
+            return _create(fallback), fallback
+
+        except RuntimeError as e2:
+            raise ModelError(f"couldn't load model '{fallback.name}': {e2}") from e2
+
+
+def _create(config):
+    model = WhisperModel(
+        config.name,
+        device=config.device,
+        compute_type=config.compute_type,
+    )
+
+    if config.device == "cuda":
+        _warm_up(model)
+
+    return model
+
+
+def _warm_up(model):
+    """
+    Run one tiny encode on the GPU.
+
+    CTranslate2 loads cuBLAS and cuDNN lazily, so missing libraries (or too
+    little GPU memory) would otherwise only surface in the middle of the
+    first transcription. Raises RuntimeError in that case.
+    """
+
+    silence = np.zeros(
+        (1, model.model.n_mels, model.feature_extractor.nb_max_frames),
+        dtype=np.float32,
+    )
+
+    model.model.encode(
+        ctranslate2.StorageView.from_array(silence),
+        to_cpu=True,
+    )
+
+
+def parse_language_argument(value):
+    """
+    Parse the -l/--language argument.
+
+    Accepts a single language code ("he") or a comma or
+    space separated list of codes ("he,en") for audio that
+    contains more than one language.
+
+    Returns (languages, multilingual), where `multilingual`
+    is True when more than one code was given.
+    """
+
+    if not value:
+        return [], False
+
+    languages = []
+
+    for code in re.split(r"[,\s]+", value.strip()):
+
+        code = code.lower()
+
+        if not code:
+            continue
+
+        if code not in languages:
+            languages.append(code)
+
+    if _LANGUAGE_CODES is not None:
+
+        unknown = [
+            code
+            for code in languages
+            if code not in _LANGUAGE_CODES
+        ]
+
+        if unknown:
+            raise ValueError(
+                f"unknown language code(s): "
+                f"{', '.join(unknown)}"
+            )
+
+    return languages, len(languages) > 1