ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: add embedding model training support

Browse source 
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
This commit is contained in:
Roland Tannous 2026-03-10 18:10:09 +00:00
commit 5a086353ab
20 changed files with 758 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

43
studio/backend/assets/configs/model_defaults/embedding/unsloth_Qwen3-Embedding-0.6B.yaml Normal file
View file

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

39
studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml Normal file
View file

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

39
studio/backend/assets/configs/model_defaults/embedding/unsloth_bge-m3.yaml Normal file
View file

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

42
studio/backend/assets/configs/model_defaults/embedding/unsloth_embeddinggemma-300m.yaml Normal file
View file

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

38
studio/backend/assets/configs/model_defaults/embedding/unsloth_gte-modernbert-base.yaml Normal file
View file

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

363
studio/backend/core/training/worker.py
View file

@ -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(),
})

2
studio/backend/models/__init__.py
View file

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

1
studio/backend/models/models.py
View file

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

6
studio/backend/models/responses.py
View file

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

1
studio/backend/models/training.py
View file

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

43
studio/backend/routes/models.py
View file

@ -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')"),

1
studio/backend/routes/training.py
View file

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

2
studio/backend/utils/models/__init__.py
View file

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

85
studio/backend/utils/models/model_config.py
View file

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

30
studio/frontend/src/features/studio/sections/params-section.tsx
View file

@ -785,7 +785,7 @@ export function ParamsSection(): ReactElement {
</SelectContent>
</Select>
</Row>
{!showVisionLora && (
{!showVisionLora && !store.isEmbeddingModel && (
<div className="flex items-center gap-2">
<Checkbox
id="packing"
@ -800,19 +800,21 @@ export function ParamsSection(): ReactElement {
</label>
</div>
)}
<div className="flex items-center gap-2">
<Checkbox
id="trainOnCompletions"
checked={store.trainOnCompletions}
onCheckedChange={(v) => store.setTrainOnCompletions(!!v)}
/>
<label
htmlFor="trainOnCompletions"
className="text-xs cursor-pointer text-muted-foreground"
>
Assistant completions only
</label>
</div>
{!store.isEmbeddingModel && (
<div className="flex items-center gap-2">
<Checkbox
id="trainOnCompletions"
checked={store.trainOnCompletions}
onCheckedChange={(v) => store.setTrainOnCompletions(!!v)}
/>
<label
htmlFor="trainOnCompletions"
className="text-xs cursor-pointer text-muted-foreground"
>
Assistant completions only
</label>
</div>
)}
</TabsContent>
</Tabs>
</CollapsibleContent>

14
studio/frontend/src/features/training/api/mappers.ts
View file

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

23
studio/frontend/src/features/training/api/models-api.ts
View file

@ -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<boolean> {
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<boolean> {
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,

10
studio/frontend/src/features/training/stores/training-config-store.ts
View file

@ -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<keyof TrainingConfigState> = new Set([
"modelType",
"isCheckingVision",
"isEmbeddingModel",
"isLoadingModelDefaults",
"modelDefaultsError",
"modelDefaultsAppliedFor",
@ -128,14 +130,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
}
// 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<TrainingConfigStore>()(
selectedModel: null,
isCheckingVision: false,
isVisionModel: false,
isEmbeddingModel: false,
isDatasetAudio: false,
isLoadingModelDefaults: false,
modelDefaultsError: null,
@ -249,6 +254,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
set({
isCheckingVision: false,
isVisionModel: false,
isEmbeddingModel: false,
isDatasetAudio: false,
isLoadingModelDefaults: false,
modelDefaultsError: null,

1
studio/frontend/src/features/training/types/api.ts
View file

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

1
studio/frontend/src/features/training/types/config.ts
View file

@ -60,6 +60,7 @@ export interface TrainingConfigState {
logFrequency: number;
isCheckingVision: boolean;
isVisionModel: boolean;
isEmbeddingModel: boolean;
isLoadingModelDefaults: boolean;
modelDefaultsError: string | null;
modelDefaultsAppliedFor: string | null;