GCS base_model from pretraining

Author: kddubey

src/grouping_trainer/utils.py

@@ -1,4 +1,7 @@
+import logging
 import math
+import os
+import subprocess
 from collections.abc import Iterable
 from typing import Literal, cast, overload
 
@@ -9,6 +12,11 @@
 from sentence_transformers import SentenceTransformer as SentenceTransformerOriginal
 from transformers import PreTrainedTokenizerBase
 
+logger = logging.getLogger(__name__)
+
+_GCS_PREFIX = "gs://"
+_DIR_BASE_MODELS_LOCAL = "_base_models"
+
 
 def concat_vertical_unordered(
     dfs: Iterable[pl.DataFrame],
@@ -122,13 +130,45 @@ def encode(self, sentences: str | list[str] | np.ndarray, **kwargs) -> np.ndarra
         return embeddings[[text_to_idx[text] for text in texts]]  # assume numpy or torch
 
 
+def is_gcs_uri(uri: str) -> bool:
+    return uri.startswith(_GCS_PREFIX)
+
+
+def assert_gcs_path_exists(uri: str) -> None:
+    """
+    Raises `CalledProcessError` if `uri` doesn't exist (or matches no objects) in GCS.
+    """
+    subprocess.run(
+        ["gcloud", "storage", "ls", uri.rstrip("/") + "/"],
+        check=True,
+        stdout=subprocess.DEVNULL,
+    )
+
+
+def _download_base_model_from_gcs(uri: str) -> str:
+    """
+    Rsync `uri` (a `gs://...` directory) into `_base_models/<basename>/` relative to CWD and return the local path.
+    """
+    basename = uri.rstrip("/").rsplit("/", 1)[-1]
+    path_local = os.path.join(_DIR_BASE_MODELS_LOCAL, basename)
+    logger.info(f"Downloading base model: {uri} -> {path_local}")
+    subprocess.run(["gcloud", "storage", "rsync", "-r", uri.rstrip("/"), path_local], check=True)
+    return path_local
+
+
 def encoder_from_base(base_model: str, use_text_prefix: bool = True) -> SentenceTransformer:
     """
     Build a SentenceTransformer encoder with standard dtype/attention settings.
 
+    `base_model` is a HuggingFace model ID, a local path, or a `gs://...` path to a custom model directory. gs:// models
+    are downloaded into `_base_models/` (relative to CWD) on first call.
+
     Handles model-specific quirks (e.g. jina v5's config_kwargs and trust_remote_code) and enables bfloat16 + SDPA when
     supported.
     """
+    if is_gcs_uri(base_model):
+        base_model = _download_base_model_from_gcs(base_model)
+
     if base_model == "jinaai/jina-embeddings-v5-text-nano-text-matching":
         return SentenceTransformer(
             base_model,

train.py

@@ -49,8 +49,9 @@ def run(
     Parameters
     ----------
     base_model
-        HuggingFace model ID or local path for the base encoder. Others we've tried: Alibaba-NLP/gte-modernbert-base,
-        Qwen/Qwen3-Embedding-0.6B, jinaai/jina-embeddings-v5-text-nano-text-matching
+        HuggingFace model ID or local path for the base encoder, or a `gs://...` path to a custom model directory in our
+        bucket (downloaded into `_base_models/` on the instance). Others we've tried:
+        Alibaba-NLP/gte-modernbert-base, Qwen/Qwen3-Embedding-0.6B, jinaai/jina-embeddings-v5-text-nano-text-matching
     run_shortname
         Short name for the run. Doesn't need to be unique b/c it's appended to the timestamp.
     per_device_token_budget_scale
@@ -79,6 +80,10 @@ def run(
     if not tiny_run:
         assert run_shortname is not None, "run_shortname is required for full training runs"
 
+    # Fail fast on a typo'd gs:// model URI before wasting time launching training.
+    if gt.utils.is_gcs_uri(base_model):
+        gt.utils.assert_gcs_path_exists(base_model)
+
     # Generate run_name up front so we can log the artifact URL locally before auto-launching. On the remote, re-use the
     # local run_name via env var so both sides log the same GCS path (rather than each generating its own timestamp).
     run_name = os.environ.get(_RUN_NAME_ENV_VAR) or gt.launch.run_name_from_shortname(run_shortname or "tiny-run")