translate.py

"""Stage-3: Translate English fine-tune JSONL to Mandarin and Arabic via Ollama.

Pipeline position: Data Recipes -> validate.py -> translate.py -> Kaggle upload.
Model choice (D-11): gemma4:26b stock — strict has num_predict=512 cap, ocr has
temperature=0.1 tuned for verbatim extraction, compiler is fine-tuned for pack
compilation. Stock gemma4:26b has 140-language training and no output-length cap.

Translation scope (D-09 / Pitfall 5): only the `output` field is translated.
`instruction` and `input` remain in English so a single English query taxonomy
can serve all target languages.

D-10: agent.py:translate_chunk is intentionally NOT reused — that function takes
raw pack chunks, not {instruction, input, output} fine-tune records.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

from openai import OpenAI

TRANSLATE_MODEL = "gemma4:26b"  # D-11
OLLAMA_BASE_URL = "http://localhost:11434/v1"

LANGUAGES: dict[str, str] = {
    "zh": "Simplified Chinese (Mandarin)",
    "ar": "Modern Standard Arabic",
}

TRANSLATE_PROMPT = (
    "You are a professional translator specialising in plain-language safety guidance. "
    "Translate the following English text into {lang_name}. "
    "Preserve numbered lists, bullet structure, and proper nouns. "
    "Return ONLY the translated text — no explanation, no preamble, no quotation marks.\n\n"
    "TEXT:\n{text}"
)

client = OpenAI(base_url=OLLAMA_BASE_URL, api_key="ollama")


def _validate_path(path: str) -> str:
    """Validate that the path is within the current working directory.

    This addresses CR-04 (Path Traversal).
    """
    root = Path.cwd().resolve()
    candidate = Path(path).resolve()
    try:
        candidate.relative_to(root)
    except ValueError as exc:
        raise ValueError(
            f"Path traversal detected: {path} is outside allowed directory."
        ) from exc
    return str(candidate)


def read_jsonl(path: Path) -> list[dict]:
    """Read a JSONL file and return a list of dicts."""
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def translate_record(record: dict, target_lang: str) -> dict:
    """Translate the `output` field of a fine-tune record to the target language.

    Per D-09: only `output` is translated. `instruction` and `input` remain in English
    so a single English query taxonomy serves all target languages.

    Per D-10: does NOT call agent.py:translate_chunk (different input contract).
    Per D-11: calls gemma4:26b on Ollama (stock — not strict/ocr/compiler variants).

    Returns dict with same instruction+input, translated output, and language tag.
    """
    if target_lang not in LANGUAGES:
        raise ValueError(f"Unsupported target language: {target_lang}. Known: {list(LANGUAGES)}")
    required_fields = ("instruction", "input", "output")
    missing = [f for f in required_fields if f not in record]
    if missing:
        raise ValueError(f"Record missing required fields: {missing}")
    lang_name = LANGUAGES[target_lang]
    output_text = record["output"]
    response = client.chat.completions.create(
        model=TRANSLATE_MODEL,
        messages=[{
            "role": "user",
            "content": TRANSLATE_PROMPT.format(lang_name=lang_name, text=output_text),
        }],
        temperature=0.3,
    )
    translated = response.choices[0].message.content.strip()
    return {
        "instruction": record["instruction"],
        "input": record["input"],
        "output": translated,
        "language": target_lang,
    }


def translate_file(input_jsonl: Path, target_lang: str, output_jsonl: Path) -> None:
    """Translate all records in input_jsonl to target_lang and write to output_jsonl.

    Per D-09: reads English JSONL, translates `output` field only, writes per-language JSONL.
    Per D-11: uses gemma4:26b for translation.
    """
    records = read_jsonl(input_jsonl)
    print(f"[translate] {target_lang}  source={input_jsonl.name}  records={len(records)}", flush=True)
    translated: list[dict] = []
    for i, record in enumerate(records):
        t0 = time.monotonic()
        try:
            translated.append(translate_record(record, target_lang))
        except Exception as exc:
            raise RuntimeError(
                f"[translate] {target_lang} record {i} failed: {exc}"
            ) from exc
        print(
            f"[translate] {target_lang} {i+1}/{len(records)}  "
            f"{time.monotonic() - t0:.1f}s",
            flush=True,
        )
    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
    output_jsonl.write_text(
        "\n".join(json.dumps(r, ensure_ascii=False) for r in translated) + "\n",
        encoding="utf-8",
    )
    print(f"[translate] done  {target_lang}  -> {output_jsonl}  {len(translated)} records", flush=True)


def merge_jsonl(files: list[Path], output: Path) -> None:
    """Merge multiple JSONL files into a single file preserving all language tags.

    Per D-04: merged multilingual JSONL preserves all language tags.
    """
    all_records: list[dict] = []
    for f in files:
        all_records += [json.loads(line) for line in f.read_text(encoding="utf-8").splitlines() if line.strip()]
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(
        "\n".join(json.dumps(r, ensure_ascii=False) for r in all_records) + "\n",
        encoding="utf-8",
    )
    print(f"[translate] merged  {len(all_records)} records  -> {output}", flush=True)


def main() -> None:
    """CLI entry point for translate.py.

    Usage:
        uv run --extra agent python translate.py \\
            --input build/train_en.jsonl \\
            --output-dir build \\
            --langs zh ar [--merge]
    """
    parser = argparse.ArgumentParser(
        description="Stage-3: Translate English fine-tune JSONL to Mandarin and Arabic via Ollama (D-09, D-11)"
    )
    parser.add_argument("--input", required=True, help="Path to English JSONL file (train_en.jsonl)")
    parser.add_argument("--output-dir", required=True, help="Directory for per-language JSONL outputs")
    parser.add_argument(
        "--langs",
        nargs="+",
        default=["zh", "ar"],
        help="Target language codes (default: zh ar)",
    )
    parser.add_argument(
        "--merge",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Emit train_multilingual.jsonl combining English + all translated languages (use --no-merge to skip)",
    )
    args = parser.parse_args()

    try:
        input_path = Path(_validate_path(args.input))
        output_dir = Path(_validate_path(args.output_dir))
    except ValueError as e:
        print(f"[translate] FATAL: {e}", flush=True)
        sys.exit(1)

    try:
        output_dir.mkdir(parents=True, exist_ok=True)

        per_lang_outputs: list[Path] = []
        for lang in args.langs:
            out_path = output_dir / f"train_{lang}.jsonl"
            translate_file(input_path, lang, out_path)
            per_lang_outputs.append(out_path)

        if args.merge:
            merge_jsonl([input_path, *per_lang_outputs], output_dir / "train_multilingual.jsonl")

        print("[translate] all done", flush=True)

    except Exception as e:
        print(f"[translate] FATAL: {e}", flush=True)
        sys.exit(1)


if __name__ == "__main__":
    main()