123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
"""Exporta model.joblib a ONNX + inference.json (sin depender de services-count)."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import joblib
|
|
from skl2onnx import convert_sklearn
|
|
from skl2onnx.common.data_types import FloatTensorType
|
|
|
|
# Alineado con services-count-facturas-medicas-ml v1.1.0 y pipeline Go.
|
|
FEATURE_COLUMNS = [
|
|
"first_token_is_code",
|
|
"first_token_is_service_code",
|
|
"has_service_table_pattern",
|
|
"has_header_keyword",
|
|
"currency_at_row_end",
|
|
"first_token_x_norm",
|
|
"is_likely_continuation",
|
|
"currency_count",
|
|
"integer_count",
|
|
"x_span_norm",
|
|
"y_norm",
|
|
"token_count",
|
|
"special_char_ratio",
|
|
"digit_ratio",
|
|
]
|
|
|
|
DEFAULT_OCR = {"dpi": 300, "lang": "spa"}
|
|
DEFAULT_DBSCAN = {"eps": 10, "min_samples": 1}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
root = Path(__file__).resolve().parents[1]
|
|
parser = argparse.ArgumentParser(description="Exportar Random Forest joblib a ONNX")
|
|
parser.add_argument(
|
|
"--joblib",
|
|
type=Path,
|
|
required=True,
|
|
help="Ruta al model.joblib (ej. models/v1.1.0/model.joblib)",
|
|
)
|
|
parser.add_argument(
|
|
"--out",
|
|
type=Path,
|
|
default=root / "models" / "latest",
|
|
help="Directorio de salida para model.onnx e inference.json",
|
|
)
|
|
parser.add_argument(
|
|
"--metadata",
|
|
type=Path,
|
|
default=None,
|
|
help="model_metadata.json opcional para pipeline_version y params OCR/DBSCAN",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def load_metadata(path: Path | None) -> dict:
|
|
if path is None or not path.exists():
|
|
return {}
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if not args.joblib.exists():
|
|
raise SystemExit(f"No se encontró: {args.joblib}")
|
|
|
|
bundle = joblib.load(args.joblib)
|
|
model = bundle["model"]
|
|
threshold = float(bundle.get("threshold", 0.5))
|
|
meta = load_metadata(args.metadata)
|
|
|
|
bundle_columns = list(bundle.get("feature_columns") or FEATURE_COLUMNS)
|
|
if bundle_columns != FEATURE_COLUMNS:
|
|
raise SystemExit(
|
|
f"feature_columns del joblib no coinciden con el pipeline Go.\n"
|
|
f" joblib: {bundle_columns}\n"
|
|
f" esperado: {FEATURE_COLUMNS}"
|
|
)
|
|
|
|
n_features = int(getattr(model, "n_features_in_", len(FEATURE_COLUMNS)))
|
|
if n_features != len(FEATURE_COLUMNS):
|
|
raise SystemExit(
|
|
f"El model.joblib tiene {n_features} features pero el pipeline Go produce "
|
|
f"{len(FEATURE_COLUMNS)}. Reentrena o usa un joblib compatible."
|
|
)
|
|
|
|
initial_type = [("float_input", FloatTensorType([None, len(FEATURE_COLUMNS)]))]
|
|
options = {id(model): {"zipmap": False}}
|
|
onnx_model = convert_sklearn(
|
|
model,
|
|
initial_types=initial_type,
|
|
options=options,
|
|
target_opset=12,
|
|
)
|
|
|
|
out_dir = args.out
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
onnx_path = out_dir / "model.onnx"
|
|
meta_path = out_dir / "inference.json"
|
|
|
|
with open(onnx_path, "wb") as f:
|
|
f.write(onnx_model.SerializeToString())
|
|
|
|
inference = {
|
|
"threshold": threshold,
|
|
"feature_columns": FEATURE_COLUMNS,
|
|
"pipeline_version": meta.get("pipeline_version", bundle.get("pipeline_version", "unknown")),
|
|
"ocr": meta.get("ocr", DEFAULT_OCR),
|
|
"dbscan": meta.get("dbscan", DEFAULT_DBSCAN),
|
|
}
|
|
meta_path.write_text(json.dumps(inference, indent=2), encoding="utf-8")
|
|
|
|
print(f"ONNX: {onnx_path}")
|
|
print(f"Meta: {meta_path}")
|
|
print(f"Umbral: {threshold}")
|
|
print(f"Features: {len(FEATURE_COLUMNS)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|