From 5a086353ab655d78cd6ffd459e67f8bf528ec8bb Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 18:10:09 +0000 Subject: [PATCH] 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;