Merge pull request #1 from redhat-ai-services/download-flags

setup flags for download script

granite-3.0-2b-instruct/Containerfile

@@ -4,8 +4,8 @@ FROM quay.io/redhat-ai-services/huggingface-modelcar-builder:latest as base
 # The model repo to download
 ENV MODEL_REPO="ibm-granite/granite-3.0-2b-instruct"
 
-# Download the necessary model files (config.json, tokenizer.json, and safetensors)
-RUN python3 download_model.py
+# Download the necessary model files
+RUN python3 download_model.py --model-repo ${MODEL_REPO}
 
 # Final image containing only the essential model files
 FROM registry.access.redhat.com/ubi9/ubi-micro:9.4

granite-3.0-8b-instruct/Containerfile

@@ -4,8 +4,8 @@ FROM quay.io/redhat-ai-services/huggingface-modelcar-builder:latest as base
 # The model repo to download
 ENV MODEL_REPO="ibm-granite/granite-3.0-8b-instruct"
 
-# Download the necessary model files (config.json, tokenizer.json, and safetensors)
-RUN python3 download_model.py
+# Download the necessary model files
+RUN python3 download_model.py --model-repo ${MODEL_REPO}
 
 # Final image containing only the essential model files
 FROM registry.access.redhat.com/ubi9/ubi-micro:9.4

huggingface-modelcar-builder/README.md

@@ -1,3 +1,19 @@
 # Huggingface-modelcar-builder
 
 [![GitHub](https://img.shields.io/badge/GitHub-repo-blue.svg)](https://github.com/redhat-ai-services/modelcar-catalog/tree/main/huggingface-modelcar-builder) [![Quay.io](https://img.shields.io/badge/Quay.io-image-blue.svg)](https://quay.io/repository/redhat-ai-services/huggingface-modelcar-builder)
+
+## download_model
+
+```
+python download_model.py -h         
+usage: download_model.py [-h] [-m MODEL_REPO] [-t TARGET_DIR] [-a ALLOW_PATTERNS [ALLOW_PATTERNS ...]]
+
+options:
+  -h, --help            show this help message and exit
+  -m MODEL_REPO, --model-repo MODEL_REPO
+                        The model repo on huggingface
+  -t TARGET_DIR, --target-dir TARGET_DIR
+                        The target directory to download the model
+  -a ALLOW_PATTERNS [ALLOW_PATTERNS ...], --allow-patterns ALLOW_PATTERNS [ALLOW_PATTERNS ...]
+                        The allowed patterns to download
+```

huggingface-modelcar-builder/download_model.py

@@ -1,13 +1,48 @@
 import os
+from typing import List
+
+import argparse
 
 from huggingface_hub import snapshot_download
 
-model_repo = os.getenv("MODEL_REPO")
+def main(
+    model_repo: str,
+    local_dir: str = "./models",
+    allow_patterns: List[str] = ["*.safetensors", "*.json", "*.txt"],
+):
+    print(f"Attempting to download the following model from huggingface: {model_repo}")
+    print(f"Target director: {local_dir}")
+    print(f"With allow-patterns: {allow_patterns}")
+
+    snapshot_download(
+        repo_id=model_repo,
+        local_dir=local_dir,
+        allow_patterns=allow_patterns,
+    )
 
-print(f"Attempting to download the following model from huggingface: {model_repo}")
 
-snapshot_download(
-    repo_id=model_repo,
-    local_dir="/models",
-    allow_patterns=["*.safetensors", "*.json", "*.txt"],
-)
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-m", "--model-repo", help="The model repo on huggingface", type=str
+    )
+    parser.add_argument(
+        "-t",
+        "--target-dir",
+        help="The target directory to download the model",
+        default="./models",
+        type=str,
+    )
+    parser.add_argument(
+        "-a",
+        "--allow-patterns",
+        help="The allowed patterns to download",
+        nargs="+",
+        default=["*.safetensors", "*.json", "*.txt"],
+    )
+    args = parser.parse_args()
+    main(
+        model_repo=args.model_repo,
+        local_dir=args.target_dir,
+        allow_patterns=args.allow_patterns,
+    )

huggingface-modelcar-builder/pyproject.toml

@@ -1,4 +1,4 @@
 [project]
 name = "huggingface-modelcar-builder"
 description = "A container image to help build modelcar images"
-version = "0.1.0"
+version = "0.2.0"

huggingface-modelcar-builder/requirements.txt

@@ -1 +1,2 @@
+argparse
 huggingface-hub

mistral-7b-instruct-v0.3/Containerfile

@@ -8,7 +8,8 @@ ARG HF_TOKEN
 ENV MODEL_REPO="mistralai/Mistral-7B-Instruct-v0.3"
 
 # Download the necessary model files (config.json, tokenizer.json, and safetensors)
-RUN python3 download_model.py
+RUN python3 download_model.py --model-repo ${MODEL_REPO} \
+    --allow-patterns "params.json" "consolidated.safetensors" "tokenizer.model.v3"
 
 # Final image containing only the essential model files
 FROM registry.access.redhat.com/ubi9/ubi-micro:9.4