File size: 4,620 Bytes
b2f5c42 f046303 b2f5c42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """Lightweight ML route/risk router for EmpathRAG Core.
This module deliberately uses small scikit-learn models so the demo can start
without GPU, internet, or heavyweight transformer loading. Hard safety policy
still owns final crisis decisions; ML routing is advisory with confidence.
"""
from __future__ import annotations
from dataclasses import dataclass
import pickle
from pathlib import Path
from typing import Any
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from .v2_schema import SafetyTier, SupportRoute
DEFAULT_MODEL_DIR = Path("models/router")
ROUTE_MODEL_PATH = DEFAULT_MODEL_DIR / "route_classifier.pkl"
TIER_MODEL_PATH = DEFAULT_MODEL_DIR / "tier_classifier.pkl"
@dataclass(frozen=True)
class MLRoutePrediction:
route_label: str
safety_tier: str
route_confidence: float
tier_confidence: float
model_available: bool
used_ml: bool
reason: str
def build_text_classifier() -> Pipeline:
return Pipeline(
steps=[
("tfidf", TfidfVectorizer(ngram_range=(1, 2), min_df=1)),
("clf", LogisticRegression(max_iter=1000, class_weight="balanced")),
]
)
def train_classifier(texts: list[str], labels: list[str]) -> Pipeline:
model = build_text_classifier()
model.fit(texts, labels)
return model
def save_models(route_model: Pipeline, tier_model: Pipeline, model_dir: Path = DEFAULT_MODEL_DIR) -> None:
model_dir.mkdir(parents=True, exist_ok=True)
with (model_dir / ROUTE_MODEL_PATH.name).open("wb") as handle:
pickle.dump(route_model, handle)
with (model_dir / TIER_MODEL_PATH.name).open("wb") as handle:
pickle.dump(tier_model, handle)
def load_models(model_dir: Path = DEFAULT_MODEL_DIR) -> tuple[Pipeline | None, Pipeline | None]:
route_path = model_dir / ROUTE_MODEL_PATH.name
tier_path = model_dir / TIER_MODEL_PATH.name
if not route_path.exists() or not tier_path.exists():
return None, None
with route_path.open("rb") as handle:
route_model = pickle.load(handle)
with tier_path.open("rb") as handle:
tier_model = pickle.load(handle)
return route_model, tier_model
class MLRouter:
def __init__(self, model_dir: Path = DEFAULT_MODEL_DIR, min_confidence: float = 0.35):
self.model_dir = model_dir
self.min_confidence = min_confidence
self.route_model, self.tier_model = load_models(model_dir)
@property
def available(self) -> bool:
return self.route_model is not None and self.tier_model is not None
def predict(
self,
text: str,
fallback_route: SupportRoute | str,
fallback_tier: SafetyTier | str,
) -> MLRoutePrediction:
fallback_route_value = fallback_route.value if isinstance(fallback_route, SupportRoute) else str(fallback_route)
fallback_tier_value = fallback_tier.value if isinstance(fallback_tier, SafetyTier) else str(fallback_tier)
if not self.available:
return MLRoutePrediction(
route_label=fallback_route_value,
safety_tier=fallback_tier_value,
route_confidence=0.0,
tier_confidence=0.0,
model_available=False,
used_ml=False,
reason="model_artifacts_missing",
)
route_label, route_conf = _predict_one(self.route_model, text)
tier_label, tier_conf = _predict_one(self.tier_model, text)
if min(route_conf, tier_conf) < self.min_confidence:
return MLRoutePrediction(
route_label=fallback_route_value,
safety_tier=fallback_tier_value,
route_confidence=route_conf,
tier_confidence=tier_conf,
model_available=True,
used_ml=False,
reason="low_confidence_fallback",
)
return MLRoutePrediction(
route_label=route_label,
safety_tier=tier_label,
route_confidence=route_conf,
tier_confidence=tier_conf,
model_available=True,
used_ml=True,
reason="ml_prediction",
)
def _predict_one(model: Any, text: str) -> tuple[str, float]:
label = str(model.predict([text])[0])
if hasattr(model, "predict_proba"):
probs = model.predict_proba([text])[0]
classes = list(model.classes_)
confidence = float(probs[classes.index(label)])
else:
confidence = 1.0
return label, confidence
|