srtforge 0.1.2__tar.gz → 0.2.0__tar.gz

{srtforge-0.1.2/src/srtforge.egg-info → srtforge-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: srtforge
-Version: 0.1.2
+Version: 0.2.0
 Summary: Generate, translate, and burn SRT subtitles locally with MLX Whisper.
 Author: rromanv
 License-Expression: MIT
@@ -22,7 +22,8 @@ Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: mlx-whisper>=0.4
-Requires-Dist: mlx-lm>=0.20
+Provides-Extra: translate
+Requires-Dist: mlx-lm>=0.20; extra == "translate"
 Dynamic: license-file
 
 # srtforge
@@ -76,6 +77,12 @@ From PyPI:
 pipx install --python python3.11 srtforge
 ```
 
+If you want local LLM translation support, install the translation extra:
+
+```bash
+pipx install --python python3.11 'srtforge[translate]'
+```
+
 `srtforge` currently supports Python 3.10-3.12. If your default Python is newer
 than that, such as Python 3.14, pass a supported interpreter explicitly with
 `--python`.
@@ -86,6 +93,12 @@ From GitHub:
 pipx install --python python3.11 git+https://github.com/rromanv/srtforge.git
 ```
 
+With translation support from GitHub:
+
+```bash
+pipx install --python python3.11 'srtforge[translate] @ git+https://github.com/rromanv/srtforge.git'
+```
+
 For local development:
 
 ```bash
@@ -94,6 +107,7 @@ cd srtforge
 python3.11 -m venv .venv
 source .venv/bin/activate
 pip install -e .
+pip install -e '.[translate]'  # optional, for `srtforge -t ...`
 ```
 
 Or into a virtual environment:
@@ -102,6 +116,7 @@ Or into a virtual environment:
 python3.11 -m venv .venv
 source .venv/bin/activate
 pip install srtforge
+pip install 'srtforge[translate]'  # optional, for `srtforge -t ...`
 ```
 
 ## Troubleshooting
@@ -177,6 +192,8 @@ srtforge video.mp4 -t "Brazilian Portuguese"
 srtforge video.mp4 -t ja --translate-model mlx-community/Qwen3.5-9B-OptiQ-4bit
 ```
 
+Translation requires installing `srtforge[translate]`.
+
 Burn subtitles into a video:
 
 ```bash
@@ -184,6 +201,9 @@ srtforge merge video.mp4 video.srt -o final.mp4
 srtforge merge video.mp4 video.es.srt --crf 16
 ```
 
+Existing outputs are not overwritten by default. Pass `--force` when you intend
+to replace an existing `.srt` or merged video output.
+
 Run:
 
 ```bash
@@ -225,6 +245,7 @@ src/srtforge/
   segment.py     # sentence-aware re-cueing, wrapping, pacing
   translate.py   # context-aware local-LLM translation
   transcript.py  # reference transcript reading + correction
+  _types.py      # shared typed subtitle data shapes
   merge.py       # burn subtitles into video with ffmpeg/libass
   srt.py         # SRT rendering
   cli.py         # argparse CLI
@@ -234,6 +255,7 @@ tests/
   test_segment.py
   test_srt.py
   test_transcript.py
+  test_translate.py
 ```
 
 ## Development

{srtforge-0.1.2 → srtforge-0.2.0}/README.md

@@ -49,6 +49,12 @@ From PyPI:
 pipx install --python python3.11 srtforge
 ```
 
+If you want local LLM translation support, install the translation extra:
+
+```bash
+pipx install --python python3.11 'srtforge[translate]'
+```
+
 `srtforge` currently supports Python 3.10-3.12. If your default Python is newer
 than that, such as Python 3.14, pass a supported interpreter explicitly with
 `--python`.
@@ -59,6 +65,12 @@ From GitHub:
 pipx install --python python3.11 git+https://github.com/rromanv/srtforge.git
 ```
 
+With translation support from GitHub:
+
+```bash
+pipx install --python python3.11 'srtforge[translate] @ git+https://github.com/rromanv/srtforge.git'
+```
+
 For local development:
 
 ```bash
@@ -67,6 +79,7 @@ cd srtforge
 python3.11 -m venv .venv
 source .venv/bin/activate
 pip install -e .
+pip install -e '.[translate]'  # optional, for `srtforge -t ...`
 ```
 
 Or into a virtual environment:
@@ -75,6 +88,7 @@ Or into a virtual environment:
 python3.11 -m venv .venv
 source .venv/bin/activate
 pip install srtforge
+pip install 'srtforge[translate]'  # optional, for `srtforge -t ...`
 ```
 
 ## Troubleshooting
@@ -150,6 +164,8 @@ srtforge video.mp4 -t "Brazilian Portuguese"
 srtforge video.mp4 -t ja --translate-model mlx-community/Qwen3.5-9B-OptiQ-4bit
 ```
 
+Translation requires installing `srtforge[translate]`.
+
 Burn subtitles into a video:
 
 ```bash
@@ -157,6 +173,9 @@ srtforge merge video.mp4 video.srt -o final.mp4
 srtforge merge video.mp4 video.es.srt --crf 16
 ```
 
+Existing outputs are not overwritten by default. Pass `--force` when you intend
+to replace an existing `.srt` or merged video output.
+
 Run:
 
 ```bash
@@ -198,6 +217,7 @@ src/srtforge/
   segment.py     # sentence-aware re-cueing, wrapping, pacing
   translate.py   # context-aware local-LLM translation
   transcript.py  # reference transcript reading + correction
+  _types.py      # shared typed subtitle data shapes
   merge.py       # burn subtitles into video with ffmpeg/libass
   srt.py         # SRT rendering
   cli.py         # argparse CLI
@@ -207,6 +227,7 @@ tests/
   test_segment.py
   test_srt.py
   test_transcript.py
+  test_translate.py
 ```
 
 ## Development

{srtforge-0.1.2 → srtforge-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "srtforge"
-version = "0.1.2"
+version = "0.2.0"
 description = "Generate, translate, and burn SRT subtitles locally with MLX Whisper."
 readme = "README.md"
 requires-python = ">=3.10,<3.13"
@@ -23,6 +23,10 @@ classifiers = [
 ]
 dependencies = [
     "mlx-whisper>=0.4",
+]
+
+[project.optional-dependencies]
+translate = [
     "mlx-lm>=0.20",
 ]
 

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/__init__.py

@@ -1,3 +1,3 @@
 """srtforge: generate SRT subtitles from video using local Whisper (MLX)."""
 
-__version__ = "0.1.2"
+__version__ = "0.2.0"

srtforge-0.2.0/src/srtforge/_types.py

@@ -0,0 +1,24 @@
+"""Shared subtitle data shapes used across the pipeline."""
+
+from __future__ import annotations
+
+from typing import TypedDict
+
+
+class WhisperWord(TypedDict, total=False):
+    word: str
+    start: float
+    end: float
+
+
+class Segment(TypedDict, total=False):
+    start: float
+    end: float
+    text: str
+    words: list[WhisperWord]
+
+
+class Cue(TypedDict, total=False):
+    start: float
+    end: float
+    text: str

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/cli.py

@@ -17,6 +17,10 @@ from .transcript import correct_segments, initial_prompt_from_transcript, read_t
 from .translate import DEFAULT_TRANSLATE_MODEL, TranslationError, translate_segments
 
 
+class CliError(RuntimeError):
+    """Raised for invalid CLI arguments that argparse cannot express alone."""
+
+
 def _log(msg: str) -> None:
     print(msg, file=sys.stderr, flush=True)
 
@@ -49,6 +53,29 @@ def _translated_path(output: Path, language: str) -> Path:
     return output.with_suffix(f".{tag}{output.suffix}")
 
 
+def _same_path(a: Path, b: Path) -> bool:
+    return a.resolve(strict=False) == b.resolve(strict=False)
+
+
+def _ensure_srt_output(path: Path) -> None:
+    if path.suffix.lower() != ".srt":
+        raise CliError(f"Subtitle output must end in .srt: {path}")
+
+
+def _ensure_can_write_output(
+    output: Path,
+    *,
+    force: bool,
+    protected: list[Path],
+    kind: str,
+) -> None:
+    for path in protected:
+        if _same_path(output, path):
+            raise CliError(f"Refusing to overwrite input {kind}: {output}")
+    if output.exists() and not force:
+        raise CliError(f"Output already exists: {output}. Pass --force to overwrite.")
+
+
 def build_parser() -> argparse.ArgumentParser:
     p = argparse.ArgumentParser(
         prog="srtforge",
@@ -61,6 +88,10 @@ def build_parser() -> argparse.ArgumentParser:
         "-o", "--output", type=Path, default=None,
         help="Output .srt path (default: alongside the video).",
     )
+    p.add_argument(
+        "--force", action="store_true",
+        help="Overwrite existing output files.",
+    )
     p.add_argument(
         "-m", "--model", default=DEFAULT_MODEL,
         help=f"Whisper model / HF repo (default: {DEFAULT_MODEL}).",
@@ -122,6 +153,23 @@ def _gen_main(argv: list[str] | None = None) -> int:
     need_words = args.word_timestamps or args.resegment
     wav_path = None
     try:
+        _ensure_srt_output(output)
+        _ensure_can_write_output(
+            output,
+            force=args.force,
+            protected=[args.video],
+            kind="video",
+        )
+        if translating:
+            t_out = _translated_path(output, args.translate)
+            _ensure_srt_output(t_out)
+            _ensure_can_write_output(
+                t_out,
+                force=args.force,
+                protected=[args.video, output],
+                kind="video or subtitle output",
+            )
+
         transcript_text = read_transcript(args.transcript) if args.transcript else ""
         _log(f"[1/{total}] Extracting audio from {args.video} ...")
         wav_path = extract_audio(args.video)
@@ -172,7 +220,7 @@ def _gen_main(argv: list[str] | None = None) -> int:
         _log("Done.")
         return 0
 
-    except (FFmpegError, FileNotFoundError, TranslationError) as e:
+    except (CliError, FFmpegError, FileNotFoundError, TranslationError) as e:
         _log(f"Error: {e}")
         return 1
     except KeyboardInterrupt:
@@ -202,6 +250,10 @@ def build_merge_parser() -> argparse.ArgumentParser:
         "--preset", default="slow",
         help="x264 speed/efficiency preset (default: slow).",
     )
+    p.add_argument(
+        "--force", action="store_true",
+        help="Overwrite the output video if it already exists.",
+    )
     return p
 
 
@@ -209,6 +261,12 @@ def _merge_main(argv: list[str]) -> int:
     args = build_merge_parser().parse_args(argv)
     output = args.output or args.video.with_suffix(f".subbed{args.video.suffix}")
     try:
+        _ensure_can_write_output(
+            output,
+            force=args.force,
+            protected=[args.video, args.subtitles],
+            kind="video or subtitle file",
+        )
         _log(f"Burning {args.subtitles} into {args.video} (crf={args.crf}, preset={args.preset}) ...")
         burn_subtitles(
             args.video,
@@ -219,7 +277,7 @@ def _merge_main(argv: list[str]) -> int:
         )
         _log(f"Done. Wrote {output}")
         return 0
-    except (FFmpegError, FileNotFoundError) as e:
+    except (CliError, FFmpegError, FileNotFoundError) as e:
         _log(f"Error: {e}")
         return 1
     except KeyboardInterrupt:

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/segment.py

@@ -12,7 +12,8 @@ from __future__ import annotations
 
 import re
 from dataclasses import dataclass, field
-from typing import Any
+
+from ._types import Cue, Segment
 
 # --- Netflix-style defaults ------------------------------------------------- #
 MAX_CPL = 42          # max characters per line
@@ -71,7 +72,7 @@ class _Sentence:
 # --------------------------------------------------------------------------- #
 # Word flattening & sentence splitting
 # --------------------------------------------------------------------------- #
-def _flatten_words(segments: list[dict[str, Any]]) -> list[_Word]:
+def _flatten_words(segments: list[Segment]) -> list[_Word]:
     words: list[_Word] = []
     for seg in segments:
         for w in seg.get("words") or []:
@@ -197,7 +198,7 @@ def _split_long_sentence(sent: _Sentence, opts: Opts) -> list[_Sentence]:
     return chunks
 
 
-def _build_cues(sentences: list[_Sentence], opts: Opts) -> list[dict[str, Any]]:
+def _build_cues(sentences: list[_Sentence], opts: Opts) -> list[Cue]:
     # 1) Ensure every unit fits within the line limits: split over-long ones.
     units: list[_Sentence] = []
     for s in sentences:
@@ -207,7 +208,7 @@ def _build_cues(sentences: list[_Sentence], opts: Opts) -> list[dict[str, Any]]:
             units.extend(_split_long_sentence(s, opts))
 
     # 2) Mostly one sentence per cue; merge a SHORT trailing unit if both fit.
-    cues: list[dict[str, Any]] = []
+    cues: list[Cue] = []
     i = 0
     while i < len(units):
         cur = units[i]
@@ -234,7 +235,7 @@ def _build_cues(sentences: list[_Sentence], opts: Opts) -> list[dict[str, Any]]:
 # --------------------------------------------------------------------------- #
 # Pacing / timing pass
 # --------------------------------------------------------------------------- #
-def _apply_pacing(cues: list[dict[str, Any]], opts: Opts) -> list[dict[str, Any]]:
+def _apply_pacing(cues: list[Cue], opts: Opts) -> list[Cue]:
     n = len(cues)
     for k, cue in enumerate(cues):
         chars = len(cue["text"].replace("\n", " "))
@@ -267,7 +268,7 @@ def _apply_pacing(cues: list[dict[str, Any]], opts: Opts) -> list[dict[str, Any]
 # --------------------------------------------------------------------------- #
 # Public API
 # --------------------------------------------------------------------------- #
-def fit_cues(cues: list[dict[str, Any]], opts: Opts | None = None) -> list[dict[str, Any]]:
+def fit_cues(cues: list[Segment], opts: Opts | None = None) -> list[Cue]:
     """Fit each cue's text within the line limits, splitting over-capacity cues.
 
     Used for translated text, whose length differs from the source. A cue whose
@@ -275,7 +276,7 @@ def fit_cues(cues: list[dict[str, Any]], opts: Opts | None = None) -> list[dict[
     boundaries, with the original time span apportioned by character count.
     """
     opts = opts or Opts()
-    out: list[dict[str, Any]] = []
+    out: list[Cue] = []
     for c in cues:
         text = " ".join((c.get("text") or "").split())
         if not text:
@@ -308,24 +309,38 @@ def fit_cues(cues: list[dict[str, Any]], opts: Opts | None = None) -> list[dict[
     return _apply_pacing(out, opts)
 
 
-def resegment(segments: list[dict[str, Any]], opts: Opts | None = None) -> list[dict[str, Any]]:
-    """Rebuild Whisper segments into readable, sentence-aligned cues.
-
-    Requires word-level timestamps (segment['words']). Segments lacking words
-    fall back to char-based splitting of their own text.
-    """
-    opts = opts or Opts()
-    has_words = any(seg.get("words") for seg in segments)
-    if not has_words:
-        return segments
-
-    words = _flatten_words(segments)
-    if not words:
-        return []
+def _resegment_words(words: list[_Word], opts: Opts) -> list[Cue]:
     sentences = _split_sentences(words)
     cues = _build_cues(sentences, opts)
     cues = _apply_pacing(cues, opts)
-    # Wrap text into <=max_lines lines.
     for cue in cues:
         cue["text"] = wrap_lines(cue["text"], opts.max_cpl, opts.max_lines)
     return cues
+
+
+def resegment(segments: list[Segment], opts: Opts | None = None) -> list[Cue]:
+    """Rebuild Whisper segments into readable, sentence-aligned cues.
+
+    Uses word-level timestamps (segment['words']) where available. Segments
+    lacking words fall back to char-based splitting of their own text.
+    """
+    opts = opts or Opts()
+    out: list[Cue] = []
+    pending_words: list[_Word] = []
+
+    def flush_words() -> None:
+        if pending_words:
+            out.extend(_resegment_words(pending_words, opts))
+            pending_words.clear()
+
+    for seg in segments:
+        words = _flatten_words([seg])
+        if words:
+            pending_words.extend(words)
+            continue
+
+        flush_words()
+        out.extend(fit_cues([seg], opts))
+
+    flush_words()
+    return out

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/srt.py

@@ -2,7 +2,11 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterable
+from typing import Iterable
+
+from ._types import Segment
+
+MIN_CUE_DURATION = 0.001
 
 
 def _format_timestamp(seconds: float) -> str:
@@ -16,16 +20,29 @@ def _format_timestamp(seconds: float) -> str:
     return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
 
 
-def segments_to_srt(segments: Iterable[dict[str, Any]]) -> str:
+def _normalize_timing(start: float, end: float, previous_end: float) -> tuple[float, float]:
+    start = max(0.0, start, previous_end)
+    end = max(start + MIN_CUE_DURATION, end)
+    return start, end
+
+
+def segments_to_srt(segments: Iterable[Segment]) -> str:
     """Render Whisper segments to an SRT document string."""
     blocks: list[str] = []
     index = 1
+    previous_end = 0.0
     for seg in segments:
         text = (seg.get("text") or "").strip()
         if not text:
             continue
-        start = _format_timestamp(float(seg.get("start", 0.0)))
-        end = _format_timestamp(float(seg.get("end", 0.0)))
+        start_seconds, end_seconds = _normalize_timing(
+            float(seg.get("start", 0.0)),
+            float(seg.get("end", 0.0)),
+            previous_end,
+        )
+        previous_end = end_seconds
+        start = _format_timestamp(start_seconds)
+        end = _format_timestamp(end_seconds)
         blocks.append(f"{index}\n{start} --> {end}\n{text}\n")
         index += 1
     return "\n".join(blocks)

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/transcribe.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any
+
+from ._types import Segment
 
 DEFAULT_MODEL = "mlx-community/whisper-large-v3-turbo"
 
@@ -14,7 +15,7 @@ def transcribe(
     language: str | None = None,
     word_timestamps: bool = False,
     initial_prompt: str | None = None,
-) -> list[dict[str, Any]]:
+) -> list[Segment]:
     """Transcribe audio and return Whisper segments.
 
     Each segment is a dict with at least 'start', 'end', 'text'.

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/transcript.py

@@ -5,11 +5,15 @@ import re
 import sys
 from pathlib import Path
 
+from ._types import Segment
+
 MIN_GLOBAL_SIMILARITY = 0.30
 MIN_BLOCK_SIMILARITY = 0.60
+MIN_TOKEN_SIMILARITY = 0.72
 MAX_INITIAL_PROMPT_CHARS = 1_000
 
 _TRAILING_PUNCT_RE = re.compile(r"([^\w]+)$")
+_EDGE_PUNCT_RE = re.compile(r"^\W+|\W+$")
 
 
 def read_transcript(path: Path) -> str:
@@ -41,6 +45,22 @@ def _transfer_trailing_punct(whisper_token: str, replacement: str) -> str:
     return stripped + punct
 
 
+def _normalize_token(token: str) -> str:
+    return _EDGE_PUNCT_RE.sub("", token.casefold())
+
+
+def _tokens_close(whisper_token: str, reference_token: str) -> bool:
+    whisper = _normalize_token(whisper_token)
+    reference = _normalize_token(reference_token)
+    if not whisper or not reference or whisper == reference:
+        return False
+    max_len = max(len(whisper), len(reference))
+    if abs(len(whisper) - len(reference)) > max(2, max_len // 2):
+        return False
+    similarity = difflib.SequenceMatcher(None, whisper, reference).ratio()
+    return similarity >= MIN_TOKEN_SIMILARITY
+
+
 def initial_prompt_from_transcript(text: str, max_chars: int = MAX_INITIAL_PROMPT_CHARS) -> str:
     """Return a short transcript prefix suitable for Whisper's initial prompt."""
     text = " ".join(text.split())
@@ -51,7 +71,7 @@ def initial_prompt_from_transcript(text: str, max_chars: int = MAX_INITIAL_PROMP
     return truncated or text[:max_chars].strip()
 
 
-def correct_segments(segments: list[dict], reference_text: str) -> list[dict]:
+def correct_segments(segments: list[Segment], reference_text: str) -> list[Segment]:
     if not reference_text or not reference_text.strip():
         return segments
 
@@ -98,7 +118,7 @@ def correct_segments(segments: list[dict], reference_text: str) -> list[dict]:
         w_slice = whisper_lower[i1:i2]
         r_slice = ref_lower[j1:j2]
         if len(w_slice) == 1 and len(r_slice) == 1:
-            if w_slice[0] != r_slice[0]:
+            if _tokens_close(whisper_tokens[i1], ref_tokens[j1]):
                 corrected[i1] = _transfer_trailing_punct(
                     whisper_tokens[i1], ref_tokens[j1]
                 )
@@ -109,12 +129,12 @@ def correct_segments(segments: list[dict], reference_text: str) -> list[dict]:
             if block_sm.ratio() >= MIN_BLOCK_SIMILARITY:
                 for k in range(len(w_slice)):
                     if k < len(r_slice):
-                        if w_slice[k] != r_slice[k]:
+                        if _tokens_close(whisper_tokens[i1 + k], ref_tokens[j1 + k]):
                             corrected[i1 + k] = _transfer_trailing_punct(
                                 whisper_tokens[i1 + k], ref_tokens[j1 + k]
                             )
 
-    new_segments: list[dict] = []
+    new_segments: list[Segment] = []
     for idx, seg in enumerate(segments):
         start_idx, end_idx, source = seg_offsets[idx]
         seg_tokens = corrected[start_idx:end_idx]

{srtforge-0.1.2 → srtforge-0.2.0}/src/srtforge/translate.py

@@ -10,7 +10,8 @@ Uses ``mlx_lm`` text models.
 from __future__ import annotations
 
 import re
-from typing import Any
+
+from ._types import Segment
 
 DEFAULT_TRANSLATE_MODEL = "mlx-community/Qwen3.5-9B-OptiQ-4bit"
 
@@ -111,6 +112,8 @@ def _parse_numbered(text: str, expected: int) -> dict[int, str] | None:
             continue
         idx = int(m.group(1))
         if 1 <= idx <= expected:
+            if idx in out:
+                return None
             out[idx] = m.group(2).strip()
     if len(out) == expected and all(i in out for i in range(1, expected + 1)):
         return out
@@ -121,13 +124,13 @@ def _parse_numbered(text: str, expected: int) -> dict[int, str] | None:
 # Public API
 # --------------------------------------------------------------------------- #
 def translate_segments(
-    segments: list[dict[str, Any]],
+    segments: list[Segment],
     target_language: str,
     model: str = DEFAULT_TRANSLATE_MODEL,
     window: int = 40,
     context_cues: int = 6,
     backend: _Backend | None = None,
-) -> list[dict[str, Any]]:
+) -> list[Segment]:
     """Return new segments with text translated into ``target_language``.
 
     Timestamps are preserved. Cue count is preserved (1:1 mapping). Translated