From c58426523521045231e8b1e1435daa00677e561e Mon Sep 17 00:00:00 2001 From: Abdussamet Kocak Date: Thu, 25 Jun 2026 08:46:51 +0300 Subject: [PATCH] feat(subtitle-translator): add contextual chunk translation --- subtitle_translator.py | 390 +++++++++++++++++++++++++++-------------- 1 file changed, 256 insertions(+), 134 deletions(-) diff --git a/subtitle_translator.py b/subtitle_translator.py index ac8f83e..5e2471d 100755 --- a/subtitle_translator.py +++ b/subtitle_translator.py @@ -4,6 +4,7 @@ # /// import argparse +from concurrent.futures import ThreadPoolExecutor import json import logging import os @@ -13,15 +14,21 @@ from typing import Iterator import srt import httpx + +logger = logging.getLogger(__name__) + api_key = os.getenv("OPENROUTER_API_KEY") if not api_key: raise RuntimeError("OPENROUTER_API_KEY environment variable not set") + client = httpx.Client( headers={"Authorization": f"Bearer {api_key}"}, base_url="https://openrouter.ai/api/v1", - timeout=120, + timeout=30, ) +pool = ThreadPoolExecutor(max_workers=5) + # retry is a decorator to retry the function on failure def retry(max_retries: int = 3, delay: float = 1.0): @@ -31,7 +38,7 @@ def retry(max_retries: int = 3, delay: float = 1.0): try: return func(*args, **kwargs) except Exception as e: - logging.debug(f"Error: {e}, retrying {attempt + 1}/{max_retries}...") + logger.warning(f"Error: {e}, retrying {attempt + 1}/{max_retries}...") if attempt < max_retries - 1: time.sleep(delay) raise RuntimeError("Max retries exceeded") @@ -41,92 +48,82 @@ def retry(max_retries: int = 3, delay: float = 1.0): return decorator -def combine_entries(entries: list[srt.Subtitle], max_entries: float = 2) -> list[srt.Subtitle]: +def combine_entries(entries: list[srt.Subtitle], max_entries: int = 2) -> list[srt.Subtitle]: """ - Merge adjacent subtitles in groups of 'factor' if: - - They are no more than 1s apart (start of next - end of prev <= 1s) - - The total word count is fewer than 8 + Merge adjacent subtitles in groups of up to 'max_entries' if: + - The gap between the end of one and the start of the next is <= 1 second + - The total word count of the group is fewer than 8 """ if not entries: return [] - merged = [] + merged_subtitles = [] + current_index = 1 i = 0 - idx = 1 - n = len(entries) - while i < n: + total_subs = len(entries) + + while i < total_subs: group = [entries[i]] - j = 1 - while j < int(max_entries) and (i + j) < n: - prev = group[-1] - curr = entries[i + j] - gap = (curr.start - prev.end).total_seconds() - total_words = sum(len(s.content.split()) for s in group) + len(curr.content.split()) - if gap <= 1 and total_words < 8: - group.append(curr) - j += 1 + group_size = 1 + + while group_size < max_entries and (i + group_size) < total_subs: + last_sub = group[-1] + next_sub = entries[i + group_size] + time_gap = (next_sub.start - last_sub.end).total_seconds() + group_word_count = sum(len(sub.content.split()) for sub in group) + len(next_sub.content.split()) + + if time_gap <= 1 and group_word_count < 18: + group.append(next_sub) + group_size += 1 else: break + start_time = group[0].start end_time = group[-1].end - combined_text = " ".join(x.content.replace("\n", " ") for x in group) + combined_text = " ".join(sub.content.replace("\n", " ") for sub in group) words = combined_text.split() formatted_text = format_subtitle_text(words) - merged.append(srt.Subtitle(index=idx, start=start_time, end=end_time, content=formatted_text, proprietary="")) - idx += 1 + merged_subtitles.append(srt.Subtitle(index=current_index, start=start_time, end=end_time, content=formatted_text, proprietary="")) + current_index += 1 i += len(group) - return merged + + return merged_subtitles -def format_subtitle_text(words: list[str], max_words_per_line: int = 10) -> str: +def format_subtitle_text(words: list[str], max_words_per_line: int = 8) -> str: """ - Format a list of words into subtitle text with appropriate line breaks. - - Args: - words: List of words to format - max_words_per_line: Maximum words allowed per line - - Returns: - Formatted subtitle text with line breaks + Format a list of words into subtitle text with balanced line breaks. + Try to keep lines with the same amount of words, but put at most max_words_per_line words per line. + If there's 1 word left for the last line, join it with the previous line. """ if not words: return "" - if len(words) <= max_words_per_line: + n = len(words) + if n <= max_words_per_line: return " ".join(words) + num_lines = (n + max_words_per_line - 1) // max_words_per_line + + # Try to distribute words as evenly as possible + base = n // num_lines + extra = n % num_lines + lines = [] - i = 0 - - while i < len(words): - remaining = len(words) - i - - # Handle orphan words (single word on last line) - if remaining <= 1 and lines: - lines[-1] += " " + words[i] + idx = 0 + for i in range(num_lines): + # Distribute the remainder words to the first 'extra' lines + count = base + (1 if i < extra else 0) + # If this is the last line and only 1 word left, join it with previous + if i == num_lines - 1 and count == 1 and lines: + lines[-1] += " " + words[idx] break - - # Calculate words for this line - words_for_line = min(max_words_per_line, remaining) - lines.append(" ".join(words[i : i + words_for_line])) - i += words_for_line + lines.append(" ".join(words[idx : idx + count])) + idx += count return "\n".join(lines) -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) - p.add_argument("input", type=Path, help="Input SRT file") - p.add_argument("output", type=Path, help="Output SRT file") - p.add_argument("--translate", type=str, help="Language code for translation") - p.add_argument( - "--condense", - type=int, - help="Merge adjacent short subtitles in groups of N if close together", - ) - return p.parse_args() - - def clean_entries(subs: list[srt.Subtitle]) -> list[srt.Subtitle]: """ Remove all occurrences of '\\h' from subtitle content. @@ -145,12 +142,6 @@ def remove_overlaps(subs: list[srt.Subtitle]) -> list[srt.Subtitle]: If the beginning of a subtitle entry matches the end of the previous entry, those duplicate lines are removed from the current subtitle. - - Args: - subs: List of subtitle entries - - Returns: - List of cleaned subtitle entries with duplicates removed """ if not subs: return [] @@ -219,107 +210,238 @@ TOOL_SCHEMA_SAVE_SUBTITLES = { } -def translate_entries(entries: list[srt.Subtitle], lang: str) -> list[srt.Subtitle]: - logging.info(f"Translating subtitles to {lang} via OpenRouter") - - @retry(max_retries=5) - def translate_chunk(chunk: list[srt.Subtitle]) -> list[srt.Subtitle]: - prompt = f""" +@retry(max_retries=8) +def _translate_chunk_with_context( + chunk: list[srt.Subtitle], + context_before: list[srt.Subtitle], + context_after: list[srt.Subtitle], + lang: str, + model: str, +) -> list[srt.Subtitle]: + """ + Translates a single chunk of subtitles using preceding and succeeding subtitles as context. + """ + prompt = f""" You are an expert subtitle translator. -Translate the following subtitle entries from English to {lang}. -For each entry, only translate the 'content' field. Do not change the 'id'. -Do NOT translate proper nouns, names, or technical terms; keep them in English. -Do NOT add, remove, or merge entries. Do not include any explanations or comments. +You will be provided with a JSON object containing three lists: 'context_before', 'main_chunk', and 'context_after'. -Return exactly the same number of entries as provided. For each input entry, there must be one output entry with the same 'id'. -If an entry is empty, return it as empty. +Your task is to translate ONLY the subtitles in the 'main_chunk' from English to {lang}. +Use the 'context_before' and 'context_after' lists to understand the full conversational flow, resolve ambiguities, and ensure tonal consistency. +DO NOT translate the content of the 'context_before' or 'context_after' lists in your final output. + +- Translate idiomatically, as a native speaker would. +- If the content is too verbose, rewrite it to be more concise while preserving the original meaning. +- Do not censor profanity or slang; keep the original tone and style. +- Add punctuation and capitalization as needed. +- Do NOT translate proper nouns, names, or technical terms; keep them in English. +- Do NOT add, remove, or merge entries. Maintain the exact number of entries. +- Ensure the translated subtitles are in the same order as the original 'main_chunk'. +- If an entry is part of a continuing sentence, you may use '...' at the end/start if it helps readability. +- The output must contain {len(chunk)} entries. +- Each entry must have the same 'id' as the corresponding entry in 'main_chunk'. +- You are translating for a {lang}-speaking audience. Because {lang} word order is different from English, you are allowed to move information between entries within the 'main_chunk' to ensure the Turkish is grammatically correct and the verb is properly placed. However, you MUST return exactly the same number of entries as provided. + +Return ONLY the translated entries for the 'main_chunk'. The number of entries in your response must exactly match the number of entries in the input 'main_chunk'. +For each input entry in 'main_chunk', there must be one output entry with the same 'id'. Use your save_subtitles tool to return the results. """ - chunk_json = json.dumps( - [ - { - "id": s.index, - "content": s.content.strip(), - } - for s in chunk - ], - ensure_ascii=False, + # Create the structured payload for the API + payload_data = { + "context_before": [{"id": s.index, "content": s.content.strip()} for s in context_before], + "main_chunk": [{"id": s.index, "content": s.content.strip()} for s in chunk], + "context_after": [{"id": s.index, "content": s.content.strip()} for s in context_after], + } + + payload = { + "model": model, + "messages": [ + {"role": "system", "content": prompt}, + {"role": "user", "content": json.dumps(payload_data, ensure_ascii=False)}, + ], + "tools": [TOOL_SCHEMA_SAVE_SUBTITLES], + "tool_choice": {"type": "function", "function": {"name": TOOL_SCHEMA_SAVE_SUBTITLES["function"]["name"]}}, + } + + res = client.post(url="/chat/completions", json=payload) + if res.is_error: + logger.error(f"Error translating subtitles: {res.text}") + res.raise_for_status() + + data = res.json() + tool_calls = data["choices"][0]["message"]["tool_calls"] + + for call in tool_calls: + if call["function"]["name"] != TOOL_SCHEMA_SAVE_SUBTITLES["function"]["name"]: + continue + + parsed = json.loads(call["function"]["arguments"]) + translated_entries = parsed["entries"] + + if len(translated_entries) != len(chunk): + raise ValueError(f"Translation returned {len(translated_entries)} entries, but the chunk had {len(chunk)}.") + + # Reassemble subtitles with translated content + id_to_sub = {s.index: s for s in chunk} + result = [] + for entry in translated_entries: + original_sub = id_to_sub.get(entry["id"]) + if original_sub: + original_sub.content = entry["translated"].strip() + result.append(original_sub) + return result + + raise ValueError("The model did not return any translated subtitles via the specified tool.") + + +def translate_entries( + entries: list[srt.Subtitle], + model: str, + lang: str, + chunk_size: int = 30, + context_size: int = 4, +) -> list[srt.Subtitle]: + """ + Translates a list of subtitle entries using a sliding window for context. + """ + logger.info(f"Translating {len(entries)} entries to {lang}") + + if not entries: + return [] + + chunks = list(chunkify(entries, chunk_size)) + futures = [] + + for i, chunk in enumerate(chunks): + logger.info(f"Submitting chunk {i + 1}/{len(chunks)} for translation...") + + # Get context from the previous chunk + context_before = [] + if i > 0: + prev_chunk = chunks[i - 1] + context_before = prev_chunk[-context_size:] + + # Get context from the next chunk + context_after = [] + if i < len(chunks) - 1: + next_chunk = chunks[i + 1] + context_after = next_chunk[:context_size] + + # Submit the translation task to the thread pool + future = pool.submit( + _translate_chunk_with_context, + chunk=chunk, + context_before=context_before, + context_after=context_after, + lang=lang, + model=model, ) - payload = { - "model": "google/gemini-2.5-flash-lite", - "messages": [ - {"role": "system", "content": prompt}, - {"role": "user", "content": chunk_json}, - ], - "tools": [TOOL_SCHEMA_SAVE_SUBTITLES], - "tool_choice": {"type": "function", "function": {"name": TOOL_SCHEMA_SAVE_SUBTITLES["function"]["name"]}}, - } - res = client.post(url="/chat/completions", json=payload) - if res.is_error: - logging.error(f"Error translating subtitles: {res.text}") - res.raise_for_status() + futures.append(future) - data = res.json() + # Collect results as they complete + translated_subs = [] + for i, future in enumerate(futures): + logger.info(f"Processing result of chunk {i + 1}/{len(chunks)}...") + try: + translated_chunk = future.result() + translated_subs.extend(translated_chunk) + except Exception as e: + logger.error(f"Error translating chunk {i + 1}: {e}") + # Optionally, decide how to handle failed chunks, e.g., skip or retry. + # For now, we'll just log the error and continue. - tool_calls = data["choices"][0]["message"]["tool_calls"] + # Sort subtitles by start time to ensure correct order + translated_subs.sort(key=lambda s: s.start) - for call in tool_calls: - if call["function"]["name"] != TOOL_SCHEMA_SAVE_SUBTITLES["function"]["name"]: - continue - parsed = json.loads(call["function"]["arguments"]) - translated_entries = parsed["entries"] - if len(translated_entries) != len(chunk): - raise ValueError(f"Expected {len(chunk)} translated entries, got {len(translated_entries)}") - # Reassemble subtitles with translated content - id_to_sub = {s.index: s for s in chunk} - result = [] - for entry in translated_entries: - orig = id_to_sub.get(entry["id"]) - if orig: - orig.content = entry["translated"].strip() - result.append(orig) - return result - - combined = [] - chunks = chunkify(entries, 20) - for i, chunk in enumerate(chunks, start=1): - logging.info(f"Translating chunk {i}/{len(chunks)}") - translated = translate_chunk(chunk) - combined.extend(translated) - # Reindex - for i, sub in enumerate(combined, 1): + # Re-index all subtitles to ensure they are sequential + for i, sub in enumerate(translated_subs, 1): sub.index = i - return combined + + logger.info("Translation complete.") + return translated_subs -def translate(subtitle_path: Path, lang: str, save_path: Path, condense: int | None = None) -> None: +def translate( + subtitle_path: Path, + lang: str, + save_path: Path, + model: str, + condense: int | None = None, + chunk_size: int | None = None, +) -> None: """ Translate the subtitles in the given SRT file to the specified language. Returns a new SRT file with translated subtitles. """ - entries = list(srt.parse(subtitle_path.read_text())) - entries = clean_entries(entries) - entries = remove_overlaps(entries) + original = list(srt.parse(subtitle_path.read_text())) + cleaned = clean_entries(original) + without_overlaps = remove_overlaps(cleaned) + entries = without_overlaps if condense: - entries = combine_entries(entries, max_entries=condense) + condensed = combine_entries(entries, max_entries=condense) + save_path.with_name(f"{save_path.stem}.condensed.srt").write_text(srt.compose(condensed)) + logger.info(f"Condensed subtitles from {len(entries)} to {len(condensed)} entries") + entries = condensed if not entries: raise ValueError("No valid subtitle entries found to translate") - translated = translate_entries(entries, lang=lang) + translated = translate_entries( + entries=entries, + lang=lang, + model=model, + chunk_size=chunk_size, + ) + # wrap the translated version as well + if condense: + translated = combine_entries(entries=translated, max_entries=condense) save_path.write_text(srt.compose(translated)) +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description=""" +Translate subtitles from SRT file using OpenRouter API. +This script reads an SRT file, translates the subtitles to the specified language, +and saves the translated subtitles to a new SRT file. +""", + ) + p.add_argument("input", type=Path, help="Input SRT file") + p.add_argument("output", type=Path, help="Output SRT file") + p.add_argument("--translate", type=str, help="Language code for translation") + p.add_argument( + "--model", + type=str, + default="google/gemini-2.5-flash", + help="Model name for translation", + ) + p.add_argument( + "--condense", + type=int, + help="Merge adjacent short subtitles in groups of N if close together", + ) + p.add_argument( + "--chunk-size", + type=int, + default=30, + help="Number of subtitle entries to process in each chunk", + ) + return p.parse_args() + + def main(): - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") + logging.basicConfig(level=logging.INFO, format="%(name)s: %(asctime)s %(levelname)s: %(message)s") + logging.getLogger("httpx").setLevel(logging.WARNING) # Suppress httpx debug logs args = parse_args() translate( subtitle_path=args.input, lang=args.translate, save_path=args.output, + model=args.model, condense=args.condense, + chunk_size=args.chunk_size, )