Added model download script for mixtral

anandhu-eng · Jun 23, 2024 · 375c676 · 375c676
1 parent 72265bc
commit 375c676
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 0 deletions.
diff --git a/script/get-ml-model-mixtral/README.md b/script/get-ml-model-mixtral/README.md
@@ -0,0 +1 @@
+Please see [https://docs.mlcommons.org/cm4mlops/scripts/AI-ML-models/get-ml-model-llama2](https://docs.mlcommons.org/cm4mlops/scripts/AI-ML-models/get-ml-model-llama2) for the documentation of this CM script.
diff --git a/script/get-ml-model-mixtral/_cm.json b/script/get-ml-model-mixtral/_cm.json
@@ -0,0 +1,95 @@
+{
+  "alias": "get-ml-model-mixtral",
+  "automation_alias": "script",
+  "automation_uid": "5b4e0237da074764",
+  "cache": true,
+  "category": "AI/ML models",
+  "env": {
+    "CM_ML_MODEL_DATASET": "",
+    "CM_ML_MODEL_WEIGHT_TRANSFORMATIONS": "no"
+  },
+  "input_mapping": {
+    "checkpoint": "MIXTRAL_CHECKPOINT_PATH"
+  },
+  "new_env_keys": [
+    "CM_ML_MODEL_*",
+    "MIXTRAL_CHECKPOINT_PATH"
+  ],
+  "prehook_deps": [
+    {
+      "enable_if_env": {
+        "CM_TMP_REQUIRE_DOWNLOAD": [
+          "yes"
+        ]
+      },
+      "env": {},
+      "extra_cache_tags": "mixtral",
+      "force_env_keys": [
+        "CM_GIT_CHECKOUT_FOLDER"
+      ],
+      "names": [
+        "hf-zoo"
+      ],
+      "tags": "get,ml-model,huggingface,zoo,_clone-repo"
+    }
+  ],
+  "print_env_at_the_end": {
+    "MIXTRAL_CHECKPOINT_PATH": "MIXTRAL checkpoint path"
+  },
+  "tags": [
+    "get",
+    "raw",
+    "ml-model",
+    "language-processing",
+    "mixtral",
+    "mixtral-7b"
+  ],
+  "uid": "0c14127677f34ea2",
+  "variations": {
+    "batch_size.#": {
+      "env": {
+        "CM_ML_MODEL_BATCH_SIZE": "#"
+      }
+    },
+    "fp32": {
+      "default": true,
+      "env": {
+        "CM_ML_MODEL_INPUT_DATA_TYPES": "fp32",
+        "CM_ML_MODEL_PRECISION": "fp32",
+        "CM_ML_MODEL_WEIGHT_DATA_TYPES": "fp32"
+      },
+      "group": "precision"
+    },
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": {
+      "adr": {
+        "hf-zoo": {
+          "tags": "_model-stub.mistralai/Mixtral-8x7B-Instruct-v0.1"
+        }
+      },
+      "default": true,
+      "env": {
+        "CM_GIT_CHECKOUT_FOLDER": "Mixtral-8x7B-Instruct-v0.1",
+        "CM_MODEL_ZOO_ENV_KEY": "MIXTRAL"
+      },
+      "group": "huggingface-stub"
+    },
+    "pytorch": {
+      "default": true,
+      "env": {
+        "CM_ML_MODEL_FRAMEWORK": "pytorch"
+      },
+      "group": "framework"
+    },
+    "stub.#": {
+      "adr": {
+        "hf-zoo": {
+          "tags": "_model-stub.#"
+        }
+      },
+      "env": {
+        "CM_MODEL_ZOO_ENV_KEY": "MIXTRAL"
+      },
+      "group": "huggingface-stub"
+    }
+  }
+}
diff --git a/script/get-ml-model-mixtral/customize.py b/script/get-ml-model-mixtral/customize.py
@@ -0,0 +1,24 @@
+from cmind import utils
+import os
+
+def preprocess(i):
+
+    os_info = i['os_info']
+    env = i['env']
+
+    path = env.get('MIXTRAL_CHECKPOINT_PATH', '').strip()
+
+    if path == '' or not os.path.exists(path):
+        env['CM_TMP_REQUIRE_DOWNLOAD'] = 'yes'
+
+    return {'return':0}
+
+def postprocess(i):
+
+    env = i['env']
+
+    env['MIXTRAL_CHECKPOINT_PATH'] = env['CM_ML_MODEL_PATH']
+    env['CM_ML_MODEL_PATH'] = env['MIXTRAL_CHECKPOINT_PATH']
+    env['CM_GET_DEPENDENT_CACHED_PATH'] = env['CM_ML_MODEL_PATH']
+
+    return {'return':0}