diff --git a/ai_mouse/server/routes_train.py b/ai_mouse/server/routes_train.py index 14a66e1..2c4d511 100644 --- a/ai_mouse/server/routes_train.py +++ b/ai_mouse/server/routes_train.py @@ -23,13 +23,17 @@ router = APIRouter() # --------------------------------------------------------------------------- -def _paths() -> tuple[Path, Path]: +def _paths() -> tuple[Path, Path, Path]: data_dir = get_data_dir() - return data_dir / "traces.jsonl", data_dir / "models_v2" + return ( + data_dir / "traces.jsonl", + data_dir / "models_v2", + data_dir / "models_v2_pretrained", + ) def _trace_count() -> int: - traces_path, _ = _paths() + traces_path, _, _ = _paths() if not traces_path.exists(): return 0 return sum( @@ -40,7 +44,7 @@ def _trace_count() -> int: def _model_trained() -> bool: - _, models_dir = _paths() + _, models_dir, _ = _paths() return (models_dir / "flow_model.pt").exists() @@ -75,16 +79,31 @@ async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]: async def run_training_async() -> None: from ai_mouse.trainer import train - traces_path, models_dir = _paths() + 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)})