From 5a086353ab655d78cd6ffd459e67f8bf528ec8bb Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 18:10:09 +0000 Subject: [PATCH 1/9] feat: add embedding model training support Add end-to-end embedding/sentence-transformer training pipeline using FastSentenceTransformer, SentenceTransformerTrainer, and MultipleNegativesRankingLoss with BatchSamplers.NO_DUPLICATES. Backend: - Add is_embedding_model() detection via HF tags + pipeline_tag - Add /check-embedding/ API route and EmbeddingCheckResponse - Extend derive_model_type() to return "embeddings" - Add _run_embedding_training() in worker.py with progress callbacks, stop handling, LoRA (task_type=FEATURE_EXTRACTION), and model saving - Add is_embedding field to TrainingStartRequest and ModelDetails - Add YAML configs for 5 models: all-MiniLM-L6-v2, bge-m3, embeddinggemma-300m, gte-modernbert-base, Qwen3-Embedding-0.6B Frontend: - Wire isEmbeddingModel flag through store, API types, and mappers - Force packing=false, train_on_completions=false, warmup_ratio=0.03 - Hide packing and train_on_completions checkboxes for embedding models - Auto-set modelType to "embeddings" from backend model_type response --- .../unsloth_Qwen3-Embedding-0.6B.yaml | 43 +++ .../embedding/unsloth_all-MiniLM-L6-v2.yaml | 39 ++ .../embedding/unsloth_bge-m3.yaml | 39 ++ .../unsloth_embeddinggemma-300m.yaml | 42 ++ .../unsloth_gte-modernbert-base.yaml | 38 ++ studio/backend/core/training/worker.py | 363 ++++++++++++++++++ studio/backend/models/__init__.py | 2 + studio/backend/models/models.py | 1 + studio/backend/models/responses.py | 6 + studio/backend/models/training.py | 1 + studio/backend/routes/models.py | 43 ++- studio/backend/routes/training.py | 1 + studio/backend/utils/models/__init__.py | 2 + studio/backend/utils/models/model_config.py | 85 ++++ .../studio/sections/params-section.tsx | 30 +- .../src/features/training/api/mappers.ts | 14 +- .../src/features/training/api/models-api.ts | 23 ++ .../training/stores/training-config-store.ts | 10 +- .../src/features/training/types/api.ts | 1 + .../src/features/training/types/config.ts | 1 + 20 files changed, 758 insertions(+), 26 deletions(-) create mode 100644 studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml create mode 100644 studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml create mode 100644 studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml create mode 100644 studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml create mode 100644 studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml new file mode 100644 index 0000000000..1219648a8a --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml @@ -0,0 +1,43 @@ +# Model defaults for unsloth/Qwen3-Embedding-0.6B +# Based on Qwen3_Embedding_(0_6B).py embedding notebook +# Also applies to: unsloth/Qwen3-Embedding-4B + +training: + max_seq_length: 512 + # num_epochs: 2 + num_epochs: 0 + learning_rate: 3e-5 + batch_size: 256 + gradient_accumulation_steps: 1 + warmup_ratio: 0.03 + max_steps: 30 + save_steps: 30 + weight_decay: 0.01 + random_seed: 3407 + packing: false + train_on_completions: false + gradient_checkpointing: false + optim: "adamw_8bit" + lr_scheduler_type: "constant_with_warmup" + +lora: + lora_r: 32 + lora_alpha: 32 + lora_dropout: 0.0 + target_modules: + - "q_proj" + - "k_proj" + - "v_proj" + - "o_proj" + - "gate_proj" + - "up_proj" + - "down_proj" + use_rslora: false + use_loftq: false + +logging: + enable_wandb: false + wandb_project: "embedding-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 50 diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml new file mode 100644 index 0000000000..db742e11b5 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml @@ -0,0 +1,39 @@ +# Model defaults for unsloth/all-MiniLM-L6-v2 +# Based on All_MiniLM_L6_v2.py embedding notebook + +training: + max_seq_length: 512 + # num_epochs: 2 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 256 + gradient_accumulation_steps: 1 + warmup_ratio: 0.03 + max_steps: 30 + save_steps: 30 + weight_decay: 0.01 + random_seed: 3407 + packing: false + train_on_completions: false + gradient_checkpointing: false + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 64 + lora_alpha: 128 + lora_dropout: 0.0 + target_modules: + - "value" + - "key" + - "dense" + - "query" + use_rslora: false + use_loftq: false + +logging: + enable_wandb: false + wandb_project: "embedding-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 50 diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml new file mode 100644 index 0000000000..499c112929 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml @@ -0,0 +1,39 @@ +# Model defaults for unsloth/bge-m3 +# Based on BGE_M3.py embedding notebook + +training: + max_seq_length: 512 + # num_epochs: 2 + num_epochs: 0 + learning_rate: 3e-5 + batch_size: 256 + gradient_accumulation_steps: 1 + warmup_ratio: 0.03 + max_steps: 30 + save_steps: 30 + weight_decay: 0.01 + random_seed: 3407 + packing: false + train_on_completions: false + gradient_checkpointing: false + optim: "adamw_8bit" + lr_scheduler_type: "constant_with_warmup" + +lora: + lora_r: 32 + lora_alpha: 64 + lora_dropout: 0.0 + target_modules: + - "key" + - "query" + - "dense" + - "value" + use_rslora: false + use_loftq: false + +logging: + enable_wandb: false + wandb_project: "embedding-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 50 diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml new file mode 100644 index 0000000000..016d284a16 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml @@ -0,0 +1,42 @@ +# Model defaults for unsloth/embeddinggemma-300m +# Based on EmbeddingGemma_(300M).py embedding notebook + +training: + max_seq_length: 1024 + # num_epochs: 1 + num_epochs: 0 + learning_rate: 2e-5 + batch_size: 64 + gradient_accumulation_steps: 2 + warmup_ratio: 0.03 + max_steps: 30 + save_steps: 30 + weight_decay: 0.01 + random_seed: 3407 + packing: false + train_on_completions: false + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 32 + lora_alpha: 64 + lora_dropout: 0.0 + target_modules: + - "q_proj" + - "k_proj" + - "v_proj" + - "o_proj" + - "gate_proj" + - "up_proj" + - "down_proj" + use_rslora: false + use_loftq: false + +logging: + enable_wandb: false + wandb_project: "embedding-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 5 diff --git a/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml new file mode 100644 index 0000000000..eefaecc8cd --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml @@ -0,0 +1,38 @@ +# Model defaults for unsloth/gte-modernbert-base +# Based on ModernBert.py embedding notebook + +training: + max_seq_length: 512 + # num_epochs: 2 + num_epochs: 0 + learning_rate: 3e-5 + batch_size: 256 + gradient_accumulation_steps: 1 + warmup_ratio: 0.03 + max_steps: 30 + save_steps: 30 + weight_decay: 0.01 + random_seed: 3407 + packing: false + train_on_completions: false + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "constant_with_warmup" + +lora: + lora_r: 64 + lora_alpha: 128 + lora_dropout: 0.0 + target_modules: + - "Wi" + - "Wo" + - "Wqkv" + use_rslora: false + use_loftq: false + +logging: + enable_wandb: false + wandb_project: "embedding-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 50 diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 7b98865eea..3bd5eb7653 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -133,6 +133,22 @@ def run_training_process( }) return + # ── 2b. EMBEDDING MODEL FAST-PATH ── + # Embedding models use a completely different pipeline (FastSentenceTransformer + # + SentenceTransformerTrainer + MultipleNegativesRankingLoss) so we branch + # early and handle the entire flow in a self-contained function. + if config.get("is_embedding", False): + try: + _run_embedding_training(event_queue, stop_queue, config) + except Exception as exc: + event_queue.put({ + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + # ── 3. Create a fresh trainer instance ── trainer = UnslothTrainer() @@ -372,3 +388,350 @@ def _send_status(event_queue: Any, message: str) -> None: "message": message, "ts": time.time(), }) + + +def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None: + """Self-contained embedding model training pipeline. + + Uses FastSentenceTransformer + SentenceTransformerTrainer + + MultipleNegativesRankingLoss — completely separate from the + LLM/VLM/audio paths in UnslothTrainer. + + Mirrors the pattern from the reference embedding notebooks: + All_MiniLM_L6_v2.py, BGE_M3.py, EmbeddingGemma_300M.py, + ModernBert.py, Qwen3_Embedding_0_6B.py + """ + import math + import queue as _queue + import threading + + model_name = config["model_name"] + training_start_time = time.time() + + # ── 1. Import embedding-specific libraries ── + _send_status(event_queue, "Importing embedding libraries...") + try: + from unsloth import FastSentenceTransformer, is_bfloat16_supported + from sentence_transformers import ( + SentenceTransformerTrainer, + SentenceTransformerTrainingArguments, + ) + from sentence_transformers.losses import MultipleNegativesRankingLoss + from sentence_transformers.training_args import BatchSamplers + from datasets import load_dataset, Dataset + from transformers import TrainerCallback + except ImportError as e: + event_queue.put({ + "type": "error", + "error": f"Failed to import embedding libraries: {e}. " + "Ensure 'sentence_transformers' and 'unsloth' are installed.", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── Stop signal handling ── + _should_stop = False + _save_on_stop = True + + def _poll_stop(): + nonlocal _should_stop, _save_on_stop + while True: + try: + msg = stop_queue.get(timeout=1.0) + if msg and msg.get("type") == "stop": + _save_on_stop = msg.get("save", True) + _should_stop = True + logger.info("Embedding training: stop signal received (save=%s)", _save_on_stop) + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target=_poll_stop, daemon=True) + stop_thread.start() + + # ── 2. Load model ── + _send_status(event_queue, "Loading embedding model...") + try: + max_seq_length = config.get("max_seq_length", 512) + training_type = config.get("training_type", "LoRA/QLoRA") + use_lora = (training_type == "LoRA/QLoRA") + + model = FastSentenceTransformer.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + full_finetuning=not use_lora, + ) + except Exception as e: + event_queue.put({ + "type": "error", + "error": f"Failed to load embedding model '{model_name}': {e}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + if _should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + return + + # ── 3. Apply LoRA ── + if use_lora: + _send_status(event_queue, "Configuring LoRA adapters (FEATURE_EXTRACTION)...") + try: + gradient_checkpointing = config.get("gradient_checkpointing", False) + # Normalize: "none" or empty → False + if gradient_checkpointing in ("none", "", None): + gradient_checkpointing = False + + model = FastSentenceTransformer.get_peft_model( + model, + r=config.get("lora_r", 32), + target_modules=config.get("target_modules") or ["q_proj", "k_proj", "v_proj", "o_proj"], + lora_alpha=config.get("lora_alpha", 64), + lora_dropout=config.get("lora_dropout", 0.0), + bias="none", + use_gradient_checkpointing=gradient_checkpointing, + random_state=config.get("random_seed", 3407), + use_rslora=config.get("use_rslora", False), + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if config.get("use_loftq") else None, + task_type="FEATURE_EXTRACTION", + ) + except Exception as e: + event_queue.put({ + "type": "error", + "error": f"Failed to configure LoRA for embedding model: {e}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + if _should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + return + + # ── 4. Load dataset ── + _send_status(event_queue, "Loading dataset...") + try: + hf_dataset = config.get("hf_dataset", "") + local_datasets = config.get("local_datasets") or [] + subset = config.get("subset") or None + train_split = config.get("train_split", "train") or "train" + + if hf_dataset and hf_dataset.strip(): + hf_token = config.get("hf_token", "") + hf_token = hf_token if hf_token and hf_token.strip() else None + dataset = load_dataset( + hf_dataset.strip(), + subset, + split=train_split, + token=hf_token, + ) + elif local_datasets: + # Load from local file(s) + local_path = local_datasets[0] + if local_path.endswith(".csv"): + dataset = load_dataset("csv", data_files=local_path, split="train") + elif local_path.endswith(".json") or local_path.endswith(".jsonl"): + dataset = load_dataset("json", data_files=local_path, split="train") + elif local_path.endswith(".parquet"): + dataset = load_dataset("parquet", data_files=local_path, split="train") + else: + dataset = load_dataset(local_path, split="train") + else: + event_queue.put({ + "type": "error", + "error": "No dataset specified for embedding training.", + "stack": "", "ts": time.time(), + }) + return + + # Apply dataset slicing if specified + slice_start = config.get("dataset_slice_start") + slice_end = config.get("dataset_slice_end") + if slice_start is not None or slice_end is not None: + start = slice_start or 0 + end = slice_end or len(dataset) + dataset = dataset.select(range(start, min(end + 1, len(dataset)))) + + logger.info(f"Embedding dataset loaded: {len(dataset)} samples") + except Exception as e: + event_queue.put({ + "type": "error", + "error": f"Failed to load dataset: {e}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + if _should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + return + + # ── 5. Create loss function ── + loss = MultipleNegativesRankingLoss(model) + + # ── 6. Build training arguments ── + _send_status(event_queue, "Configuring training...") + try: + lr_value = float(config.get("learning_rate", "2e-4")) + except ValueError: + event_queue.put({ + "type": "error", + "error": f"Invalid learning rate: {config.get('learning_rate')}", + "stack": "", "ts": time.time(), + }) + return + + output_dir = config.get("output_dir") + if not output_dir: + output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}" + + num_epochs = config.get("num_epochs", 2) + batch_size = config.get("batch_size", 256) + gradient_accumulation_steps = config.get("gradient_accumulation_steps", 1) + max_steps_val = config.get("max_steps", 0) + save_steps_val = config.get("save_steps", 0) + warmup_ratio = config.get("warmup_ratio", 0.03) + warmup_steps_val = config.get("warmup_steps") + log_frequency = config.get("log_frequency", 50) + + # Build args dict + training_args_kwargs = { + "output_dir": output_dir, + "per_device_train_batch_size": batch_size, + "gradient_accumulation_steps": gradient_accumulation_steps, + "learning_rate": lr_value, + "fp16": not is_bfloat16_supported(), + "bf16": is_bfloat16_supported(), + "logging_steps": max(1, log_frequency) if log_frequency else 1, + "report_to": ["wandb"] if config.get("enable_wandb") else "none", + "lr_scheduler_type": config.get("lr_scheduler_type", "linear"), + "batch_sampler": BatchSamplers.NO_DUPLICATES, + "optim": config.get("optim", "adamw_8bit"), + "weight_decay": config.get("weight_decay", 0.01), + "seed": config.get("random_seed", 3407), + } + + # max_steps vs epochs + if max_steps_val and max_steps_val > 0: + training_args_kwargs["max_steps"] = max_steps_val + else: + training_args_kwargs["num_train_epochs"] = num_epochs if num_epochs > 0 else 2 + + # warmup: prefer warmup_ratio (standard for embedding scripts), fallback to steps + if warmup_ratio is not None and warmup_ratio > 0: + training_args_kwargs["warmup_ratio"] = warmup_ratio + elif warmup_steps_val is not None and warmup_steps_val > 0: + training_args_kwargs["warmup_steps"] = warmup_steps_val + + # save_steps + if save_steps_val and save_steps_val > 0: + training_args_kwargs["save_steps"] = save_steps_val + training_args_kwargs["save_strategy"] = "steps" + + args = SentenceTransformerTrainingArguments(**training_args_kwargs) + + # ── 7. Calculate total steps for progress tracking ── + if max_steps_val and max_steps_val > 0: + total_steps = max_steps_val + else: + effective_epochs = num_epochs if num_epochs > 0 else 2 + len_dataloader = math.ceil(len(dataset) / batch_size) + steps_per_epoch = max(len_dataloader // gradient_accumulation_steps, 1) + total_steps = steps_per_epoch * effective_epochs + + # ── 8. Create progress callback ── + class _EmbeddingProgressCallback(TrainerCallback): + """Sends training progress events to the parent process via event_queue.""" + + def on_log(self, args, state, control, logs=None, **kwargs): + if not logs: + return + loss_value = logs.get("loss", logs.get("train_loss", 0.0)) + current_step = state.global_step + + elapsed = time.time() - training_start_time + eta = None + if current_step > 0 and total_steps > 0: + remaining = total_steps - current_step + if remaining > 0: + eta = (elapsed / current_step) * remaining + + event_queue.put({ + "type": "progress", + "step": current_step, + "epoch": round(state.epoch, 2) if state.epoch else 0, + "loss": loss_value, + "learning_rate": logs.get("learning_rate", 0.0), + "total_steps": total_steps, + "elapsed_seconds": elapsed, + "eta_seconds": eta, + "grad_norm": logs.get("grad_norm"), + "num_tokens": getattr(state, "num_input_tokens_seen", None), + "eval_loss": logs.get("eval_loss"), + "status_message": "", + "ts": time.time(), + }) + + def on_step_end(self, args, state, control, **kwargs): + if _should_stop: + logger.info("Embedding training: stop at step %d", state.global_step) + control.should_training_stop = True + return control + + # ── 9. Create trainer and train ── + _send_status(event_queue, "Starting embedding training...") + try: + trainer = SentenceTransformerTrainer( + model=model, + train_dataset=dataset, + loss=loss, + args=args, + callbacks=[_EmbeddingProgressCallback()], + ) + + trainer.train() + except Exception as e: + event_queue.put({ + "type": "error", + "error": f"Embedding training failed: {e}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 10. Save model ── + if _should_stop and not _save_on_stop: + event_queue.put({ + "type": "complete", + "output_dir": None, + "status_message": "Training cancelled", + "ts": time.time(), + }) + return + + _send_status(event_queue, "Saving model...") + try: + model.save_pretrained(output_dir) + model.tokenizer.save_pretrained(output_dir) + logger.info("Embedding model saved to %s", output_dir) + except Exception as e: + logger.error("Failed to save embedding model: %s", e) + event_queue.put({ + "type": "error", + "error": f"Training completed but failed to save: {e}", + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + return + + # ── 11. Done ── + event_queue.put({ + "type": "complete", + "output_dir": output_dir, + "status_message": "Embedding training completed", + "ts": time.time(), + }) diff --git a/studio/backend/models/__init__.py b/studio/backend/models/__init__.py index a6c1b38b34..49ef8998ee 100644 --- a/studio/backend/models/__init__.py +++ b/studio/backend/models/__init__.py @@ -54,6 +54,7 @@ from .responses import ( TrainingMetricsResponse, LoRABaseModelResponse, VisionCheckResponse, + EmbeddingCheckResponse, ) from .data_recipe import ( RecipePayload, @@ -108,6 +109,7 @@ __all__ = [ "TrainingMetricsResponse", "LoRABaseModelResponse", "VisionCheckResponse", + "EmbeddingCheckResponse", # Data recipe "RecipePayload", "PreviewResponse", diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 67eaf3fb96..294134d430 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -57,6 +57,7 @@ class ModelDetails(BaseModel): name: Optional[str] = Field(None, description="Display name for the model") config: Optional[Dict[str, Any]] = Field(None, description="Model configuration dictionary") is_vision: bool = Field(False, description="Whether model is a vision model") + is_embedding: bool = Field(False, description="Whether model is an embedding/sentence-transformer model") is_lora: bool = Field(False, description="Whether model is a LoRA adapter") is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp format)") is_audio: bool = Field(False, description="Whether model is a TTS audio model") diff --git a/studio/backend/models/responses.py b/studio/backend/models/responses.py index c157717262..8e16e201fc 100644 --- a/studio/backend/models/responses.py +++ b/studio/backend/models/responses.py @@ -41,3 +41,9 @@ class VisionCheckResponse(BaseModel): """Response for checking if a model is a vision model""" model_name: str = Field(..., description="Model identifier") is_vision: bool = Field(..., description="Whether the model is a vision model") + + +class EmbeddingCheckResponse(BaseModel): + """Response for checking if a model is an embedding model""" + model_name: str = Field(..., description="Model identifier") + is_embedding: bool = Field(..., description="Whether the model is an embedding/sentence-transformer model") diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index c3b8d99bb4..c2724a73b1 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -76,6 +76,7 @@ class TrainingStartRequest(BaseModel): finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules") is_dataset_image: bool = Field(False, description="Whether the dataset contains image data") is_dataset_audio: bool = Field(False, description="Whether the dataset contains audio data") + is_embedding: bool = Field(False, description="Whether model is an embedding/sentence-transformer model") # Logging parameters enable_wandb: bool = Field(False, description="Enable Weights & Biases logging") diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 05b18aa606..2fb0171c94 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -25,6 +25,7 @@ try: load_model_defaults, get_base_model_from_lora, is_vision_model, + is_embedding_model, scan_checkpoints, list_gguf_variants, ModelConfig, @@ -42,6 +43,7 @@ except ImportError: load_model_defaults, get_base_model_from_lora, is_vision_model, + is_embedding_model, scan_checkpoints, list_gguf_variants, ModelConfig, @@ -61,14 +63,16 @@ from models import ( ModelListResponse, ) from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType -from models.responses import LoRABaseModelResponse, VisionCheckResponse +from models.responses import LoRABaseModelResponse, VisionCheckResponse, EmbeddingCheckResponse router = APIRouter() logger = logging.getLogger(__name__) -def derive_model_type(is_vision: bool, audio_type: Optional[str]) -> ModelType: +def derive_model_type(is_vision: bool, audio_type: Optional[str], is_embedding: bool = False) -> ModelType: """Collapse individual capability flags into a single model modality string.""" + if is_embedding: + return "embeddings" if audio_type is not None: return "audio" if is_vision: @@ -299,6 +303,7 @@ async def get_model_config( # Detect model capabilities (pass HF token for gated models) is_vision = is_vision_model(model_name) + is_embedding = is_embedding_model(model_name, hf_token=hf_token) audio_type = detect_audio_type(model_name, hf_token=hf_token) # Check if it's a LoRA adapter @@ -311,17 +316,18 @@ async def get_model_config( except Exception: pass - logger.info(f"Model config result for {model_name}: is_vision={is_vision}, audio_type={audio_type}, is_lora={is_lora}") + logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_embedding={is_embedding}, audio_type={audio_type}, is_lora={is_lora}") return ModelDetails( id=model_name, model_name=model_name, config=config_dict, is_vision=is_vision, + is_embedding=is_embedding, is_lora=is_lora, is_audio=audio_type is not None, audio_type=audio_type, has_audio_input=is_audio_input_type(audio_type), - model_type=derive_model_type(is_vision, audio_type), + model_type=derive_model_type(is_vision, audio_type, is_embedding), base_model=base_model, ) @@ -444,6 +450,35 @@ async def check_vision_model( detail=f"Failed to check vision model: {str(e)}" ) +@router.get("/check-embedding/{model_name:path}", response_model=EmbeddingCheckResponse) +async def check_embedding_model( + model_name: str, + hf_token: Optional[str] = Query(None), + current_subject: str = Depends(get_current_subject), +): + """ + Check if a model is an embedding model. + + This endpoint wraps the backend is_embedding_model function. + """ + try: + logger.info(f"Checking if embedding model: {model_name}") + is_embedding = is_embedding_model(model_name, hf_token=hf_token) + + logger.info(f"Embedding check result for {model_name}: is_embedding={is_embedding}") + return EmbeddingCheckResponse( + model_name=model_name, + is_embedding=is_embedding, + ) + + except Exception as e: + logger.error(f"Error checking embedding model: {e}", exc_info=True) + raise HTTPException( + status_code=500, + detail=f"Failed to check embedding model: {str(e)}" + ) + + @router.get("/gguf-variants", response_model=GgufVariantsResponse) async def get_gguf_variants( repo_id: str = Query(..., description="HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"), diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index c9dd63a6a4..e3438824f6 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -200,6 +200,7 @@ async def start_training( "finetune_mlp_modules": request.finetune_mlp_modules, "is_dataset_image": request.is_dataset_image, "is_dataset_audio": request.is_dataset_audio, + "is_embedding": request.is_embedding, "enable_wandb": request.enable_wandb, "wandb_token": request.wandb_token or "", "wandb_project": request.wandb_project or "", diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 5463d98222..8674d6851a 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -8,6 +8,7 @@ from .model_config import ( ModelConfig, GgufVariantInfo, is_vision_model, + is_embedding_model, detect_audio_type, is_audio_input_type, VALID_AUDIO_TYPES, @@ -26,6 +27,7 @@ __all__ = [ 'ModelConfig', 'GgufVariantInfo', 'is_vision_model', + 'is_embedding_model', 'detect_audio_type', 'is_audio_input_type', 'VALID_AUDIO_TYPES', diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 8724f4b00b..7b619e46b9 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -25,6 +25,30 @@ logger = logging.getLogger(__name__) # Format: "canonical_model_name.yaml": [list of all equivalent model names] # Based on the model mapper provided - canonical filename is based on the first model name in the mapper MODEL_NAME_MAPPING = { + # ── Embedding models ── + "unsloth_all-MiniLM-L6-v2.yaml": [ + "unsloth/all-MiniLM-L6-v2", + "sentence-transformers/all-MiniLM-L6-v2", + ], + "unsloth_bge-m3.yaml": [ + "unsloth/bge-m3", + "BAAI/bge-m3", + ], + "unsloth_embeddinggemma-300m.yaml": [ + "unsloth/embeddinggemma-300m", + "google/embeddinggemma-300m", + ], + "unsloth_gte-modernbert-base.yaml": [ + "unsloth/gte-modernbert-base", + "Alibaba-NLP/gte-modernbert-base", + ], + "unsloth_Qwen3-Embedding-0.6B.yaml": [ + "unsloth/Qwen3-Embedding-0.6B", + "Qwen/Qwen3-Embedding-0.6B", + "unsloth/Qwen3-Embedding-4B", + "Qwen/Qwen3-Embedding-4B", + ], + # ── Other models ── "unsloth_answerdotai_ModernBERT-large.yaml": [ "answerdotai/ModernBERT-large", ], @@ -894,6 +918,67 @@ def download_gguf_file( return local_path +# Cache embedding detection results per session to avoid repeated HF API calls +_embedding_detection_cache: Dict[str, bool] = {} + + +def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: + """ + Detect embedding/sentence-transformer models using HuggingFace model metadata. + + Uses a belt-and-suspenders approach combining three signals: + 1. "sentence-transformers" in model tags + 2. "feature-extraction" in model tags + 3. pipeline_tag is "sentence-similarity" or "feature-extraction" + + This catches all known embedding models including those like gte-modernbert + whose library_name is "transformers" rather than "sentence-transformers". + + Args: + model_name: Model identifier (HF repo or local path) + hf_token: Optional HF token for accessing gated/private models + + Returns: + 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] + + # Local paths have no HF metadata to query + if is_local_path(model_name): + _embedding_detection_cache[model_name] = False + return False + + try: + from huggingface_hub import model_info as hf_model_info + + info = hf_model_info(model_name, token=hf_token) + tags = set(info.tags or []) + pipeline_tag = info.pipeline_tag or "" + + is_emb = ( + "sentence-transformers" in tags + or "feature-extraction" in tags + or pipeline_tag in ("sentence-similarity", "feature-extraction") + ) + + _embedding_detection_cache[model_name] = is_emb + if is_emb: + logger.info( + f"Model {model_name} detected as embedding model: " + f"pipeline_tag={pipeline_tag}, " + f"sentence-transformers in tags={('sentence-transformers' in tags)}, " + f"feature-extraction in tags={('feature-extraction' in tags)}" + ) + return is_emb + + except Exception as e: + logger.warning(f"Could not determine if {model_name} is embedding model: {e}") + _embedding_detection_cache[model_name] = False + return False + + def scan_trained_loras(outputs_dir: str = "./outputs") -> List[Tuple[str, str]]: """ Scan outputs folder for trained LoRA adapters. diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 51d4c56821..6928e4c4ea 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -785,7 +785,7 @@ export function ParamsSection(): ReactElement { - {!showVisionLora && ( + {!showVisionLora && !store.isEmbeddingModel && (
)} -
- store.setTrainOnCompletions(!!v)} - /> - -
+ {!store.isEmbeddingModel && ( +
+ store.setTrainOnCompletions(!!v)} + /> + +
+ )} diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index fed17a538c..f9fb07d6dc 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -25,6 +25,7 @@ export function buildTrainingStartPayload( ): TrainingStartRequest { const adapterMethod = config.trainingMethod !== "full"; const isQloraMethod = config.trainingMethod === "qlora"; + const isEmbedding = config.isEmbeddingModel; const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null; const localDatasets = config.datasetSource === "upload" && config.uploadedFile @@ -53,14 +54,14 @@ export function buildTrainingStartPayload( learning_rate: String(config.learningRate), batch_size: config.batchSize, gradient_accumulation_steps: config.gradientAccumulation, - warmup_steps: config.warmupSteps, - warmup_ratio: null, + warmup_steps: isEmbedding ? null : config.warmupSteps, + warmup_ratio: isEmbedding ? 0.03 : null, max_steps: config.maxSteps, save_steps: config.saveSteps, eval_steps: config.evalSteps, weight_decay: config.weightDecay, random_seed: config.randomSeed, - packing: config.packing, + packing: isEmbedding ? false : config.packing, optim: config.optimizerType, lr_scheduler_type: config.lrSchedulerType, use_lora: adapterMethod, @@ -71,13 +72,14 @@ export function buildTrainingStartPayload( gradient_checkpointing: config.gradientCheckpointing, use_rslora: config.loraVariant === "rslora", use_loftq: config.loraVariant === "loftq", - train_on_completions: config.trainOnCompletions, + train_on_completions: isEmbedding ? false : config.trainOnCompletions, finetune_vision_layers: config.finetuneVisionLayers, finetune_language_layers: config.finetuneLanguageLayers, finetune_attention_modules: config.finetuneAttentionModules, finetune_mlp_modules: config.finetuneMLPModules, - is_dataset_image: !!config.isDatasetImage, - is_dataset_audio: config.isDatasetAudio, + is_dataset_image: isEmbedding ? false : !!config.isDatasetImage, + is_dataset_audio: isEmbedding ? false : config.isDatasetAudio, + is_embedding: isEmbedding, enable_wandb: config.enableWandb, wandb_token: config.enableWandb ? config.wandbToken.trim() || null : null, wandb_project: config.enableWandb diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index f40f29fda2..64ac3884e3 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -8,6 +8,11 @@ interface VisionCheckResponse { is_vision: boolean; } +interface EmbeddingCheckResponse { + model_name: string; + is_embedding: boolean; +} + interface BackendTrainingDefaults { max_seq_length?: number; num_epochs?: number; @@ -61,6 +66,7 @@ export interface ModelConfigResponse { model_name?: string | null; config?: BackendModelConfig | null; is_vision: boolean; + is_embedding?: boolean; is_lora: boolean; is_audio?: boolean; base_model?: string | null; @@ -97,6 +103,23 @@ export async function checkVisionModel(modelName: string): Promise { return data.is_vision; } +/** + * Check whether a model is an embedding model by asking the backend. + * Calls GET /api/models/check-embedding/{model_name}. + */ +export async function checkEmbeddingModel( + modelName: string, +): Promise { + const encoded = encodeURIComponent(modelName); + const response = await authFetch(`/api/models/check-embedding/${encoded}`); + if (!response.ok) { + // If the check fails (e.g. network error), default to non-embedding + return false; + } + const data = (await response.json()) as EmbeddingCheckResponse; + return data.is_embedding; +} + export async function getModelConfig( modelName: string, signal?: AbortSignal, diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 502889ee1e..4ee25ca668 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -36,6 +36,7 @@ const initialState: TrainingConfigState = { uploadedFile: null, isCheckingVision: false, isVisionModel: false, + isEmbeddingModel: false, isLoadingModelDefaults: false, modelDefaultsError: null, modelDefaultsAppliedFor: null, @@ -58,6 +59,7 @@ let _trainOnCompletionsManuallySet = false; const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set([ "modelType", "isCheckingVision", + "isEmbeddingModel", "isLoadingModelDefaults", "modelDefaultsError", "modelDefaultsAppliedFor", @@ -128,14 +130,16 @@ export const useTrainingConfigStore = create()( } // Use backend-provided model_type when available, otherwise - // infer from is_vision (temporary until backend ships model_type). + // infer from capability flags. + const isEmbedding = !!modelDetails.is_embedding; const inferredModelType: ModelType = modelDetails.model_type - ?? (modelDetails.is_vision ? "vision" : modelDetails.is_audio ? "audio" : "text"); + ?? (isEmbedding ? "embeddings" : modelDetails.is_vision ? "vision" : modelDetails.is_audio ? "audio" : "text"); set({ ...patch, modelType: inferredModelType, isVisionModel: modelDetails.is_vision, + isEmbeddingModel: isEmbedding, isLoadingModelDefaults: false, isCheckingVision: false, modelDefaultsError: null, @@ -233,6 +237,7 @@ export const useTrainingConfigStore = create()( selectedModel: null, isCheckingVision: false, isVisionModel: false, + isEmbeddingModel: false, isDatasetAudio: false, isLoadingModelDefaults: false, modelDefaultsError: null, @@ -249,6 +254,7 @@ export const useTrainingConfigStore = create()( set({ isCheckingVision: false, isVisionModel: false, + isEmbeddingModel: false, isDatasetAudio: false, isLoadingModelDefaults: false, modelDefaultsError: null, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index f222f4352e..54d5b98e72 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -47,6 +47,7 @@ export interface TrainingStartRequest { finetune_mlp_modules: boolean; is_dataset_image: boolean; is_dataset_audio: boolean; + is_embedding: boolean; enable_wandb: boolean; wandb_token: string | null; wandb_project: string | null; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index bc8305430b..66d5e339ce 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -60,6 +60,7 @@ export interface TrainingConfigState { logFrequency: number; isCheckingVision: boolean; isVisionModel: boolean; + isEmbeddingModel: boolean; isLoadingModelDefaults: boolean; modelDefaultsError: string | null; modelDefaultsAppliedFor: string | null; From d6ae910edc1b161fb7d5622a8b159dcff67d2ef9 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 19:05:47 +0000 Subject: [PATCH 2/9] fix: propagate is_embedding into worker subprocess config start_training() cherry-picks kwargs into a config dict but was missing is_embedding, so config.get("is_embedding", False) in worker.py always returned False and embedding training never ran. --- studio/backend/core/training/training.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 14ff57bbd7..5f9cb86b0c 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -150,6 +150,7 @@ class TrainingBackend: "custom_format_mapping": kwargs.get("custom_format_mapping"), "is_dataset_image": kwargs.get("is_dataset_image", False), "is_dataset_audio": kwargs.get("is_dataset_audio", False), + "is_embedding": kwargs.get("is_embedding", False), "num_epochs": kwargs.get("num_epochs", 3), "learning_rate": kwargs.get("learning_rate", "2e-4"), "batch_size": kwargs.get("batch_size", 2), From b8678a3ed69822152d85fe7dd4a908400af360b5 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 19:20:12 +0000 Subject: [PATCH 3/9] 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 From c986174c56f74c72320ede51ad9f7d58b3fa0f50 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 19:33:10 +0000 Subject: [PATCH 4/9] fix: preserve zero-valued dataset slice boundaries in embedding worker Use explicit None checks instead of falsy `or` for slice_start and slice_end so that a valid slice_end=0 is not replaced with the full dataset length. --- studio/backend/core/training/worker.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 48b99ca9ce..368e756887 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -555,8 +555,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> slice_start = config.get("dataset_slice_start") slice_end = config.get("dataset_slice_end") if slice_start is not None or slice_end is not None: - start = slice_start or 0 - end = slice_end or len(dataset) + start = slice_start if slice_start is not None else 0 + end = slice_end if slice_end is not None else len(dataset) dataset = dataset.select(range(start, min(end + 1, len(dataset)))) logger.info(f"Embedding dataset loaded: {len(dataset)} samples") From 3b0b002b340337b6aba383aa09ec1e69c275b502 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 20:32:10 +0000 Subject: [PATCH 5/9] fixing logging for each step --- studio/backend/core/training/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 368e756887..34af8d9f94 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -609,7 +609,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> "learning_rate": lr_value, "fp16": not is_bfloat16_supported(), "bf16": is_bfloat16_supported(), - "logging_steps": max(1, log_frequency) if log_frequency else 1, + "logging_steps": 1, "report_to": ["wandb"] if config.get("enable_wandb") else "none", "lr_scheduler_type": config.get("lr_scheduler_type", "linear"), "batch_sampler": BatchSamplers.NO_DUPLICATES, From 9523e5c1f9721c2ed246adf5a6fb501da43875d7 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 21:12:24 +0000 Subject: [PATCH 6/9] fixing embedding model search --- studio/frontend/src/hooks/use-hf-model-search.ts | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index 5f603a71e8..330455e819 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -25,6 +25,13 @@ const EXCLUDED_TAGS = new Set([ "ctranslate2", ]); +// Embedding / sentence-transformer models ship with onnx/openvino as additional +// export formats — they should not be excluded by the tag check above. +const EMBEDDING_TAGS = new Set([ + "sentence-transformers", + "feature-extraction", +]); + function withPopularitySort( input: Parameters[0], init?: Parameters[1], @@ -56,7 +63,8 @@ function makeMapModel(excludeGguf: boolean) { safetensors?: { total: number }; tags?: string[]; }; - if (m.tags?.some((t) => EXCLUDED_TAGS.has(t))) { + const isEmbedding = m.tags?.some((t) => EMBEDDING_TAGS.has(t)); + if (!isEmbedding && m.tags?.some((t) => EXCLUDED_TAGS.has(t))) { return null; } if (excludeGguf && m.tags?.includes("gguf")) { From bc5a72dd8cc548ac30b37c487df1f5991c1e047e Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 21:29:51 +0000 Subject: [PATCH 7/9] fix: local directory dataset loading --- studio/backend/core/training/worker.py | 49 ++++++++++++++++++++------ 1 file changed, 39 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 34af8d9f94..8cc6e9785f 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -533,16 +533,45 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> token=hf_token, ) elif local_datasets: - # Load from local file(s) - local_path = local_datasets[0] - if local_path.endswith(".csv"): - dataset = load_dataset("csv", data_files=local_path, split="train") - elif local_path.endswith(".json") or local_path.endswith(".jsonl"): - dataset = load_dataset("json", data_files=local_path, split="train") - elif local_path.endswith(".parquet"): - dataset = load_dataset("parquet", data_files=local_path, split="train") - else: - dataset = load_dataset(local_path, split="train") + # Load from local file(s) — mirrors the non-embedding pipeline's + # directory handling so recipe outputs (parquet-files/) work. + all_files: list[str] = [] + for dataset_file in local_datasets: + file_path = dataset_file if os.path.isabs(dataset_file) else os.path.join( + project_root, "studio", "backend", "assets", "datasets", dataset_file, + ) + if os.path.isdir(file_path): + file_path_obj = Path(file_path) + parquet_dir = ( + file_path_obj / "parquet-files" + if (file_path_obj / "parquet-files").exists() + else file_path_obj + ) + parquet_files = sorted(parquet_dir.glob("*.parquet")) + if parquet_files: + all_files.extend(str(p) for p in parquet_files) + continue + candidates: list[Path] = [] + for ext in (".json", ".jsonl", ".csv", ".parquet"): + candidates.extend(sorted(file_path_obj.glob(f"*{ext}"))) + if candidates: + all_files.extend(str(c) for c in candidates) + continue + raise ValueError(f"No supported data files in directory: {file_path_obj}") + else: + all_files.append(file_path) + + if all_files: + first_ext = Path(all_files[0]).suffix.lower() + if first_ext in (".json", ".jsonl"): + loader = "json" + elif first_ext == ".csv": + loader = "csv" + elif first_ext == ".parquet": + loader = "parquet" + else: + raise ValueError(f"Unsupported local dataset format: {all_files[0]}") + dataset = load_dataset(loader, data_files=all_files, split="train") else: event_queue.put({ "type": "error", From 294a3d3e478e619188040a5e0f1601ca5fb5f07a Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 21:33:13 +0000 Subject: [PATCH 8/9] fix: reset isEmbeddingModel in error fallback paths to prevent stale state --- .../src/features/training/stores/training-config-store.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 56460137a9..cbaf9f375c 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -165,6 +165,7 @@ export const useTrainingConfigStore = create()( set({ isLoadingModelDefaults: false, + isEmbeddingModel: false, isAudioModel: false, modelDefaultsError: error instanceof Error @@ -179,13 +180,14 @@ export const useTrainingConfigStore = create()( set({ modelType: isVision ? "vision" : "text", isVisionModel: isVision, + isEmbeddingModel: false, isAudioModel: false, isCheckingVision: false, }); }) .catch(() => { if (get().selectedModel !== modelName) return; - set({ isCheckingVision: false, isAudioModel: false }); + set({ isCheckingVision: false, isEmbeddingModel: false, isAudioModel: false }); }); }); }; From 983c20bbb29d8333662cfe7fa0a93651e02f38f6 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 10 Mar 2026 21:58:45 +0000 Subject: [PATCH 9/9] local model's embedding nature check --- studio/backend/utils/models/model_config.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index f9f43b7191..909df0fa03 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -946,10 +946,12 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] - # Local paths have no HF metadata to query + # Local paths: check for sentence-transformer marker file (modules.json) if is_local_path(model_name): - _embedding_detection_cache[cache_key] = False - return False + local_dir = normalize_path(model_name) + is_emb = os.path.isfile(os.path.join(local_dir, "modules.json")) + _embedding_detection_cache[cache_key] = is_emb + return is_emb try: from huggingface_hub import model_info as hf_model_info