From e6fc0bc18fa6cd15c58581236cbfdb066453acba Mon Sep 17 00:00:00 2001 From: erjieyong Date: Sat, 17 Feb 2024 19:05:19 +0800 Subject: [PATCH 1/3] convert string keys to int keys --- unstructured_inference/models/base.py | 24 +++++++++--------------- 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/unstructured_inference/models/base.py b/unstructured_inference/models/base.py index 26336e23..69c20cb0 100644 --- a/unstructured_inference/models/base.py +++ b/unstructured_inference/models/base.py @@ -7,25 +7,15 @@ from unstructured_inference.models.detectron2 import ( MODEL_TYPES as DETECTRON2_MODEL_TYPES, ) -from unstructured_inference.models.detectron2 import ( - UnstructuredDetectronModel, -) +from unstructured_inference.models.detectron2 import UnstructuredDetectronModel from unstructured_inference.models.detectron2onnx import ( MODEL_TYPES as DETECTRON2_ONNX_MODEL_TYPES, ) -from unstructured_inference.models.detectron2onnx import ( - UnstructuredDetectronONNXModel, -) -from unstructured_inference.models.super_gradients import ( - UnstructuredSuperGradients, -) +from unstructured_inference.models.detectron2onnx import UnstructuredDetectronONNXModel +from unstructured_inference.models.super_gradients import UnstructuredSuperGradients from unstructured_inference.models.unstructuredmodel import UnstructuredModel -from unstructured_inference.models.yolox import ( - MODEL_TYPES as YOLOX_MODEL_TYPES, -) -from unstructured_inference.models.yolox import ( - UnstructuredYoloXModel, -) +from unstructured_inference.models.yolox import MODEL_TYPES as YOLOX_MODEL_TYPES +from unstructured_inference.models.yolox import UnstructuredYoloXModel DEFAULT_MODEL = "yolox" @@ -58,6 +48,10 @@ def get_model(model_name: Optional[str] = None) -> UnstructuredModel: if initialize_param_json is not None: with open(initialize_param_json) as fp: initialize_params = json.load(fp) + label_map_int_keys = { + int(key): value for key, value in initialize_params["label_map"].items() + } + initialize_params["label_map"] = label_map_int_keys else: if model_name in DETECTRON2_MODEL_TYPES: initialize_params = DETECTRON2_MODEL_TYPES[model_name] From cae2001cc3d2b6cb849ed4afc58f785a7284e9f4 Mon Sep 17 00:00:00 2001 From: erjieyong Date: Sat, 17 Feb 2024 19:26:50 +0800 Subject: [PATCH 2/3] update test --- .../models/test_model.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/test_unstructured_inference/models/test_model.py b/test_unstructured_inference/models/test_model.py index fc7c18aa..57bf84f8 100644 --- a/test_unstructured_inference/models/test_model.py +++ b/test_unstructured_inference/models/test_model.py @@ -1,3 +1,4 @@ +import json from typing import Any from unittest import mock @@ -45,7 +46,9 @@ def test_model_initializes_once(): from unstructured_inference.inference import layout with mock.patch.dict(models.model_class_map, {"yolox": MockModel}), mock.patch.object( - models, "models", {} + models, + "models", + {}, ): doc = layout.DocumentLayout.from_file("sample-docs/loremipsum.pdf") doc.pages[0].detection_model.initializer.assert_called_once() @@ -143,23 +146,32 @@ def test_env_variables_override_default_model(monkeypatch): # args, we should get back the model the env var calls for monkeypatch.setattr(models, "models", {}) with mock.patch.dict( - models.os.environ, {"UNSTRUCTURED_DEFAULT_MODEL_NAME": "checkbox"} + models.os.environ, + {"UNSTRUCTURED_DEFAULT_MODEL_NAME": "checkbox"}, ), mock.patch.dict(models.model_class_map, {"checkbox": MockModel}): model = models.get_model() assert isinstance(model, MockModel) -def test_env_variables_override_intialization_params(monkeypatch): +def test_env_variables_override_initialization_params(monkeypatch): # When initialization params are specified in an environment variable, and we call get_model, we # should see that the model was initialized with those params monkeypatch.setattr(models, "models", {}) + fake_label_map = {"1": "label1", "2": "label2"} with mock.patch.dict( models.os.environ, {"UNSTRUCTURED_DEFAULT_MODEL_INITIALIZE_PARAMS_JSON_PATH": "fake_json.json"}, ), mock.patch.object(models, "DEFAULT_MODEL", "fake"), mock.patch.dict( - models.model_class_map, {"fake": mock.MagicMock()} + models.model_class_map, + {"fake": mock.MagicMock()}, ), mock.patch( - "builtins.open", mock.mock_open(read_data='{"date": "3/26/81"}') + "builtins.open", + mock.mock_open( + read_data='{"model_path": "fakepath", "label_map": ' + json.dumps(fake_label_map) + "}", + ), ): model = models.get_model() - model.initialize.assert_called_once_with(date="3/26/81") + model.initialize.assert_called_once_with( + model_path="fakepath", + label_map={1: "label1", 2: "label2"}, + ) From 85f68fbbb5fc7c6c3f216fad40292f29cd244e63 Mon Sep 17 00:00:00 2001 From: erjieyong Date: Sat, 17 Feb 2024 19:31:26 +0800 Subject: [PATCH 3/3] update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bf85833..95d264b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## 0.7.25-dev0 * chore: removes `install-detectron2` from the `Makefile` +* fix: convert label_map keys read from os.environment `UNSTRUCTURED_DEFAULT_MODEL_INITIALIZE_PARAMS_JSON_PATH` to int type ## 0.7.24