diff --git a/ocr-service/main.py b/ocr-service/main.py index 16b2d955..754ed75b 100644 --- a/ocr-service/main.py +++ b/ocr-service/main.py @@ -305,20 +305,25 @@ def _parse_best_checkpoint(checkpoint_dir: str) -> tuple[float | None, int]: def _find_best_model(checkpoint_dir: str) -> str | None: - """Return the checkpoint file with the highest validation metric, or any model file.""" - pattern = re.compile(r"checkpoint_(\d+)-([0-9.]+)\.(ckpt|mlmodel)$") + """Return the best final model file produced by ketos train. + + With --weights-format coreml, ketos writes ``best_.mlmodel``. + Falls back to any .mlmodel in the directory. + """ + # Prefer the named best file (e.g. best_0.8256.mlmodel or best_0.8256.safetensors) + best_pattern = re.compile(r"best_([0-9.]+)\.(mlmodel|safetensors)$") best_acc: float | None = None best_path: str | None = None for fname in os.listdir(checkpoint_dir): - m = pattern.match(fname) + m = best_pattern.match(fname) if m: - acc = float(m.group(2)) + acc = float(m.group(1)) if best_acc is None or acc > best_acc: best_acc = acc best_path = os.path.join(checkpoint_dir, fname) if best_path: return best_path - # Fallback: any .mlmodel file in the directory + # Fallback: any .mlmodel file for fname in os.listdir(checkpoint_dir): if fname.endswith(".mlmodel"): return os.path.join(checkpoint_dir, fname) @@ -369,6 +374,7 @@ async def train_model( "ketos", "--workers", "0", "--device", "cpu", "--threads", "2", "train", "-f", "page", + "--weights-format", "coreml", "-o", checkpoint_dir, "-q", "fixed", "-N", "10",