unsloth/studio/backend/assets/configs/model_defaults/embedding/unsloth_all-MiniLM-L6-v2.yaml
Roland Tannous 5a086353ab 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
2026-03-10 18:10:09 +00:00

39 lines
777 B
YAML

# 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