132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
# ai_mouse/server/routes_train.py
|
|
"""Training and status routes."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import AsyncGenerator
|
|
|
|
from fastapi import APIRouter
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
|
|
from .deps import get_data_dir
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _paths() -> tuple[Path, Path, Path]:
|
|
data_dir = get_data_dir()
|
|
return (
|
|
data_dir / "traces.jsonl",
|
|
data_dir / "models_v2",
|
|
data_dir / "models_v2_pretrained",
|
|
)
|
|
|
|
|
|
def _trace_count() -> int:
|
|
traces_path, _, _ = _paths()
|
|
if not traces_path.exists():
|
|
return 0
|
|
return sum(
|
|
1
|
|
for line in traces_path.read_text(encoding="utf-8").splitlines()
|
|
if line.strip()
|
|
)
|
|
|
|
|
|
def _model_trained() -> bool:
|
|
_, models_dir, _ = _paths()
|
|
return (models_dir / "flow_model.pt").exists()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request models
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TrainRequest(BaseModel):
|
|
epochs: int = 200
|
|
data_path: str | None = None
|
|
output_dir: str | None = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Routes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@router.get("/status")
|
|
def get_status() -> dict:
|
|
return {"trace_count": _trace_count(), "model_trained": _model_trained()}
|
|
|
|
|
|
async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]:
|
|
"""Run training in a thread via asyncio.to_thread, yield SSE events via asyncio.Queue."""
|
|
queue: asyncio.Queue[dict] = asyncio.Queue()
|
|
|
|
def callback(msg: dict) -> None:
|
|
queue.put_nowait(msg)
|
|
|
|
async def run_training_async() -> None:
|
|
from tools.trainer import train
|
|
|
|
traces_path, models_dir, pretrained_dir = _paths()
|
|
data_path = Path(req.data_path) if req.data_path else traces_path
|
|
output_dir = Path(req.output_dir) if req.output_dir else models_dir
|
|
|
|
# Auto-detect pretrained checkpoint and switch to fine-tune mode
|
|
resume_from: Path | None = None
|
|
effective_lr = 3e-4
|
|
if (pretrained_dir / "flow_model.pt").exists():
|
|
resume_from = pretrained_dir
|
|
effective_lr = 1e-5 # fine-tune lr
|
|
logger.info("Detected pretrained checkpoint, fine-tuning at lr=%g", effective_lr)
|
|
queue.put_nowait({
|
|
"info": f"Detected pretrained checkpoint at {pretrained_dir.name}, "
|
|
f"running fine-tune at lr={effective_lr}",
|
|
})
|
|
|
|
try:
|
|
await asyncio.to_thread(
|
|
train,
|
|
data_path=data_path,
|
|
output_dir=output_dir,
|
|
epochs=req.epochs,
|
|
lr=effective_lr,
|
|
progress_callback=callback,
|
|
resume_from=resume_from,
|
|
)
|
|
except Exception as exc: # noqa: BLE001
|
|
queue.put_nowait({"error": str(exc)})
|
|
|
|
task = asyncio.create_task(run_training_async())
|
|
|
|
while True:
|
|
msg = await queue.get()
|
|
yield f"data: {json.dumps(msg)}\n\n"
|
|
if msg.get("done") or msg.get("error"):
|
|
break
|
|
|
|
await task
|
|
|
|
|
|
@router.post("/train")
|
|
async def train_model(req: TrainRequest) -> StreamingResponse:
|
|
return StreamingResponse(
|
|
_train_sse_generator(req),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|