mirror of
https://github.com/napnap75/multiarch-docker-images.git
synced 2026-11-13 09:31:46 +01:00
Add local ML models and document suggestions (#13)
This commit is contained in:
@@ -0,0 +1,363 @@
|
||||
"""Offline ML inference for document field extraction and routing decisions.
|
||||
|
||||
Model artifacts are trusted, locally trained joblib objects. No rule-based or
|
||||
configuration-based classifier is used when an artifact is unavailable.
|
||||
Each single-label artifact exposes predict([sample]) and, optionally,
|
||||
predict_proba([sample]). The group selector returns an iterable of group labels
|
||||
for a (company, document_type) sample.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModelUnavailableError(RuntimeError):
|
||||
"""Raised when a required local ML model is missing or invalid."""
|
||||
|
||||
|
||||
class AmbiguousPredictionError(RuntimeError):
|
||||
"""Raised when an ML prediction does not meet the configured confidence."""
|
||||
|
||||
|
||||
class LocalMLPipeline:
|
||||
"""Load local model artifacts and provide confidence-gated predictions."""
|
||||
|
||||
def __init__(self, model_dir: str | None = None, confidence_threshold: float | None = None):
|
||||
self.model_dir = Path(model_dir or os.environ.get("MODEL_DIR", "./models"))
|
||||
self.confidence_threshold = (
|
||||
float(confidence_threshold)
|
||||
if confidence_threshold is not None
|
||||
else float(os.environ.get("CONFIDENCE_THRESHOLD", "0.8"))
|
||||
)
|
||||
self._models: dict[str, Any] = {}
|
||||
self._model_versions: dict[str, str] = {}
|
||||
|
||||
def _load(self, name: str) -> Any:
|
||||
active_path = self.model_dir / "active_models.json"
|
||||
try:
|
||||
active = __import__("json").loads(active_path.read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError):
|
||||
active = {}
|
||||
version = active.get(name)
|
||||
if not version:
|
||||
raise ModelUnavailableError(
|
||||
f"No active model artifact for {name}; retrain all models from Maintenance."
|
||||
)
|
||||
if name in self._models and self._model_versions.get(name) == version:
|
||||
return self._models[name]
|
||||
path = self.model_dir / "sets" / version / f"{name}.joblib"
|
||||
if not path.is_file():
|
||||
# Keep the pre-retraining model layout usable until the first full replacement.
|
||||
path = self.model_dir / "versions" / name / f"{version}.joblib"
|
||||
if not path.is_file():
|
||||
raise ModelUnavailableError(
|
||||
f"Active model artifact is missing for {name} version {version}."
|
||||
)
|
||||
try:
|
||||
import joblib
|
||||
model = joblib.load(path)
|
||||
except Exception as exc:
|
||||
raise ModelUnavailableError(f"Could not load ML model {path.name}: {exc}") from exc
|
||||
self._models[name] = model
|
||||
self._model_versions[name] = version
|
||||
return model
|
||||
|
||||
def _predict_label(self, name: str, sample: Any) -> tuple[str, float]:
|
||||
model = self._load(name)
|
||||
try:
|
||||
predictions = model.predict([sample])
|
||||
logger.debug("ML model %s predicted %s for sample=%r", name, predictions, sample)
|
||||
if len(predictions) != 1:
|
||||
raise ValueError("Expected exactly one prediction")
|
||||
label = predictions[0]
|
||||
if not isinstance(label, str) or not label.strip():
|
||||
raise ValueError("The model returned an empty or non-string label")
|
||||
confidence = 0.0
|
||||
if hasattr(model, "predict_proba"):
|
||||
probabilities = model.predict_proba([sample])
|
||||
confidence = float(max(probabilities[0]))
|
||||
elif hasattr(model, "decision_function"):
|
||||
# Decision scores are not calibrated probabilities; do not
|
||||
# misrepresent them as confidence.
|
||||
raise ModelUnavailableError(
|
||||
f"Model {name} must expose calibrated predict_proba() scores"
|
||||
)
|
||||
else:
|
||||
raise ModelUnavailableError(
|
||||
f"Model {name} must expose predict_proba() for confidence gating"
|
||||
)
|
||||
if confidence < self.confidence_threshold:
|
||||
raise AmbiguousPredictionError(
|
||||
f"Model {name} confidence {confidence:.3f} is below "
|
||||
f"threshold {self.confidence_threshold:.3f}"
|
||||
)
|
||||
return label.strip(), confidence
|
||||
except (ModelUnavailableError, AmbiguousPredictionError) as exc:
|
||||
logger.debug("ML model %s failed for sample=%r: %s", name, sample, exc)
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ModelUnavailableError(f"ML inference failed for {name}: {exc}") from exc
|
||||
|
||||
def predict_labels(
|
||||
self,
|
||||
fields: dict[str, tuple[str, Any]],
|
||||
) -> tuple[dict[str, tuple[str, float]], dict[str, Exception]]:
|
||||
"""Predict independent fields without letting one failure block the others."""
|
||||
predictions: dict[str, tuple[str, float]] = {}
|
||||
errors: dict[str, Exception] = {}
|
||||
for field, (model_name, sample) in fields.items():
|
||||
try:
|
||||
predictions[field] = self._predict_label(model_name, sample)
|
||||
except (ModelUnavailableError, AmbiguousPredictionError) as exc:
|
||||
errors[field] = exc
|
||||
return predictions, errors
|
||||
|
||||
def extract_initial_fields(self, text: str) -> dict[str, tuple[str, float]]:
|
||||
"""Predict common fields directly from full document text."""
|
||||
fields, errors = self.predict_labels({
|
||||
field: (f"field_{field}", text)
|
||||
for field in ("emitting_company", "document_type")
|
||||
})
|
||||
if errors:
|
||||
raise next(iter(errors.values()))
|
||||
logger.debug("Extracted initial fields %r from text=%r", fields, text)
|
||||
return fields
|
||||
|
||||
def predict_additional_groups(self, company: str, document_type: str) -> tuple[list[str], float]:
|
||||
"""Predict applicable additional-field groups from company and type."""
|
||||
model = self._load("additional_field_groups")
|
||||
sample = f"{company}\n{document_type}"
|
||||
try:
|
||||
predictions = model.predict([sample])
|
||||
if len(predictions) != 1:
|
||||
raise ValueError("Expected exactly one group prediction")
|
||||
predicted = predictions[0]
|
||||
if isinstance(predicted, str):
|
||||
groups = [predicted] if predicted.strip() else []
|
||||
else:
|
||||
groups = [str(value).strip() for value in predicted if str(value).strip()]
|
||||
if not groups:
|
||||
raise AmbiguousPredictionError("ML model predicted no additional-field groups")
|
||||
if not hasattr(model, "predict_proba"):
|
||||
raise ModelUnavailableError(
|
||||
"Model additional_field_groups must expose calibrated predict_proba() scores"
|
||||
)
|
||||
probabilities = model.predict_proba([sample])
|
||||
first = probabilities[0]
|
||||
if hasattr(first, "toarray"):
|
||||
first = first.toarray().ravel()
|
||||
try:
|
||||
confidence = float(max(first))
|
||||
except (TypeError, ValueError):
|
||||
confidence = float(getattr(model, "confidence_threshold", 0.0))
|
||||
if confidence < self.confidence_threshold:
|
||||
raise AmbiguousPredictionError(
|
||||
f"Additional-field group confidence {confidence:.3f} is below "
|
||||
f"threshold {self.confidence_threshold:.3f}"
|
||||
)
|
||||
logger.debug('Predicted additional-field group %r (confidence=%.3f)', groups[0], confidence)
|
||||
return groups, confidence
|
||||
except (ModelUnavailableError, AmbiguousPredictionError):
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ModelUnavailableError(f"ML inference failed for additional_field_groups: {exc}") from exc
|
||||
|
||||
def predict_additional_values(
|
||||
self,
|
||||
field_names: list[str],
|
||||
text: str,
|
||||
*,
|
||||
allow_partial: bool = False,
|
||||
) -> dict[str, tuple[str, float]]:
|
||||
"""Extract requested fields, optionally omitting individual low-confidence results."""
|
||||
model = self._load("additional_field_values")
|
||||
if not hasattr(model, "predict_values"):
|
||||
raise ModelUnavailableError(
|
||||
"The active additional_field_values artifact does not support shared extraction; retrain it."
|
||||
)
|
||||
try:
|
||||
predictions = model.predict_values(text, field_names)
|
||||
results: dict[str, tuple[str, float]] = {}
|
||||
for field_name in field_names:
|
||||
value, confidence = predictions[field_name]
|
||||
if not value or confidence < self.confidence_threshold:
|
||||
error = AmbiguousPredictionError(
|
||||
f"Additional field '{field_name}' confidence {confidence:.3f} is below "
|
||||
f"threshold {self.confidence_threshold:.3f}"
|
||||
)
|
||||
if allow_partial:
|
||||
logger.debug("Omitting low-confidence additional field %s: %s", field_name, error)
|
||||
continue
|
||||
raise error
|
||||
results[field_name] = (str(value), float(confidence))
|
||||
logger.debug('Predicted additional field %s=%r (confidence=%.3f)', field_name, value, confidence)
|
||||
return results
|
||||
except (ModelUnavailableError, AmbiguousPredictionError):
|
||||
raise
|
||||
except ValueError as exc:
|
||||
raise ModelUnavailableError(str(exc)) from exc
|
||||
except Exception as exc:
|
||||
raise ModelUnavailableError(f"ML inference failed for additional_field_values: {exc}") from exc
|
||||
|
||||
def predict_template_and_period_format(
|
||||
self,
|
||||
emitting_company: str,
|
||||
document_type: str,
|
||||
additional_fields: dict[str, str],
|
||||
) -> tuple[str, float, str, float]:
|
||||
"""Predict template and period format from company, type, and extracted fields."""
|
||||
fields_text = "\n".join(f"{key}: {value}" for key, value in sorted(additional_fields.items()))
|
||||
sample = f"{emitting_company}\n{document_type}\n{fields_text}"
|
||||
predictions, errors = self.predict_labels({
|
||||
"template_name": ("template_name", sample),
|
||||
"period_format": ("period_format", sample),
|
||||
})
|
||||
if errors:
|
||||
raise next(iter(errors.values()))
|
||||
template, template_confidence = predictions["template_name"]
|
||||
period_format, period_confidence = predictions["period_format"]
|
||||
logger.debug(
|
||||
"Predicted template=%s (confidence=%.3f), period_format=%s (confidence=%.3f) "
|
||||
"for company=%r, document_type=%r, additional_fields=%r",
|
||||
template,
|
||||
template_confidence,
|
||||
period_format,
|
||||
period_confidence,
|
||||
emitting_company,
|
||||
document_type,
|
||||
additional_fields,
|
||||
)
|
||||
return template, template_confidence, period_format, period_confidence
|
||||
|
||||
@staticmethod
|
||||
def match_configured_additional_value(
|
||||
text: str,
|
||||
predicted_value: str,
|
||||
configured_values: list[str],
|
||||
) -> str | None:
|
||||
"""Resolve an extracted span to a complete, unambiguous configured value."""
|
||||
matched_values = []
|
||||
for raw_value in configured_values:
|
||||
value = str(raw_value).strip()
|
||||
words = value.split()
|
||||
if not words:
|
||||
continue
|
||||
pattern = r"(?<!\w)" + r"\s+".join(re.escape(word) for word in words) + r"(?!\w)"
|
||||
if re.search(pattern, text, flags=re.IGNORECASE):
|
||||
matched_values.append(value)
|
||||
|
||||
if len(matched_values) == 1:
|
||||
return matched_values[0]
|
||||
if not matched_values:
|
||||
return None
|
||||
|
||||
predicted_words = set(re.findall(r"\w+", predicted_value.casefold()))
|
||||
scores = {
|
||||
value: len(predicted_words & set(re.findall(r"\w+", value.casefold())))
|
||||
for value in matched_values
|
||||
}
|
||||
best_score = max(scores.values())
|
||||
best_matches = [value for value, score in scores.items() if score == best_score]
|
||||
return best_matches[0] if best_score and len(best_matches) == 1 else None
|
||||
|
||||
def suggest_document_fields(
|
||||
self,
|
||||
text: str,
|
||||
company: str,
|
||||
document_type: str,
|
||||
current_fields: dict[str, str],
|
||||
configured_values: dict[str, list[str]],
|
||||
) -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
|
||||
"""Collect independent model and configured-value suggestions for one document."""
|
||||
suggestions: dict[str, dict[str, Any]] = {}
|
||||
decisions = {
|
||||
"emitting_company": "skipped: model not evaluated",
|
||||
"document_type": "skipped: model not evaluated",
|
||||
"template_name": "skipped: prerequisites not met",
|
||||
"period_format": "skipped: prerequisites not met",
|
||||
"additional_fields": "skipped: groups not evaluated",
|
||||
}
|
||||
initial, initial_errors = self.predict_labels({
|
||||
"emitting_company": ("field_emitting_company", text),
|
||||
"document_type": ("field_document_type", text),
|
||||
})
|
||||
for field, (value, confidence) in initial.items():
|
||||
suggestions[field] = {"value": value, "confidence": confidence, "source": "model"}
|
||||
decisions[field] = f"suggested: {value!r} ({confidence:.3f})"
|
||||
for field, error in initial_errors.items():
|
||||
decisions[field] = f"unavailable: {error}" if isinstance(error, ModelUnavailableError) else f"ambiguous: {error}"
|
||||
|
||||
suggested_company = initial.get("emitting_company", (company, 0.0))[0]
|
||||
suggested_type = initial.get("document_type", (document_type, 0.0))[0]
|
||||
if not suggested_company or not suggested_type:
|
||||
missing = [
|
||||
name for name, value in (("emitting_company", suggested_company), ("document_type", suggested_type))
|
||||
if not value
|
||||
]
|
||||
reason = f"skipped: missing {', '.join(missing)}"
|
||||
decisions["template_name"] = reason
|
||||
decisions["period_format"] = reason
|
||||
return suggestions, decisions
|
||||
|
||||
try:
|
||||
groups, _ = self.predict_additional_groups(suggested_company, suggested_type)
|
||||
value_model = self._load("additional_field_values")
|
||||
trained_fields = getattr(value_model, "trained_fields", None)
|
||||
if trained_fields is None:
|
||||
extractable_groups = groups
|
||||
unsupported_groups: list[str] = []
|
||||
else:
|
||||
supported_fields = set(trained_fields)
|
||||
extractable_groups = [name for name in groups if name in supported_fields]
|
||||
unsupported_groups = sorted(set(groups) - supported_fields)
|
||||
decisions["additional_fields"] = (
|
||||
f"predicted={groups}; extracting={extractable_groups}; unsupported={unsupported_groups}"
|
||||
)
|
||||
extracted = self.predict_additional_values(
|
||||
extractable_groups, text, allow_partial=True
|
||||
) if extractable_groups else {}
|
||||
matched = {
|
||||
name: value
|
||||
for name, (span, _confidence) in extracted.items()
|
||||
if (value := self.match_configured_additional_value(
|
||||
text, span, configured_values.get(name, [])
|
||||
)) is not None
|
||||
}
|
||||
if matched:
|
||||
suggestions["optional_fields"] = {"value": matched, "source": "configured"}
|
||||
model_spans = {
|
||||
name: {"value": value, "confidence": confidence}
|
||||
for name, (value, confidence) in extracted.items()
|
||||
}
|
||||
decisions["additional_fields"] += (
|
||||
f"; model_spans={model_spans}; configured_matches={matched}; "
|
||||
f"omitted_low_confidence={sorted(set(extractable_groups) - set(extracted))}; "
|
||||
f"omitted_without_configured_match={sorted(set(extracted) - set(matched))}"
|
||||
)
|
||||
inference_fields = {**matched, **current_fields}
|
||||
fields_text = "\n".join(
|
||||
f"{key}: {value}" for key, value in sorted(inference_fields.items())
|
||||
)
|
||||
sample = f"{suggested_company}\n{suggested_type}\n{fields_text}"
|
||||
final, final_errors = self.predict_labels({
|
||||
"template_name": ("template_name", sample),
|
||||
"period_format": ("period_format", sample),
|
||||
})
|
||||
for field, (value, confidence) in final.items():
|
||||
suggestions[field] = {"value": value, "confidence": confidence, "source": "model"}
|
||||
decisions[field] = f"suggested: {value!r} ({confidence:.3f})"
|
||||
for field, error in final_errors.items():
|
||||
decisions[field] = f"unavailable: {error}" if isinstance(error, ModelUnavailableError) else f"ambiguous: {error}"
|
||||
except (ModelUnavailableError, AmbiguousPredictionError) as exc:
|
||||
reason = f"skipped: additional-field inference failed ({exc})"
|
||||
decisions["additional_fields"] = f"unavailable: {exc}"
|
||||
decisions["template_name"] = reason
|
||||
decisions["period_format"] = reason
|
||||
return suggestions, decisions
|
||||
Reference in New Issue
Block a user