convert_model.py

"""Fetch a pre-converted CTranslate2 NLLB model from the HF Hub.

The HF Space builder OOMs trying to convert NLLB-200-3.3B in place
(loading FP32 weights to quantize requires ~24 GB RAM, builders cap
around 16 GB), so conversion happens once on a beefier local machine
and the converted artifact is hosted on the Hub. At build time we
just download it.

Override the source repo with NLLB_CT2_REPO (default points at the
Napron-hosted conversion). Output dir is NLLB_CT2_DIR.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

DEFAULT_REPO = "Napron/nllb-200-3.3B-ct2-int8"

SOURCE_REPO = os.environ.get("NLLB_CT2_REPO", DEFAULT_REPO)
OUTPUT_DIR = Path(os.environ.get("NLLB_CT2_DIR", "/home/user/app/model"))


def main() -> int:
    from huggingface_hub import snapshot_download

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

    if (OUTPUT_DIR / "model.bin").exists():
        print(f"[fetch] CT2 model already present at {OUTPUT_DIR}, skipping.")
        return 0

    print(f"[fetch] downloading pre-converted CT2 model from {SOURCE_REPO} → {OUTPUT_DIR}")
    snapshot_download(
        repo_id=SOURCE_REPO,
        repo_type="model",
        local_dir=str(OUTPUT_DIR),
        token=os.environ.get("HF_TOKEN"),
    )
    print(f"[fetch] done — model at {OUTPUT_DIR}")
    return 0


if __name__ == "__main__":
    sys.exit(main())