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
This commit is contained in:
Roland Tannous 2026-03-10 19:20:12 +00:00
commit b8678a3ed6
2 changed files with 10 additions and 6 deletions

View file

@ -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({

View file

@ -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