yellowdog/ai_mouse
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(server): auto-resume from pretrained checkpoint when available

Browse Source
This commit is contained in:
Huang Qi
2026-05-10 13:17:24 +08:00
parent 1d9dff4f2a
commit b01f6512c0
1 changed files with 24 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
ai_mouse/server/routes_train.py
View File

@@ -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)})