257 lines
11 KiB
Python
Executable File
257 lines
11 KiB
Python
Executable File
#!/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()
|