unsloth/studio/backend/utils/models/__init__.py
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

43 lines
1,003 B
Python

# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
# Copyright © 2025 Unsloth AI
"""
Model and LoRA configuration handling
"""
from .model_config import (
ModelConfig,
GgufVariantInfo,
is_vision_model,
is_embedding_model,
detect_audio_type,
is_audio_input_type,
VALID_AUDIO_TYPES,
scan_trained_loras,
scan_exported_models,
load_model_defaults,
get_base_model_from_lora,
load_model_config,
list_gguf_variants,
MODEL_NAME_MAPPING,
UI_STATUS_INDICATORS,
)
from .checkpoints import scan_checkpoints
__all__ = [
'ModelConfig',
'GgufVariantInfo',
'is_vision_model',
'is_embedding_model',
'detect_audio_type',
'is_audio_input_type',
'VALID_AUDIO_TYPES',
'scan_trained_loras',
'scan_exported_models',
'load_model_defaults',
'get_base_model_from_lora',
'load_model_config',
'list_gguf_variants',
'MODEL_NAME_MAPPING',
'UI_STATUS_INDICATORS',
'scan_checkpoints',
]