unsloth/studio/backend/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

119 lines
2.6 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
# Copyright © 2025 Unsloth AI
"""
Pydantic models for API request/response schemas
"""
from .training import (
TrainingStartRequest,
TrainingJobResponse,
TrainingStatus,
TrainingProgress,
)
from .models import (
CheckpointInfo,
ModelCheckpoints,
CheckpointListResponse,
ModelDetails,
LocalModelInfo,
LocalModelListResponse,
LoRAInfo,
LoRAScanResponse,
ModelListResponse,
)
from .auth import (
AuthSetupRequest,
AuthLoginRequest,
RefreshTokenRequest,
AuthStatusResponse,
)
from .export import (
LoadCheckpointRequest,
ExportStatusResponse,
ExportOperationResponse,
ExportMergedModelRequest,
ExportBaseModelRequest,
ExportGGUFRequest,
ExportLoRAAdapterRequest,
)
from .users import Token
from .datasets import (
CheckFormatRequest,
CheckFormatResponse,
)
from .inference import (
LoadRequest,
UnloadRequest,
GenerateRequest,
LoadResponse,
UnloadResponse,
InferenceStatusResponse,
)
from .responses import (
TrainingStopResponse,
TrainingMetricsResponse,
LoRABaseModelResponse,
VisionCheckResponse,
EmbeddingCheckResponse,
)
from .data_recipe import (
RecipePayload,
PreviewResponse,
ValidateError,
ValidateResponse,
JobCreateResponse,
)
__all__ = [
# Training schemas
"TrainingStartRequest",
"TrainingJobResponse",
"TrainingStatus",
"TrainingProgress",
# Model management schemas
"ModelDetails",
"LocalModelInfo",
"LocalModelListResponse",
"LoRAInfo",
"LoRAScanResponse",
"ModelListResponse",
# Auth schemas
"AuthSetupRequest",
"AuthLoginRequest",
"RefreshTokenRequest",
"AuthStatusResponse",
# Export schemas
"CheckpointInfo",
"ModelCheckpoints",
"CheckpointListResponse",
"LoadCheckpointRequest",
"ExportStatusResponse",
"ExportOperationResponse",
"ExportMergedModelRequest",
"ExportBaseModelRequest",
"ExportGGUFRequest",
"ExportLoRAAdapterRequest",
"Token",
# Dataset schemas
"CheckFormatRequest",
"CheckFormatResponse",
# Inference schemas
"LoadRequest",
"UnloadRequest",
"GenerateRequest",
"LoadResponse",
"UnloadResponse",
"InferenceStatusResponse",
# Response schemas
"TrainingStopResponse",
"TrainingMetricsResponse",
"LoRABaseModelResponse",
"VisionCheckResponse",
"EmbeddingCheckResponse",
# Data recipe
"RecipePayload",
"PreviewResponse",
"ValidateError",
"ValidateResponse",
"JobCreateResponse",
]