Merge nucleic/sleek-ember-seal-uady into dev

tests/test_convert_coreml.py

@@ -0,0 +1,84 @@
+import json
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+from transformers import BertConfig, BertForSequenceClassification
+
+
+MODULE_DIR = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(MODULE_DIR))
+
+import convert_coreml
+from purpose_data import LABELS, DataError
+
+
+class FixedShapeBertForCoreMLTests(unittest.TestCase):
+    def test_conversion_forward_matches_transformers(self):
+        torch.manual_seed(7)
+        config = BertConfig(
+            vocab_size=64,
+            hidden_size=16,
+            num_hidden_layers=1,
+            num_attention_heads=4,
+            intermediate_size=32,
+            max_position_embeddings=128,
+            type_vocab_size=2,
+            hidden_dropout_prob=0.0,
+            attention_probs_dropout_prob=0.0,
+            num_labels=len(LABELS),
+        )
+        model = BertForSequenceClassification(config).eval()
+        wrapper = convert_coreml.FixedShapeBertForCoreML(model).eval()
+        input_ids = torch.randint(0, config.vocab_size, (1, 128), dtype=torch.int32)
+        attention_mask = torch.zeros((1, 128), dtype=torch.int32)
+        attention_mask[:, :83] = 1
+        token_type_ids = torch.zeros((1, 128), dtype=torch.int32)
+        token_type_ids[:, 43:83] = 1
+        with torch.inference_mode():
+            reference = model(
+                input_ids=input_ids.long(),
+                attention_mask=attention_mask.long(),
+                token_type_ids=token_type_ids.long(),
+            ).logits
+            candidate = wrapper(input_ids, attention_mask, token_type_ids)
+            torch.testing.assert_close(candidate, reference, rtol=1e-5, atol=2e-5)
+
+            traced = torch.jit.trace(
+                wrapper,
+                (input_ids, attention_mask, token_type_ids),
+                strict=True,
+            )
+            torch.testing.assert_close(
+                traced(input_ids, attention_mask, token_type_ids),
+                reference,
+                rtol=1e-5,
+                atol=2e-5,
+            )
+
+
+class CheckpointConfigTests(unittest.TestCase):
+    def test_rejects_changed_label_order(self):
+        config = {
+            "model_type": "bert",
+            "hidden_size": 384,
+            "num_hidden_layers": 6,
+            "id2label": {
+                str(index): label
+                for index, label in enumerate(reversed(LABELS))
+            },
+        }
+        with tempfile.TemporaryDirectory() as temp:
+            model_dir = Path(temp)
+            (model_dir / "config.json").write_text(
+                json.dumps(config),
+                encoding="utf-8",
+            )
+            with self.assertRaisesRegex(DataError, "label order"):
+                convert_coreml._checkpoint_config(model_dir)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_eval.py

@@ -9,6 +9,29 @@ sys.path.insert(0, str(MODULE_DIR))
 import eval as purpose_eval
 
 
+class CoreMLComputeUnitTests(unittest.TestCase):
+    class CoreMLTools:
+        class ComputeUnit:
+            ALL = "all-value"
+            CPU_ONLY = "cpu-value"
+            CPU_AND_GPU = "gpu-value"
+            CPU_AND_NE = "ne-value"
+
+    def test_maps_cli_compute_policies(self):
+        expected = {
+            "all": "all-value",
+            "cpu-only": "cpu-value",
+            "cpu-and-gpu": "gpu-value",
+            "cpu-and-ne": "ne-value",
+        }
+        for requested, value in expected.items():
+            with self.subTest(requested=requested):
+                self.assertEqual(
+                    value,
+                    purpose_eval._coreml_compute_unit(self.CoreMLTools, requested),
+                )
+
+
 class TierDriftTests(unittest.TestCase):
     def test_current_routing_matrix_bounds_every_label_pair(self):
         records = [{"prompt": f"prompt {index}"} for index in range(8 * 8)]

tests/test_inspect_coreml.py

@@ -0,0 +1,26 @@
+import sys
+import unittest
+from pathlib import Path
+
+
+MODULE_DIR = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(MODULE_DIR))
+
+import inspect_coreml
+
+
+class CoreMLDeviceCategoryTests(unittest.TestCase):
+    def test_classifies_compute_device_types(self):
+        NeuralEngineDevice = type("MLNeuralEngineComputeDevice", (), {})
+        GPUDevice = type("MLGPUComputeDevice", (), {})
+        CPUDevice = type("MLCPUComputeDevice", (), {})
+        self.assertEqual(
+            "neuralEngine",
+            inspect_coreml._device_category(NeuralEngineDevice()),
+        )
+        self.assertEqual("gpu", inspect_coreml._device_category(GPUDevice()))
+        self.assertEqual("cpu", inspect_coreml._device_category(CPUDevice()))
+
+
+if __name__ == "__main__":
+    unittest.main()