From b8678a3ed69822152d85fe7dd4a908400af360b5 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Tue, 10 Mar 2026 19:20:12 +0000
Subject: [PATCH] fix: pass hf_token for gated embedding models and key cache
 by token

- Forward hf_token to FastSentenceTransformer.from_pretrained() so
  private/gated embedding repos authenticate correctly
- Key _embedding_detection_cache by (model_name, hf_token) tuple so
  unauthenticated lookups don't shadow subsequent authenticated ones
---
 studio/backend/core/training/worker.py      |  3 +++
 studio/backend/utils/models/model_config.py | 13 +++++++------
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 3bd5eb7653..48b99ca9ce 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -455,6 +455,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
     # ── 2. Load model ──
     _send_status(event_queue, "Loading embedding model...")
     try:
+        hf_token = config.get("hf_token", "")
+        hf_token = hf_token if hf_token and hf_token.strip() else None
         max_seq_length = config.get("max_seq_length", 512)
         training_type = config.get("training_type", "LoRA/QLoRA")
         use_lora = (training_type == "LoRA/QLoRA")
@@ -463,6 +465,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
             model_name=model_name,
             max_seq_length=max_seq_length,
             full_finetuning=not use_lora,
+            token=hf_token,
         )
     except Exception as e:
         event_queue.put({
diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py
index 7b619e46b9..f9f43b7191 100644
--- a/studio/backend/utils/models/model_config.py
+++ b/studio/backend/utils/models/model_config.py
@@ -919,7 +919,7 @@ def download_gguf_file(
 
 
 # Cache embedding detection results per session to avoid repeated HF API calls
-_embedding_detection_cache: Dict[str, bool] = {}
+_embedding_detection_cache: Dict[tuple, bool] = {}
 
 
 def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
@@ -942,12 +942,13 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
         True if the model is an embedding model, False otherwise.
         Defaults to False for local paths or on errors.
     """
-    if model_name in _embedding_detection_cache:
-        return _embedding_detection_cache[model_name]
+    cache_key = (model_name, hf_token)
+    if cache_key in _embedding_detection_cache:
+        return _embedding_detection_cache[cache_key]
 
     # Local paths have no HF metadata to query
     if is_local_path(model_name):
-        _embedding_detection_cache[model_name] = False
+        _embedding_detection_cache[cache_key] = False
         return False
 
     try:
@@ -963,7 +964,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
             or pipeline_tag in ("sentence-similarity", "feature-extraction")
         )
 
-        _embedding_detection_cache[model_name] = is_emb
+        _embedding_detection_cache[cache_key] = is_emb
         if is_emb:
             logger.info(
                 f"Model {model_name} detected as embedding model: "
@@ -975,7 +976,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
 
     except Exception as e:
         logger.warning(f"Could not determine if {model_name} is embedding model: {e}")
-        _embedding_detection_cache[model_name] = False
+        _embedding_detection_cache[cache_key] = False
         return False