Translate MM Tipitaka translation into English using Deepseek v4 flash via NIM

mm_tipitaka_translate.py

"""
Tipitaka Myanmar → English Translation Pipeline
Uses DeepSeek via NVIDIA NIM API to translate HTML files.

Usage:
    export NVIDIA_API_KEY="your_api_key_here"
    python translate.py                        # translate all files
    python translate.py --file 01_vinaya_01    # translate one file
    python translate.py --resume               # resume from last checkpoint
    python translate.py --status               # show progress status

On VPS:

    nohup python translate.py > nohup.log 2>&1 &

Monitor it later (when you log back in)

# See if it's still running
ps aux | grep translate.py

# Watch live output
tail -f nohup.log

# Or check the translation log
tail -f translation.log

# Check progress status
python translate.py --status

Stop it if needed
# Find the process ID
ps aux | grep translate.py

# Kill it (replace 12345 with actual PID)
kill 12345


"""

import os
import re
import json
import time
import argparse
import logging
from pathlib import Path
from copy import copy
from typing import Optional

from openai import OpenAI
from bs4 import BeautifulSoup, NavigableString, Tag

# ─── Configuration ────────────────────────────────────────────────────────────

BASE_DIR        = Path(__file__).parent
INPUT_DIR       = BASE_DIR / "assets" / "books"
OUTPUT_DIR      = BASE_DIR / "assets" / "books_en"
PROGRESS_FILE   = BASE_DIR / "translation_progress.json"
LOG_FILE        = BASE_DIR / "translation.log"

NVIDIA_API_KEY  = os.environ.get("NVIDIA_API_KEY", "")
MODEL           = "deepseek-ai/deepseek-v4-flash"

# How many HTML text nodes to batch into one API call (balance cost vs speed)
BATCH_SIZE      = 20        # number of paragraphs per API request
MAX_TOKENS      = 16384
TEMPERATURE     = 1
TOP_P           = 0.95

# Retry settings
MAX_RETRIES     = 5
RETRY_DELAY     = 10        # seconds between retries
RATE_LIMIT_DELAY = 2        # seconds between API calls

# Tags whose text content will be translated
TRANSLATABLE_TAGS = {"p", "h1", "h2", "h3", "h4", "title", "span"}

# ─── Logging ──────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(LOG_FILE, encoding="utf-8"),
        logging.StreamHandler(),
    ],
)
log = logging.getLogger(__name__)

# ─── NVIDIA / DeepSeek Client ─────────────────────────────────────────────────

def make_client() -> OpenAI:
    if not NVIDIA_API_KEY:
        raise ValueError(
            "NVIDIA_API_KEY environment variable is not set.\n"
            "Run:  export NVIDIA_API_KEY='your_key_here'"
        )
    return OpenAI(
        base_url="https://integrate.api.nvidia.com/v1",
        api_key=NVIDIA_API_KEY,
    )


# ─── Progress Tracking ────────────────────────────────────────────────────────

def load_progress() -> dict:
    if PROGRESS_FILE.exists():
        with open(PROGRESS_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}


def save_progress(progress: dict):
    with open(PROGRESS_FILE, "w", encoding="utf-8") as f:
        json.dump(progress, f, indent=2, ensure_ascii=False)


# ─── Translation ──────────────────────────────────────────────────────────────

SYSTEM_PROMPT = """You are an expert translator specializing in Theravada Buddhist texts.
Translate the following Myanmar (Burmese) Buddhist scripture into clear, accurate English.

Rules:
- Preserve Pali technical terms (e.g., Dhamma, Sangha, Nibbana, Bhikkhu, sutta names) — keep them as-is or transliterate
- Keep paragraph numbers like ၁, ၂, ၃ etc. — translate them to 1, 2, 3
- Preserve the meaning and reverent tone of the original
- Do NOT add explanations or commentary — translate only
- Return ONLY the translations, one per line, in the same order as input
- Use |||SEPARATOR||| between each translated item

Input format: numbered items separated by |||SEPARATOR|||
Output format: same numbered translations separated by |||SEPARATOR|||"""


