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

move inline pydantic models - fix existing models routes integration

Browse source
This commit is contained in:
Roland Tannous 2026-02-11 12:39:58 +00:00
commit 7ee4381936
7 changed files with 179 additions and 113 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -16,8 +16,28 @@ from .models import (
from .auth import (
AuthSetupRequest,
AuthLoginRequest,
RefreshTokenRequest,
AuthStatusResponse,
)
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,
)
__all__ = [
# Training schemas
@ -33,6 +53,22 @@ __all__ = [
# Auth schemas
"AuthSetupRequest",
"AuthLoginRequest",
"RefreshTokenRequest",
"AuthStatusResponse",
"Token",
# Dataset schemas
"CheckFormatRequest",
"CheckFormatResponse",
# Inference schemas
"LoadRequest",
"UnloadRequest",
"GenerateRequest",
"LoadResponse",
"UnloadResponse",
"InferenceStatusResponse",
# Response schemas
"TrainingStopResponse",
"TrainingMetricsResponse",
"LoRABaseModelResponse",
"VisionCheckResponse",
]

54
studio/backend/models/inference.py Normal file
View file

@ -0,0 +1,54 @@
"""
Pydantic schemas for Inference API
"""
from pydantic import BaseModel, Field
from typing import Optional, List
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
model_path: str = Field(..., description="Model identifier or local path")
hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models")
max_seq_length: int = Field(2048, ge=128, le=32768, description="Maximum sequence length")
load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
class UnloadRequest(BaseModel):
"""Request to unload a model"""
model_path: str = Field(..., description="Model identifier to unload")
class GenerateRequest(BaseModel):
"""Request for text generation"""
messages: List[dict] = Field(..., description="Chat messages in OpenAI format")
system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt")
temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature")
top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling")
top_k: int = Field(40, ge=1, le=100, description="Top-k sampling")
max_new_tokens: int = Field(512, ge=1, le=4096, description="Maximum tokens to generate")
repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty")
image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models")
class LoadResponse(BaseModel):
"""Response after loading a model"""
status: str = Field(..., description="Load status")
model: str = Field(..., description="Model identifier")
display_name: str = Field(..., description="Display name of the model")
is_vision: bool = Field(False, description="Whether model is a vision model")
is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
class UnloadResponse(BaseModel):
"""Response after unloading a model"""
status: str = Field(..., description="Unload status")
model: str = Field(..., description="Model identifier that was unloaded")
class InferenceStatusResponse(BaseModel):
"""Current inference backend status"""
active_model: Optional[str] = Field(None, description="Currently active model identifier")
is_vision: bool = Field(False, description="Whether the active model is a vision model")
loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
loaded: List[str] = Field(default_factory=list, description="Models currently loaded")

38
studio/backend/models/responses.py Normal file
View file

@ -0,0 +1,38 @@
"""
Pydantic response schemas for endpoints that previously returned raw dicts.
These are small response models for training and model management routes.
"""
from pydantic import BaseModel, Field
from typing import Optional, List
# --- Training route response models ---
class TrainingStopResponse(BaseModel):
"""Response for stopping a training job"""
status: str = Field(..., description="Current status: 'stopped' or 'idle'")
message: str = Field(..., description="Human-readable status message")
class TrainingMetricsResponse(BaseModel):
"""Response for training metrics history"""
loss_history: List[float] = Field(default_factory=list, description="Loss values per step")
lr_history: List[float] = Field(default_factory=list, description="Learning rate per step")
step_history: List[int] = Field(default_factory=list, description="Step numbers")
current_loss: Optional[float] = Field(None, description="Most recent loss value")
current_lr: Optional[float] = Field(None, description="Most recent learning rate")
current_step: Optional[int] = Field(None, description="Most recent step number")
# --- Model management route response models ---
class LoRABaseModelResponse(BaseModel):
"""Response for getting a LoRA's base model"""
lora_path: str = Field(..., description="Path to the LoRA adapter")
base_model: str = Field(..., description="Base model identifier")
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")

29
studio/backend/models/users.py
View file

@ -1,34 +1,15 @@
"""Pydantic models for user-related API endpoints.
"""Pydantic models for authentication tokens.
This module defines the data models used for user authentication and management
in the FastAPI application.
This module defines the Token response model used by auth routes.
"""
from pydantic import BaseModel, Field
class User(BaseModel):
"""Basic user model containing username."""
username: str
class UserInDB(BaseModel):
"""User model with password for database storage."""
password: str
class Token(BaseModel):
"""Authentication token model with access and refresh tokens."""
access_token: str
refresh_token: str
token_type: str
class TokenData(BaseModel):
"""Token payload model containing username."""
username: str | None = None
access_token: str = Field(..., description="JWT access token (60 min expiry)")
refresh_token: str = Field(..., description="Opaque refresh token (7 day expiry)")
token_type: str = Field(..., description="Token type, always 'bearer'")

71
studio/backend/routes/inference.py
View file

@ -5,8 +5,7 @@ import sys
from pathlib import Path
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from typing import Optional, List
from typing import Optional
import json
import logging
@ -26,6 +25,15 @@ except ImportError:
from core.inference import get_inference_backend
from utils.models import ModelConfig
from models.inference import (
LoadRequest,
UnloadRequest,
GenerateRequest,
LoadResponse,
UnloadResponse,
InferenceStatusResponse,
)
router = APIRouter()
logger = logging.getLogger(__name__)
@ -39,57 +47,6 @@ if not logger.handlers:
logger.setLevel(logging.INFO)
# ============================================
# Request/Response Models
# ============================================
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
model_path: str = Field(..., description="Model identifier or local path")
hf_token: Optional[str] = Field(None, description="HuggingFace token for gated models")
max_seq_length: int = Field(2048, ge=128, le=32768, description="Maximum sequence length")
load_in_4bit: bool = Field(True, description="Load model in 4-bit quantization")
is_lora: bool = Field(False, description="Whether this is a LoRA adapter")
class UnloadRequest(BaseModel):
"""Request to unload a model"""
model_path: str = Field(..., description="Model identifier to unload")
class GenerateRequest(BaseModel):
"""Request for text generation"""
messages: List[dict] = Field(..., description="Chat messages in OpenAI format")
system_prompt: str = Field("You are a helpful AI assistant.", description="System prompt")
temperature: float = Field(0.7, ge=0.0, le=2.0, description="Sampling temperature")
top_p: float = Field(0.9, ge=0.0, le=1.0, description="Top-p sampling")
top_k: int = Field(40, ge=1, le=100, description="Top-k sampling")
max_new_tokens: int = Field(512, ge=1, le=4096, description="Maximum tokens to generate")
repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="Repetition penalty")
image_base64: Optional[str] = Field(None, description="Base64 encoded image for vision models")
class LoadResponse(BaseModel):
"""Response after loading a model"""
status: str
model: str
display_name: str
is_vision: bool
is_lora: bool
class StatusResponse(BaseModel):
"""Current inference backend status"""
active_model: Optional[str]
is_vision: bool
loading: List[str]
loaded: List[str]
# ============================================
# Routes
# ============================================
@router.post("/load", response_model=LoadResponse)
async def load_model(request: LoadRequest):
"""
@ -147,7 +104,7 @@ async def load_model(request: LoadRequest):
)
@router.post("/unload")
@router.post("/unload", response_model=UnloadResponse)
async def unload_model(request: UnloadRequest):
"""
Unload a model from memory.
@ -156,7 +113,7 @@ async def unload_model(request: UnloadRequest):
backend = get_inference_backend()
backend.unload_model(request.model_path)
logger.info(f"Unloaded model: {request.model_path}")
return {"status": "unloaded", "model": request.model_path}
return UnloadResponse(status="unloaded", model=request.model_path)
except Exception as e:
logger.error(f"Error unloading model: {e}", exc_info=True)
@ -239,7 +196,7 @@ async def generate_stream(request: GenerateRequest):
)
@router.get("/status", response_model=StatusResponse)
@router.get("/status", response_model=InferenceStatusResponse)
async def get_status():
"""
Get current inference backend status.
@ -252,7 +209,7 @@ async def get_status():
model_info = backend.models.get(backend.active_model_name, {})
is_vision = model_info.get("is_vision", False)
return StatusResponse(
return InferenceStatusResponse(
active_model=backend.active_model_name,
is_vision=is_vision,
loading=list(getattr(backend, 'loading_models', set())),

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

@ -7,8 +7,6 @@ from fastapi import APIRouter, Depends, HTTPException, Query
from typing import List, Optional
import logging
from pydantic import BaseModel
# Add backend directory to path
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
@ -46,6 +44,8 @@ from models import (
LoRAInfo,
ModelListResponse,
)
from models.responses import LoRABaseModelResponse, VisionCheckResponse
router = APIRouter()
logger = logging.getLogger(__name__)
@ -207,7 +207,7 @@ async def scan_loras(
)
@router.get("/loras/{lora_path:path}/base-model")
@router.get("/loras/{lora_path:path}/base-model", response_model=LoRABaseModelResponse)
async def get_lora_base_model(
lora_path: str,
current_subject: str = Depends(get_current_subject),
@ -226,10 +226,10 @@ async def get_lora_base_model(
detail=f"Could not determine base model for LoRA: {lora_path}"
)
return {
"lora_path": lora_path,
"base_model": base_model
}
return LoRABaseModelResponse(
lora_path=lora_path,
base_model=base_model,
)
except HTTPException:
raise
@ -241,7 +241,7 @@ async def get_lora_base_model(
)
@router.get("/check-vision/{model_name:path}")
@router.get("/check-vision/{model_name:path}", response_model=VisionCheckResponse)
async def check_vision_model(
model_name: str,
current_subject: str = Depends(get_current_subject),
@ -254,10 +254,10 @@ async def check_vision_model(
try:
is_vision = is_vision_model(model_name)
return {
"model_name": model_name,
"is_vision": is_vision
}
return VisionCheckResponse(
model_name=model_name,
is_vision=is_vision,
)
except Exception as e:
logger.error(f"Error checking vision model: {e}", exc_info=True)

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

@ -36,6 +36,7 @@ from models import (
TrainingStatus,
TrainingProgress,
)
from models.responses import TrainingStopResponse, TrainingMetricsResponse
router = APIRouter()
logger = logging.getLogger(__name__)
@ -242,7 +243,7 @@ async def start_training(
)
@router.post("/stop")
@router.post("/stop", response_model=TrainingStopResponse)
async def stop_training(
current_subject: str = Depends(get_current_subject),
):
@ -253,18 +254,18 @@ async def stop_training(
backend = get_training_backend()
if not backend.is_training_active():
return {
"status": "idle",
"message": "No training job is currently running"
}
return TrainingStopResponse(
status="idle",
message="No training job is currently running"
)
# Call backend stop method
backend.stop_training()
return {
"status": "stopped",
"message": "Training job stopped successfully"
}
return TrainingStopResponse(
status="stopped",
message="Training job stopped successfully"
)
except Exception as e:
logger.error(f"Error stopping training: {e}", exc_info=True)
@ -353,7 +354,7 @@ async def get_training_status(
)
@router.get("/metrics")
@router.get("/metrics", response_model=TrainingMetricsResponse)
async def get_training_metrics(
current_subject: str = Depends(get_current_subject),
):
@ -373,15 +374,14 @@ async def get_training_metrics(
current_lr = lr_history[-1] if lr_history else None
current_step = step_history[-1] if step_history else None
# Keep metrics as a simple JSON payload instead of a Pydantic model
return {
"loss_history": loss_history,
"lr_history": lr_history,
"step_history": step_history,
"current_loss": current_loss,
"current_lr": current_lr,
"current_step": current_step,
}
return TrainingMetricsResponse(
loss_history=loss_history,
lr_history=lr_history,
step_history=step_history,
current_loss=current_loss,
current_lr=current_lr,
current_step=current_step,
)
except Exception as e:
logger.error(f"Error getting training metrics: {e}", exc_info=True)