From 375c67635a991c1cfda2fce0c034a9b8b57339f2 Mon Sep 17 00:00:00 2001
From: anandhu-eng <anandhukicks@gmail.com>
Date: Sun, 23 Jun 2024 16:47:56 +0530
Subject: [PATCH] Added model download script for mixtral

---
 script/get-ml-model-mixtral/README.md    |  1 +
 script/get-ml-model-mixtral/_cm.json     | 95 ++++++++++++++++++++++++
 script/get-ml-model-mixtral/customize.py | 24 ++++++
 3 files changed, 120 insertions(+)
 create mode 100644 script/get-ml-model-mixtral/README.md
 create mode 100644 script/get-ml-model-mixtral/_cm.json
 create mode 100644 script/get-ml-model-mixtral/customize.py

diff --git a/script/get-ml-model-mixtral/README.md b/script/get-ml-model-mixtral/README.md
new file mode 100644
index 000000000..4137062c1
--- /dev/null
+++ 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
new file mode 100644
index 000000000..f01e9d96b
--- /dev/null
+++ 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
new file mode 100644
index 000000000..8b1c0af52
--- /dev/null
+++ 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}