mlsort 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

mlsort/api.py

@@ -43,6 +43,10 @@ def _ensure_thresholds(path: str) -> Thresholds:
     return th
 
 
+def _use_fast_model() -> bool:
+    return os.environ.get("MLSORT_USE_FAST_MODEL", "0").lower() in {"1", "true", "yes", "on"}
+
+
 def select_algorithm(arr: Sequence[Any], thresholds_path: str | None = None, *, key: Any = None, reverse: bool = False) -> str:
     # Input validation
     try:
@@ -88,7 +92,17 @@ def select_algorithm(arr: Sequence[Any], thresholds_path: str | None = None, *,
     thr_path = thresholds_path or os.path.join(get_artifacts_dir(), "thresholds.json")
     os.makedirs(os.path.dirname(thr_path) or ".", exist_ok=True)
     th = _ensure_thresholds(thr_path)
-    algo = decide(arr, th)
+    # Large arrays: optionally use fast model
+    if _use_fast_model():
+        try:
+            from .features import estimate_properties
+            from .fast_model import predict_fast
+            props = estimate_properties(arr)
+            algo = predict_fast(props)
+        except Exception:
+            algo = decide(arr, th)
+    else:
+        algo = decide(arr, th)
     if get_env_bool("MLSORT_DEBUG", False):
         log.debug("mlsort.select algo=%s n=%d path=%s", algo, n, thr_path)
     return algo

mlsort/cli_export_forest.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import joblib
+from sklearn.ensemble import RandomForestClassifier
+
+from .model import LABELS
+
+
+def export_forest(model: RandomForestClassifier) -> dict:
+    trees = []
+    for est in model.estimators_:
+        t = est.tree_
+        nodes = []
+        for i in range(t.node_count):
+            if t.children_left[i] == -1 and t.children_right[i] == -1:
+                value = t.value[i][0].tolist()
+                nodes.append({"value": value})
+            else:
+                nodes.append({
+                    "feature": int(t.feature[i]),
+                    "threshold": float(t.threshold[i]),
+                    "left": int(t.children_left[i]),
+                    "right": int(t.children_right[i]),
+                })
+        trees.append({"nodes": nodes})
+    return {"label_names": LABELS, "trees": trees}
+
+
+def main():
+    p = argparse.ArgumentParser(description="Export sklearn RandomForest to fast JSON format")
+    p.add_argument("--model", required=True, help="Path to model.joblib")
+    p.add_argument("--out", required=True, help="Path to write forest.json")
+    args = p.parse_args()
+
+    model: RandomForestClassifier = joblib.load(args.model)
+    spec = export_forest(model)
+    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
+    Path(args.out).write_text(json.dumps(spec))
+    print(f"Wrote {args.out}")  # noqa: T201
+
+
+if __name__ == "__main__":
+    main()

mlsort/fast_model.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Dict, List, Optional
+
+from .features import to_feature_vector
+from .model import ID_TO_LABEL
+from .config import get_artifacts_dir
+
+_FAST_MODEL: Optional[Dict] = None
+_FAST_MODEL_PATH: Optional[str] = None
+
+
+def _get_default_fast_model_path() -> str:
+    return os.path.join(get_artifacts_dir(), "forest.json")
+
+
+def load_fast_model(path: Optional[str] = None) -> Dict:
+    global _FAST_MODEL, _FAST_MODEL_PATH
+    use_path = path or _get_default_fast_model_path()
+    if _FAST_MODEL is None or _FAST_MODEL_PATH != use_path:
+        with open(use_path, "r") as f:
+            _FAST_MODEL = json.load(f)
+        _FAST_MODEL_PATH = use_path
+    return _FAST_MODEL  # type: ignore[return-value]
+
+
+def _tree_predict(tree: Dict, x: List[float]) -> int:
+    nodes = tree["nodes"]
+    i = 0
+    while True:
+        node = nodes[i]
+        if "value" in node:
+            vec: List[float] = node["value"]
+            return int(max(range(len(vec)), key=lambda k: vec[k]))
+        feat = node["feature"]
+        thr = node["threshold"]
+        left = node["left"]
+        right = node["right"]
+        i = left if x[feat] <= thr else right
+
+
+def predict_fast(props: Dict[str, float], *, model_path: Optional[str] = None) -> str:
+    fm = load_fast_model(model_path)
+    x = to_feature_vector(props)
+    votes: List[int] = []
+    for tree in fm["trees"]:
+        votes.append(_tree_predict(tree, x))
+    if not votes:
+        return "timsort"
+    counts: Dict[int, int] = {}
+    for v in votes:
+        counts[v] = counts.get(v, 0) + 1
+    best_id = max(counts.items(), key=lambda kv: kv[1])[0]
+    return ID_TO_LABEL[int(best_id)]

{mlsort-0.1.0.dist-info → mlsort-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlsort
-Version: 0.1.0
+Version: 0.1.1
 Summary: ML-guided sorting backend selector with install-time benchmarking
 Author: Siddharth Chaudhary
 License: MIT License
@@ -43,9 +43,10 @@ Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy>=1.24
-Requires-Dist: scikit-learn>=1.3
-Requires-Dist: scipy>=1.10
-Requires-Dist: joblib>=1.3
+Provides-Extra: train
+Requires-Dist: scikit-learn>=1.3; extra == "train"
+Requires-Dist: scipy>=1.10; extra == "train"
+Requires-Dist: joblib>=1.3; extra == "train"
 Dynamic: license-file
 
 # mlsort

{mlsort-0.1.0.dist-info → mlsort-0.1.1.dist-info}/RECORD

@@ -1,22 +1,24 @@
 mlsort/__init__.py,sha256=49ZFRUBmCcD_YpHDLtAvb6CjCOAUoDqczL0c5pTWhPs,1121
 mlsort/algorithms.py,sha256=MgOOe9SHy9D_af7siDS4jWtuLK6alhIv8sPusqCx9qI,4475
-mlsort/api.py,sha256=T1T_ND-ybfld0FTHyjAihhNgPTcjcWgXnv50bnVnoKw,6026
+mlsort/api.py,sha256=gyxjAYYu254QJyKDU0d0ApkWdf_O-DAOjj9zI-NZuiE,6500
 mlsort/baseline.py,sha256=2nZrEY7P5QAQ8RPOxqNz47rR_WZMyd3iONOSR38u_-Y,1104
 mlsort/benchmark.py,sha256=Ez_-HOnbvzfZD0323Nv_8vGj3xhENflztD-7IOEAalo,3713
 mlsort/cli_bench_compare.py,sha256=HH1C8H8IpWWdORT81_1gsO_je7emLW9NEIpqPedDbgw,1771
 mlsort/cli_bench_install.py,sha256=g28V9TZ_b5rJIEbV9LAkXZWYZaUMYQK10npGGxTZ5jI,936
+mlsort/cli_export_forest.py,sha256=t2qpwfU85bMv1XNphjtA1u9YETM2j7WHcIeWdMctqWE,1481
 mlsort/cli_init.py,sha256=kDcnne1lLTRg5IEZVBOYiV8kkzb_5JHmTybZEvvPpKw,1721
 mlsort/cli_optimize_cutoffs.py,sha256=6De71xb93z6JScfMIduKkDe9DZ7xaG2Pg085JQe6HB0,1216
 mlsort/config.py,sha256=3Qzumm41uCvseH9LbRaDe06ffFOsJ3k3f20pM_bdjg8,966
 mlsort/data.py,sha256=HHtffrqOE15jLRxN6sc__mXBKsptT-USLxsZhq09SIc,3246
 mlsort/decision.py,sha256=YB3epa2L7Wa1faFmbyTSNjdGj0NigO2mro7K_k5IklA,1527
+mlsort/fast_model.py,sha256=DQ1Pg3i1BeYcfvIX-6L17DO8b-FcbVRU2q7KyKcaxS8,1699
 mlsort/features.py,sha256=MJOwPnC4z8VDwiA9PXBgjwSr5r9djup7ZieGzSmwMbg,5303
 mlsort/installer.py,sha256=M7Dj2lMEecNblWY4YoMBME3fjqNf8L2BolNbvLcYUxE,4904
 mlsort/model.py,sha256=OY6b_04unIIbjrjmQ4LKrG52TmB7FvUBbflcBNg2-d0,2378
 mlsort/optimize.py,sha256=7Yi6tmiJcnj_6NtDfViuxzY1nVR3zoq-sbOj0v7yEis,2945
-mlsort-0.1.0.dist-info/licenses/LICENSE,sha256=yzOA5llIyAHw7tVsir3l5NgRm1_pkvXy2r4bUFcZY0g,1076
-mlsort-0.1.0.dist-info/METADATA,sha256=m4fy4EuvfqrB-enbeFxXrEQZb0nhbErdEaqAooD05DU,5794
-mlsort-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mlsort-0.1.0.dist-info/entry_points.txt,sha256=HKRZnDWd50NuGw9uFEhAUZE6OOEKl74MT3J3yHL-1Is,218
-mlsort-0.1.0.dist-info/top_level.txt,sha256=0tl8OhYGP3bgyXuS76DsDreFASPKloMccz5pGfteKp0,7
-mlsort-0.1.0.dist-info/RECORD,,
+mlsort-0.1.1.dist-info/licenses/LICENSE,sha256=yzOA5llIyAHw7tVsir3l5NgRm1_pkvXy2r4bUFcZY0g,1076
+mlsort-0.1.1.dist-info/METADATA,sha256=_yZ_ztOpM8KorrJdFJt43p2vP234Un8Ck0QQqFRB5Pc,5870
+mlsort-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mlsort-0.1.1.dist-info/entry_points.txt,sha256=2oMHL1Z_f3U-hEwWb369v4XTwgKCkztdc1UJUPMhaDM,271
+mlsort-0.1.1.dist-info/top_level.txt,sha256=0tl8OhYGP3bgyXuS76DsDreFASPKloMccz5pGfteKp0,7
+mlsort-0.1.1.dist-info/RECORD,,

{mlsort-0.1.0.dist-info → mlsort-0.1.1.dist-info}/entry_points.txt

@@ -1,5 +1,6 @@
 [console_scripts]
 mlsort-bench-compare = mlsort.cli_bench_compare:main
 mlsort-bench-install = mlsort.cli_bench_install:main
+mlsort-export-forest = mlsort.cli_export_forest:main
 mlsort-init = mlsort.cli_init:main
 mlsort-optimize-cutoffs = mlsort.cli_optimize_cutoffs:main