Merge pull request #5 from backend-developers-ltd/mock-image

Allow building mock image
backend-developers-ltd · Sep 20, 2024 · 4e333fe · 4e333fe
2 parents 13b8eb0 + d1cbabd
commit 4e333fe
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/.github/workflows/build_push_image.yml b/.github/workflows/build_push_image.yml
@@ -5,6 +5,7 @@ on:
     branches:
       - build-push-llama3-image
       - build-push-phi3-image
+      - build-push-mock-image
   workflow_dispatch:
 
 env:
@@ -39,6 +40,8 @@ jobs:
             echo "MODEL_NAME=llama3" >> $GITHUB_ENV
           elif [[ "${{ github.ref }}" == "refs/heads/build-push-phi3-image" ]]; then
             echo "MODEL_NAME=phi3" >> $GITHUB_ENV
+          elif [[ "${{ github.ref }}" == "refs/heads/build-push-mock-image" ]]; then
+            echo "MODEL_NAME=mock" >> $GITHUB_ENV
           fi
 
       - name: Docker build and push

diff --git a/src/compute_horde_prompt_gen/download_model.py b/src/compute_horde_prompt_gen/download_model.py
@@ -1,5 +1,6 @@
 import os
 import argparse
+import sys
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -39,6 +40,10 @@
     )
 
     args = parser.parse_args()
+
+    if args.model_name == "mock":
+        sys.exit(0)
+
     save_path = os.path.join(args.save_path, args.model_name)
     model_name = MODEL_PATHS[args.model_name]
     print(f"Saving {model_name} model to {save_path}")