Cache only specific files from SeamlessM4T

Dockerfile

@@ -36,9 +36,10 @@ ENV HOME=/home/user \
 WORKDIR ${HOME}/app
 
 COPY --chown=1000 . ${HOME}/app
+# This will cache the model into the docker image
+RUN ./cache.sh
 RUN pip install -r ${HOME}/app/requirements.txt && \
     pip install fairseq2 --pre --extra-index-url https://fair.pkg.atmeta.com/fairseq2/pt2.1.0/cu121 && \
     pip install ${HOME}/app/whl/seamless_communication-1.0.0-py3-none-any.whl
-# This will cache the model into the docker image
-RUN python -u translator.py
+
 CMD ["python", "main.py"]

cache.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+mkdir -p /home/user/app/models
+wget https://huggingface.co/facebook/seamless-m4t-v2-large/resolve/main/spm_char_lang38_tc.model -P /home/user/app/models
+wget https://huggingface.co/facebook/seamless-m4t-v2-large/resolve/main/seamlessM4T_v2_large.pt -P /home/user/app/models
+wget https://huggingface.co/facebook/seamless-m4t-large/resolve/main/tokenizer.model -P /home/user/app/models
+wget https://dl.fbaipublicfiles.com/seamless/models/vocoder_v2.pt -P /home/user/app/models

translator.py

@@ -2,20 +2,33 @@ import os
 import pathlib
 import torch
 from fairseq2.assets import InProcAssetMetadataProvider, asset_store
-from huggingface_hub import snapshot_download
 from seamless_communication.inference import Translator
 
 CHECKPOINTS_PATH = pathlib.Path(os.getenv("CHECKPOINTS_PATH", "/home/user/app/models"))
 if not CHECKPOINTS_PATH.exists():
-    snapshot_download(repo_id="facebook/seamless-m4t-v2-large", repo_type="model", local_dir=CHECKPOINTS_PATH)
+    # from huggingface_hub import snapshot_download
+    # snapshot_download(repo_id="facebook/seamless-m4t-v2-large", repo_type="model", local_dir=CHECKPOINTS_PATH)
+    raise FileNotFoundError(f"Checkpoint path {CHECKPOINTS_PATH} does not exist")
+
 asset_store.env_resolvers.clear()
 asset_store.env_resolvers.append(lambda: "demo")
 demo_metadata = [
+    # https://github.com/facebookresearch/seamless_communication/blob/dd67e71317d66752ef16cf21bd842ca3273244c9/src/seamless_communication/cards/seamlessM4T_v2_large.yaml#L10
+    # char_tokenizer: "https://huggingface.co/facebook/seamless-m4t-v2-large/resolve/main/spm_char_lang38_tc.model"
+    # checkpoint: "https://huggingface.co/facebook/seamless-m4t-v2-large/resolve/main/seamlessM4T_v2_large.pt"
     {
         "name": "seamlessM4T_v2_large@demo",
         "checkpoint": f"file://{CHECKPOINTS_PATH}/seamlessM4T_v2_large.pt",
         "char_tokenizer": f"file://{CHECKPOINTS_PATH}/spm_char_lang38_tc.model",
     },
+    # https://github.com/facebookresearch/seamless_communication/blob/dd67e71317d66752ef16cf21bd842ca3273244c9/src/seamless_communication/cards/unity_nllb-100.yaml#L9C1-L9C93
+    # tokenizer: "https://huggingface.co/facebook/seamless-m4t-large/resolve/main/tokenizer.model"
+    {
+        "name": "unity_nllb-100@demo",
+        "tokenizer": f"file://{CHECKPOINTS_PATH}/tokenizer.model",
+    },
+    # https://github.com/facebookresearch/seamless_communication/blob/dd67e71317d66752ef16cf21bd842ca3273244c9/src/seamless_communication/cards/vocoder_v2.yaml#L10
+    # checkpoint: "https://dl.fbaipublicfiles.com/seamless/models/vocoder_v2.pt"
     {
         "name": "vocoder_v2@demo",
         "checkpoint": f"file://{CHECKPOINTS_PATH}/vocoder_v2.pt",
@@ -37,3 +50,16 @@ translator = Translator(
     dtype=dtype,
     apply_mintox=True,
 )
+
+if __name__ == '__main__':
+    input_text = "Hello, how are you today?"
+    source_language_code = "eng"
+    target_language_code = "zsm"
+
+    result = translator.predict(
+        input=input_text,
+        task_str="T2TT",
+        src_lang=source_language_code,
+        tgt_lang=target_language_code,
+    )
+    print(str(result[0]))