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:
parent
d6ae910edc
commit
b8678a3ed6
2 changed files with 10 additions and 6 deletions
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue