fix

Author: mnoukhov

mason.py

@@ -477,14 +477,21 @@ def make_internal_command(command: list[str], args: argparse.Namespace, whoami:
                 try:
                     model_arg_idx = filtered_command.index("--model_name_or_path")
                     model_name_idx = model_arg_idx + 1
-                    model_name_or_path = filtered_command[model_name_idx]
+                    model_name_or_path = filtered_command[model_name_idx].rstrip("/")
 
                     if model_name_or_path.startswith("gs://"):
                         model_name_hash = hashlib.md5(model_name_or_path.encode("utf-8")).hexdigest()[:8]
-                        local_cache_folder = f"{args.auto_output_dir_path}/{whoami}/tokenizer_{model_name_hash}"
+                        local_cache_folder = f"{args.auto_output_dir_path}/{whoami}/tokenizer_{model_name_hash}/"
 
                         if not os.path.exists(local_cache_folder):
-                            download_from_gs_bucket(f"{model_name_or_path}/tokenizer*", local_cache_folder)
+                            download_from_gs_bucket(
+                                [
+                                    f"{model_name_or_path}/tokenizer.json",
+                                    f"{model_name_or_path}/tokenizer_config.json",
+                                    f"{model_name_or_path}/config.json",
+                                ],
+                                local_cache_folder,
+                            )
 
                         filtered_command[model_name_idx] = local_cache_folder
                 except ValueError:

open_instruct/utils.py

@@ -1064,7 +1064,8 @@ def download_from_hf(model_name_or_path: str, revision: str) -> None:
     return output
 
 
-def download_from_gs_bucket(src_path: str, dest_path: str) -> None:
+def download_from_gs_bucket(src_paths: str | list[str], dest_path: str) -> None:
+    os.makedirs(dest_path, exist_ok=True)
     cmd = [
         "gsutil",
         "-o",
@@ -1074,9 +1075,11 @@ def download_from_gs_bucket(src_path: str, dest_path: str) -> None:
         "-m",
         "cp",
         "-r",
-        src_path,
-        dest_path,
     ]
+    if not isinstance(src_paths, list):
+        src_paths = [src_paths]
+    cmd.extend(src_paths)
+    cmd.append(dest_path)
     print(f"Downloading from GS bucket with command: {cmd}")
     live_subprocess_output(cmd)