mirror of
https://github.com/napnap75/multiarch-docker-images.git
synced 2026-11-11 00:22:12 +01:00
125 lines
5.0 KiB
Python
125 lines
5.0 KiB
Python
"""Tests for full-set model retraining and shared additional-field extraction."""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from app.model_manager import BASE_MODELS, list_models, train_all_models
|
|
from app.ml import LocalMLPipeline, ModelUnavailableError
|
|
|
|
|
|
def validated_records():
|
|
return [
|
|
SimpleNamespace(
|
|
extracted_text="Invoice number INV-100 for account AC-10",
|
|
optional_fields={"invoice_number": "INV-100", "account_number": "AC-10"},
|
|
emitting_company="Company A",
|
|
document_type="invoice",
|
|
template_name="invoice_template",
|
|
period_format="monthly",
|
|
),
|
|
SimpleNamespace(
|
|
extracted_text="Invoice number INV-200",
|
|
optional_fields={"invoice_number": "INV-200"},
|
|
emitting_company="Company A",
|
|
document_type="invoice",
|
|
template_name="invoice_template",
|
|
period_format="quarterly",
|
|
),
|
|
SimpleNamespace(
|
|
extracted_text="Account statement for AC-30",
|
|
optional_fields={"account_number": "AC-30"},
|
|
emitting_company="Company B",
|
|
document_type="statement",
|
|
template_name="statement_template",
|
|
period_format="yearly",
|
|
),
|
|
]
|
|
|
|
|
|
class FullModelRetrainingTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp_dir.name)
|
|
self.env = patch.dict(os.environ, {"MODEL_DIR": str(self.root)})
|
|
self.env.start()
|
|
self.addCleanup(self.env.stop)
|
|
|
|
def tearDown(self):
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_retraining_replaces_all_models_without_history(self):
|
|
old_artifact = self.root / "sets" / "old-version" / "field_document_type.joblib"
|
|
old_artifact.parent.mkdir(parents=True)
|
|
old_artifact.write_bytes(b"old artifact")
|
|
(self.root / "active_models.json").write_text(
|
|
json.dumps({"field_document_type": "old-version"}), encoding="utf-8"
|
|
)
|
|
|
|
progress = []
|
|
result = train_all_models(validated_records(), lambda value, message: progress.append((value, message)))
|
|
|
|
self.assertEqual(result["status"], "completed")
|
|
self.assertEqual(result["count"], len(BASE_MODELS))
|
|
active = json.loads((self.root / "active_models.json").read_text())
|
|
self.assertEqual(set(active), set(BASE_MODELS))
|
|
self.assertEqual(len(set(active.values())), 1)
|
|
for name, version in active.items():
|
|
self.assertTrue((self.root / "sets" / version / f"{name}.joblib").is_file())
|
|
self.assertFalse(old_artifact.exists())
|
|
self.assertEqual(
|
|
{path.name.removesuffix(".joblib") for path in (self.root / "sets" / next(iter(active.values()))).glob("*.joblib")},
|
|
set(BASE_MODELS),
|
|
)
|
|
self.assertEqual(len(list((self.root / "sets").glob("*"))), 1)
|
|
self.assertFalse((self.root / "versions").exists())
|
|
self.assertEqual(progress[-1][0], 100)
|
|
|
|
def test_failed_retraining_keeps_the_previous_active_set(self):
|
|
old_generation = "existing-set"
|
|
old_artifact = self.root / "sets" / old_generation / "field_document_type.joblib"
|
|
old_artifact.parent.mkdir(parents=True)
|
|
old_artifact.write_bytes(b"previous active artifact")
|
|
active_before = {"field_document_type": old_generation}
|
|
(self.root / "active_models.json").write_text(json.dumps(active_before), encoding="utf-8")
|
|
|
|
with self.assertRaises(ValueError):
|
|
train_all_models([])
|
|
|
|
self.assertEqual(
|
|
json.loads((self.root / "active_models.json").read_text()),
|
|
active_before,
|
|
)
|
|
self.assertTrue(old_artifact.is_file())
|
|
self.assertFalse(any(path.name.startswith(".staging-") for path in self.root.iterdir()))
|
|
|
|
def test_inventory_has_no_version_history_or_emission_date_model(self):
|
|
result = list_models()
|
|
names = {model["name"] for model in result["models"]}
|
|
self.assertNotIn("field_emission_date", names)
|
|
self.assertTrue(all("versions" not in model for model in result["models"]))
|
|
|
|
def test_inference_requires_a_complete_active_model_set(self):
|
|
pipeline = LocalMLPipeline(str(self.root), 0.8)
|
|
with self.assertRaises(ModelUnavailableError):
|
|
pipeline._load("field_document_type")
|
|
|
|
def test_shared_additional_value_model_extracts_multiple_fields(self):
|
|
train_all_models(validated_records())
|
|
pipeline = LocalMLPipeline(str(self.root), 0.0)
|
|
predictions = pipeline.predict_additional_values(
|
|
["invoice_number", "account_number"],
|
|
"Invoice number INV-200 for account AC-20",
|
|
)
|
|
self.assertEqual(set(predictions), {"invoice_number", "account_number"})
|
|
self.assertTrue(predictions["invoice_number"][0])
|
|
self.assertTrue(predictions["account_number"][0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|