feat: Add subtitle translation script
This commit is contained in:
abdus committed
1 parent
f867799876
commit
f6381d4c35
1 file changed
+256
Executable
+256
@@ -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: <name>.<lang>.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()
|
||||
Reference in new issue
Block a user