def translate_batch(client: OpenAI, texts: list[str]) -> list[str]:
    """Translate a batch of text strings using DeepSeek. Returns list of translations."""
    if not texts:
        return []

    # Build numbered input
    numbered_input = "|||SEPARATOR|||".join(
        f"[{i+1}] {t}" for i, t in enumerate(texts)
    )

    prompt = (
        f"Translate each of the following {len(texts)} Myanmar Buddhist text segments "
        f"to English. Return them separated by |||SEPARATOR||| in the same order:\n\n"
        f"{numbered_input}"
    )

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            collected = []
            completion = client.chat.completions.create(
                model=MODEL,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user",   "content": prompt},
                ],
                temperature=TEMPERATURE,
                top_p=TOP_P,
                max_tokens=MAX_TOKENS,
                extra_body={"chat_template_kwargs": {"thinking": False}},
                stream=True,
            )

            for chunk in completion:
                if not getattr(chunk, "choices", None):
                    continue
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    collected.append(delta.content)

            full_response = "".join(collected).strip()
            translations = _parse_batch_response(full_response, len(texts))

            if len(translations) == len(texts):
                time.sleep(RATE_LIMIT_DELAY)
                return translations

            log.warning(
                f"Got {len(translations)} translations for {len(texts)} inputs "
                f"(attempt {attempt}). Retrying..."
            )

        except Exception as e:
            log.error(f"API error on attempt {attempt}/{MAX_RETRIES}: {e}")
            if attempt < MAX_RETRIES:
                log.info(f"Waiting {RETRY_DELAY}s before retry...")
                time.sleep(RETRY_DELAY * attempt)
            else:
                log.error("Max retries reached. Returning originals as fallback.")
                return texts  # fallback: keep original text

    return texts  # fallback


def _parse_batch_response(response: str, expected: int) -> list[str]:
    """Parse the |||SEPARATOR|||-delimited response into a list."""
    parts = response.split("|||SEPARATOR|||")
    cleaned = []
    for part in parts:
        part = part.strip()
        # Remove leading [N] numbering if model echoed it
        part = re.sub(r"^\[\d+\]\s*", "", part).strip()
        if part:
            cleaned.append(part)
    return cleaned


# ─── HTML Processing ──────────────────────────────────────────────────────────

def collect_text_nodes(soup: BeautifulSoup) -> list[tuple[Tag, str]]:
    """
    Walk the soup tree and collect (tag, original_text) for all
    translatable leaf-text elements that contain Myanmar characters.
    """
    nodes = []
    myanmar_pattern = re.compile(r"[\u1000-\u109F\uA9E0-\uA9FF\uAA60-\uAA7F]")

    for tag in soup.find_all(TRANSLATABLE_TAGS):
        # Only direct text (not nested tags) to avoid double-processing
        direct_text = "".join(
            str(c) for c in tag.children if isinstance(c, NavigableString)
        ).strip()

        if direct_text and myanmar_pattern.search(direct_text):
            nodes.append((tag, direct_text))

    return nodes


def apply_translations(nodes: list[tuple[Tag, str]], translations: list[str]):
    """Replace original text in-place within each tag."""
    for (tag, original), translated in zip(nodes, translations):
        # Replace NavigableString children with translated text,
        # preserving any child tags (e.g., <span class="paragraph">)
        new_strings = []
        for child in list(tag.children):
            if isinstance(child, NavigableString):
                new_strings.append((child, translated))
                break  # replace only the first/main text node

        for string_node, new_text in new_strings:
            string_node.replace_with(new_text)


