deploy/bake.py

#!/usr/bin/env python3
"""Pre-download every model the server needs, at image build time.

Run from the Dockerfile, never at boot. See the comment above the ARG in the
Dockerfile for why.

  python deploy/bake.py hi,bn,kn,mr
"""
import sys

# ISO-2 -> ISO-3, duplicated from src/voice.py on purpose: this runs BEFORE
# src/ is copied into the image, so it cannot import the real map. The four
# codes that are not a naive 3-letter truncation are the ones worth checking
# against src/voice.py if this list is ever edited.
MMS_LANG = {"as": "asm", "bn": "ben", "gu": "guj", "hi": "hin", "kn": "kan",
            "ml": "mal", "mr": "mar", "ne": "npi", "or": "ory", "pa": "pan",
            "sa": "san", "ta": "tam", "te": "tel", "ur": "urd", "en": "eng"}

# Embedder (src/evaluate_retrieval.py) uses AutoModel's dense CLS vector only,
# so the onnx/ export and the colbert and sparse heads are dead weight -- about
# 2 GB of it.
BGE_FILES = ["config.json", "model.safetensors", "tokenizer.json",
             "tokenizer_config.json", "sentencepiece.bpe.model",
             "special_tokens_map.json"]

# Neither has an MMS voice; TTS falls back and the text path is unaffected.
NO_VOICE = {"ne", "sa"}


def main(langs: list[str]) -> int:
    from huggingface_hub import snapshot_download

    snapshot_download("BAAI/bge-m3", allow_patterns=BGE_FILES)
    print("baked BAAI/bge-m3")

    for lg in langs:
        if lg in NO_VOICE:
            print(f"skipped {lg}: no MMS voice exists")
            continue
        iso3 = MMS_LANG.get(lg)
        if not iso3:
            print(f"skipped {lg}: not an MSMARCO-XI language")
            continue
        snapshot_download(f"facebook/mms-tts-{iso3}")
        print(f"baked facebook/mms-tts-{iso3} ({lg})")
    return 0


if __name__ == "__main__":
    arg = sys.argv[1] if len(sys.argv) > 1 else "hi,bn,kn,mr"
    raise SystemExit(main([x.strip() for x in arg.split(",") if x.strip()]))