From f6381d4c35f5f88e6b91843fb485af4fc0d77a62 Mon Sep 17 00:00:00 2001 From: Abdussamet Kocak Date: Tue, 20 Oct 2026 04:34:15 +0000 Subject: [PATCH] feat: Add subtitle translation script --- translate_srt.py | 256 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100755 translate_srt.py diff --git a/translate_srt.py b/translate_srt.py new file mode 100755 index 0000000..5e95118 --- /dev/null +++ b/translate_srt.py @@ -0,0 +1,256 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["httpx"] +# /// +import argparse +import functools +import json +import logging +import os +import re +import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path + +import httpx + +RETRY_STATUS_CODES = {429, 500, 502, 503, 504, 524, 529} +CONTEXT_CUES = 2 +MAX_VALIDATION_ATTEMPTS = 3 +MAX_REQUEST_ATTEMPTS = 5 +# Matches the start of an SRT time line such as "00:01:02,345 --> 00:01:04,000". +TIMELINE_PATTERN = re.compile( + r""" + ^\d{2}:\d{2}:\d{2}[,.]\d{3} # start: HH:MM:SS,mmm (comma or dot before the milliseconds) + \s+-->\s+ # arrow, with whitespace around it + \d{2}:\d{2}:\d{2}[,.]\d{3} # end, same layout + """, + re.VERBOSE, +) + +PROMPT = """You are a professional subtitle translator. Translate every item in "items" into {language}. + +Rules: +1. Return only one JSON object, no Markdown, no explanation: {{"translations":[{{"index":1,"text":"..."}}]}} +2. "translations" must contain exactly one entry per item, with the same "index" values. Never merge, split, skip or reorder items. +3. Translate each item at its own index even if word order differs; never move meaning to a neighbouring index. +4. Keep translations natural and concise, suitable for subtitles. Keep names, numbers, negation and tone accurate without removing profanity. +5. Use "context_before" and "context_after" only to resolve references and keep terminology consistent. Do not translate them. +6. Do not add notes. Keep line breaks only where the source has them. + +Input: +{payload}""" + + +@dataclass +class Cue: + """One SRT cue: 1-based index, the original time line (kept verbatim) and its text.""" + + index: int + timeline: str + text: str + + +logger = logging.getLogger(__name__) + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Translate an SRT file with an OpenRouter chat model, keeping cues 1:1.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("subtitle", type=Path, help="Input SRT file") + parser.add_argument("--translate", "-t", required=True, metavar="LANG", help="Target language, e.g. Turkish") + parser.add_argument("--model", "-m", default="google/gemini-3.5-flash-lite", help="OpenRouter chat model id") + parser.add_argument( + "--api-key", + default=os.getenv("OPENROUTER_API_KEY"), + help="OpenRouter API key (default: $OPENROUTER_API_KEY%(default).0s)", # %(default).0s keeps the key out of --help + ) + parser.add_argument("--output", "-o", type=Path, help="Output SRT path (default: ..srt)") + parser.add_argument("--batch-size", type=int, default=20, help="Cues per request") + parser.add_argument("--workers", type=int, default=8, help="Parallel requests") + return parser.parse_args() + + +def parse_srt(content: str) -> list[Cue]: + """Parse SRT text into cues, tolerating CRLF line endings and a missing index line. + + Cues are renumbered from 1; the time line is kept as written. Raises ValueError on a malformed block. + """ + cues = [] + # \n\s*\n = a blank line (possibly containing spaces) separates SRT blocks + for block in re.split(r"\n\s*\n", content.replace("\r\n", "\n").replace("\r", "\n").strip()): + lines = block.split("\n") + timeline_position = next((position for position, line in enumerate(lines[:3]) if TIMELINE_PATTERN.match(line.strip())), None) + if timeline_position is None: + raise ValueError(f"Invalid SRT block: {block[:80]!r}") + text = "\n".join(lines[timeline_position + 1 :]).strip() + cues.append(Cue(index=len(cues) + 1, timeline=lines[timeline_position].strip(), text=text)) + if not cues: + raise ValueError("SRT contains no cues") + return cues + + +def retry( + attempts: int = MAX_REQUEST_ATTEMPTS, + exceptions: tuple[type[Exception], ...] = (httpx.TransportError, httpx.HTTPStatusError), +) -> Callable: + """Decorator factory: call the function up to `attempts` times, sleeping 2**attempt seconds between tries. + + Only `exceptions` trigger a retry; the last one is re-raised when attempts run out. + """ + + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args, **kwargs): + for attempt in range(1, attempts + 1): + try: + return func(*args, **kwargs) + except exceptions as error: + if attempt == attempts: + raise + logger.warning("%s failed (%s), retry %d/%d", func.__name__, error, attempt, attempts) + time.sleep(2**attempt) + + return wrapper + + return decorator + + +@retry() +def post_request(client: httpx.Client, payload: dict) -> httpx.Response: + """POST to the chat completions endpoint. Retryable statuses (429/5xx) raise so `@retry` repeats them; + other statuses are returned for the caller to check.""" + response = client.post("/chat/completions", json=payload) + if response.status_code in RETRY_STATUS_CODES: + response.raise_for_status() + return response + + +def parse_translations(response_text: str, expected: int) -> list[str]: + """Extract the {"translations": [{"index", "text"}]} object from a model reply. + + Takes everything from the first "{" to the last "}" so Markdown fences around the JSON are ignored. + Returns the texts ordered by index. Raises ValueError unless there is exactly one non-empty entry per index. + """ + start, end = response_text.find("{"), response_text.rfind("}") + if start < 0 or end <= start: + raise ValueError("no JSON object in response") + items = json.loads(response_text[start : end + 1]).get("translations") + if not isinstance(items, list) or len(items) != expected: + raise ValueError(f"expected {expected} translations, got {len(items) if isinstance(items, list) else 'none'}") + result: list[str | None] = [None] * expected + for item in items: + position = item.get("index") + text = str(item.get("text", "")).strip() + if not isinstance(position, int) or not 1 <= position <= expected or result[position - 1] is not None: + raise ValueError(f"invalid or duplicate index {position!r}") + if not text: + raise ValueError(f"empty translation at index {position}") + result[position - 1] = text + return [text for text in result if text is not None] + + +def translate_batch( + client: httpx.Client, + cues: list[Cue], + before: list[Cue], + after: list[Cue], + language: str, + model: str, +) -> list[str]: + """Translate one batch of cues, with `before`/`after` cues sent only as context. + + Invalid replies (wrong count, bad indices, empty text) are retried up to MAX_VALIDATION_ATTEMPTS times with + the validation error appended to the prompt. Returns one translation per cue, in order. + """ + payload_text = json.dumps( + { + "context_before": [cue.text for cue in before], + "items": [{"index": position, "text": cue.text} for position, cue in enumerate(cues, start=1)], + "context_after": [cue.text for cue in after], + }, + ensure_ascii=False, + ) + prompt = PROMPT.format(language=language, payload=payload_text) + last_error = "" + for attempt in range(1, MAX_VALIDATION_ATTEMPTS + 1): + request_prompt = prompt + if last_error: + request_prompt += f"\n\nYour previous answer was invalid: {last_error}. Return the complete corrected JSON." + response = post_request( + client=client, + payload={ + "model": model, + "temperature": 0.2, + "response_format": {"type": "json_object"}, + "messages": [{"role": "user", "content": request_prompt}], + }, + ) + response.raise_for_status() + try: + content = response.json()["choices"][0]["message"]["content"] or "" + return parse_translations(response_text=content, expected=len(cues)) + except (ValueError, KeyError, IndexError, AttributeError) as error: + last_error = str(error) + logger.warning("Invalid response (%s), attempt %d/%d", last_error, attempt, MAX_VALIDATION_ATTEMPTS) + raise RuntimeError(f"Batch starting at cue {cues[0].index} failed: {last_error}") + + +def build_srt(cues: list[Cue], translations: list[str]) -> str: + """Render the translations as SRT using the original indices and time lines.""" + return "\n".join(f"{cue.index}\n{cue.timeline}\n{text}\n" for cue, text in zip(cues, translations)) + + +def main() -> None: + """Translate the SRT in batches with surrounding context and write it with the original timing.""" + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") + args = parse_args() + if not args.api_key: + raise SystemExit("Provide --api-key or set OPENROUTER_API_KEY") + if args.batch_size < 1: + raise SystemExit("--batch-size must be at least 1") + if args.workers < 1: + raise SystemExit("--workers must be at least 1") + + cues = parse_srt(content=args.subtitle.read_text(encoding="utf-8-sig")) + # [^A-Za-z0-9_-]+ = any run of characters outside letters, digits, "_" and "-"; becomes one "_" for the filename + slug = re.sub(r"[^A-Za-z0-9_-]+", "_", args.translate).strip("_").lower() + output = args.output or args.subtitle.with_name(f"{args.subtitle.stem}.{slug}.srt") + translations: list[str] = [] + + with httpx.Client( + base_url="https://openrouter.ai/api/v1", + headers={"Authorization": f"Bearer {args.api_key}"}, + timeout=300, + ) as client: + ranges = [(start, min(start + args.batch_size, len(cues))) for start in range(0, len(cues), args.batch_size)] + + def translate_range(bounds: tuple[int, int]) -> list[str]: + """Translate cues[start:end] with ±CONTEXT_CUES neighbours as context.""" + start, end = bounds + logger.info("Translating cues %d-%d of %d with %s", start + 1, end, len(cues), args.model) + return translate_batch( + client=client, + cues=cues[start:end], + before=cues[max(0, start - CONTEXT_CUES) : start], + after=cues[end : end + CONTEXT_CUES], + language=args.translate, + model=args.model, + ) + + # map() yields results in input order, so the translations stay aligned with the cues + with ThreadPoolExecutor(max_workers=args.workers) as executor: + for batch_translations in executor.map(translate_range, ranges): + translations.extend(batch_translations) + + output.write_text(build_srt(cues=cues, translations=translations), encoding="utf-8") + logger.info("Wrote %s", output) + + +if __name__ == "__main__": + main()