def translate_html_file(
    client: OpenAI,
    input_path: Path,
    output_path: Path,
    progress: dict,
    file_key: str,
) -> bool:
    """
    Translate a single HTML file. Returns True on success.
    Supports resuming from last completed batch.
    """
    log.info(f"Processing: {input_path.name}")

    with open(input_path, "r", encoding="utf-8") as f:
        soup = BeautifulSoup(f.read(), "html.parser")

    nodes = collect_text_nodes(soup)
    total = len(nodes)
    log.info(f"  Found {total} translatable text nodes")

    if total == 0:
        log.info("  No Myanmar text found, copying as-is.")
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(str(soup))
        return True

    # Resume from checkpoint
    file_progress = progress.get(file_key, {"completed_batches": 0, "translations": {}})
    completed_batches = file_progress.get("completed_batches", 0)
    cached_translations: dict[int, str] = {
        int(k): v for k, v in file_progress.get("translations", {}).items()
    }

    # Build batches
    batches: list[list[tuple[int, str]]] = []  # [(node_index, text), ...]
    for i in range(0, total, BATCH_SIZE):
        batch = [(j, nodes[j][1]) for j in range(i, min(i + BATCH_SIZE, total))]
        batches.append(batch)

    log.info(f"  {len(batches)} batches total, {completed_batches} already done")

    # Translate remaining batches
    for batch_idx, batch in enumerate(batches):
        if batch_idx < completed_batches:
            continue  # already translated

        indices = [item[0] for item in batch]
        texts   = [item[1] for item in batch]

        log.info(f"  Batch {batch_idx+1}/{len(batches)} ({len(texts)} items)...")
        translated = translate_batch(client, texts)

        # Cache results
        for idx, trans in zip(indices, translated):
            cached_translations[idx] = trans

        # Save checkpoint after every batch
        file_progress["completed_batches"] = batch_idx + 1
        file_progress["translations"] = {str(k): v for k, v in cached_translations.items()}
        progress[file_key] = file_progress
        save_progress(progress)

        log.info(f"  ✓ Batch {batch_idx+1} done")

    # Apply all translations to soup
    all_translations = [cached_translations.get(i, nodes[i][1]) for i in range(total)]
    apply_translations(nodes, all_translations)

    # Write output
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(str(soup))

    log.info(f"  ✓ Saved to: {output_path}")

    # Mark file as fully complete
    file_progress["done"] = True
    progress[file_key] = file_progress
    save_progress(progress)

    return True


# ─── CLI ──────────────────────────────────────────────────────────────────────

def show_status(progress: dict):
    """Print a status table of all files."""
    all_files = sorted(INPUT_DIR.glob("*.html"))
    print(f"\n{'File':<35} {'Status':<12} {'Batches Done'}")
    print("-" * 65)
    for f in all_files:
        key = f.stem
        p = progress.get(key, {})
        done = p.get("done", False)
        batches = p.get("completed_batches", 0)
        status = "✅ done" if done else ("🔄 partial" if batches > 0 else "⏳ pending")
        print(f"{f.name:<35} {status:<12} {batches}")
    print()


def main():
    parser = argparse.ArgumentParser(description="Translate Myanmar Buddhist HTML to English")
    parser.add_argument("--file",   help="Translate a single file (stem name, e.g. 01_vinaya_01)")
    parser.add_argument("--resume", action="store_true", help="Resume from last checkpoint")
    parser.add_argument("--status", action="store_true", help="Show translation progress")
    parser.add_argument("--reset",  help="Reset progress for a specific file stem")
    args = parser.parse_args()

    progress = load_progress()

    if args.status:
        show_status(progress)
        return

    if args.reset:
        if args.reset in progress:
            del progress[args.reset]
            save_progress(progress)
            log.info(f"Progress reset for: {args.reset}")
        else:
            log.info(f"No progress found for: {args.reset}")
        return

    client = make_client()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    if args.file:
        # Single file mode
        input_path  = INPUT_DIR / f"{args.file}.html"
        output_path = OUTPUT_DIR / f"{args.file}.html"
        if not input_path.exists():
            log.error(f"File not found: {input_path}")
            return
        file_key = args.file
        if progress.get(file_key, {}).get("done") and not args.resume:
            log.info(f"{args.file} is already fully translated. Use --resume to redo.")
            return
        translate_html_file(client, input_path, output_path, progress, file_key)

    else:
        # All files mode
        all_files = sorted(INPUT_DIR.glob("*.html"))
        log.info(f"Found {len(all_files)} HTML files in {INPUT_DIR}")

        success = 0
        skipped = 0
        failed  = 0

        for input_path in all_files:
            file_key    = input_path.stem
            output_path = OUTPUT_DIR / input_path.name

            # Skip already completed files (unless --resume)
            if progress.get(file_key, {}).get("done") and not args.resume:
                log.info(f"Skipping {input_path.name} (already done)")
                skipped += 1
                continue

            try:
                ok = translate_html_file(client, input_path, output_path, progress, file_key)
                if ok:
                    success += 1
                else:
                    failed += 1
            except KeyboardInterrupt:
                log.info("\n⚠️  Interrupted. Progress saved. Run again to resume.")
                break
            except Exception as e:
                log.error(f"Failed to process {input_path.name}: {e}")
                failed += 1
                continue

        log.info(f"\n{'='*50}")
        log.info(f"Done! ✅ {success} translated | ⏭ {skipped} skipped | ❌ {failed} failed")
        log.info(f"Output directory: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()