feat(subtitle-translator): add contextual chunk translation

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:
+    Merge adjacent subtitles in groups of up to 'max_entries' if:
-    - They are no more than 1s apart (start of next - end of prev <= 1s)
+    - The gap between the end of one and the start of the next is <= 1 second
-    - The total word count is fewer than 8
+    - 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
+    total_subs = len(entries)
-    n = len(entries)
+
-    while i < n:
+    while i < total_subs:
         group = [entries[i]]
-        j = 1
+        group_size = 1
-        while j < int(max_entries) and (i + j) < n:
+
-            prev = group[-1]
+        while group_size < max_entries and (i + group_size) < total_subs:
-            curr = entries[i + j]
+            last_sub = group[-1]
-            gap = (curr.start - prev.end).total_seconds()
+            next_sub = entries[i + group_size]
-            total_words = sum(len(s.content.split()) for s in group) + len(curr.content.split())
+            time_gap = (next_sub.start - last_sub.end).total_seconds()
-            if gap <= 1 and total_words < 8:
+            group_word_count = sum(len(sub.content.split()) for sub in group) + len(next_sub.content.split())
-                group.append(curr)
+
-                j += 1
+            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=""))
+        merged_subtitles.append(srt.Subtitle(index=current_index, start=start_time, end=end_time, content=formatted_text, proprietary=""))
-        idx += 1
+        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.
+    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.
-    Args:
+    If there's 1 word left for the last line, join it with the previous line.
-        words: List of words to format
-        max_words_per_line: Maximum words allowed per line
-
-    Returns:
-        Formatted subtitle text with line breaks
     """
     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
+    idx = 0
-
+    for i in range(num_lines):
-    while i < len(words):
+        # Distribute the remainder words to the first 'extra' lines
-        remaining = len(words) - i
+        count = base + (1 if i < extra else 0)
-
+        # If this is the last line and only 1 word left, join it with previous
-        # Handle orphan words (single word on last line)
+        if i == num_lines - 1 and count == 1 and lines:
-        if remaining <= 1 and lines:
+            lines[-1] += " " + words[idx]
-            lines[-1] += " " + words[i]
             break
-
+        lines.append(" ".join(words[idx : idx + count]))
-        # Calculate words for this line
+        idx += count
-        words_for_line = min(max_words_per_line, remaining)
-        lines.append(" ".join(words[i : i + words_for_line]))
-        i += words_for_line
 
     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]:
+@retry(max_retries=8)
-    logging.info(f"Translating subtitles to {lang} via OpenRouter")
+def _translate_chunk_with_context(
-
+    chunk: list[srt.Subtitle],
-    @retry(max_retries=5)
+    context_before: list[srt.Subtitle],
-    def translate_chunk(chunk: list[srt.Subtitle]) -> 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}.
+You will be provided with a JSON object containing three lists: 'context_before', 'main_chunk', and 'context_after'.
-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.
 
-Return exactly the same number of entries as provided. For each input entry, there must be one output entry with the same 'id'.
+Your task is to translate ONLY the subtitles in the 'main_chunk' from English to {lang}.
-If an entry is empty, return it as empty.
+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(
+    # Create the structured payload for the API
-            [
+    payload_data = {
-                {
+        "context_before": [{"id": s.index, "content": s.content.strip()} for s in context_before],
-                    "id": s.index,
+        "main_chunk": [{"id": s.index, "content": s.content.strip()} for s in chunk],
-                    "content": s.content.strip(),
+        "context_after": [{"id": s.index, "content": s.content.strip()} for s in context_after],
     }
-                for s in chunk
+
-            ],
-            ensure_ascii=False,
-        )
     payload = {
-            "model": "google/gemini-2.5-flash-lite",
+        "model": model,
         "messages": [
             {"role": "system", "content": prompt},
-                {"role": "user", "content": chunk_json},
+            {"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:
-            logging.error(f"Error translating subtitles: {res.text}")
+        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"Expected {len(chunk)} translated entries, got {len(translated_entries)}")
+            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:
-                orig = id_to_sub.get(entry["id"])
+            original_sub = id_to_sub.get(entry["id"])
-                if orig:
+            if original_sub:
-                    orig.content = entry["translated"].strip()
+                original_sub.content = entry["translated"].strip()
-                    result.append(orig)
+                result.append(original_sub)
         return result
 
-    combined = []
+    raise ValueError("The model did not return any translated subtitles via the specified tool.")
-    chunks = chunkify(entries, 20)
+
-    for i, chunk in enumerate(chunks, start=1):
+
-        logging.info(f"Translating chunk {i}/{len(chunks)}")
+def translate_entries(
-        translated = translate_chunk(chunk)
+    entries: list[srt.Subtitle],
-        combined.extend(translated)
+    model: str,
-    # Reindex
+    lang: str,
-    for i, sub in enumerate(combined, 1):
+    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,
+        )
+        futures.append(future)
+
+    # 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.
+
+    # Sort subtitles by start time to ensure correct order
+    translated_subs.sort(key=lambda s: s.start)
+
+    # 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()))
+    original = list(srt.parse(subtitle_path.read_text()))
-    entries = clean_entries(entries)
+    cleaned = clean_entries(original)
-    entries = remove_overlaps(entries)
+    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,
     )