Merge pull request #59 from foundation-model-stack/skip_assertions

added an option to skip assertions in validation

Author: JRosenkranz

tests/models/test_decoders.py

@@ -48,6 +48,7 @@
 USE_MICRO_MODELS = os.environ.get("FMS_TEST_SHAPES_USE_MICRO_MODELS", "1") == "1"
 USE_DISTRIBUTED = os.environ.get("FMS_TEST_SHAPES_DISTRIBUTED", "0") == "1"
 FORCE_VALIDATION_LEVEL_1 = os.environ.get("FMS_TEST_SHAPES_FORCE_VALIDATION_LEVEL_1", "0") == "1"
+skip_assertions = os.environ.get("FMS_TEST_SHAPES_SKIP_ASSERTIONS", {})
 validation_info_dir = os.environ.get(
     "FMS_TEST_SHAPES_VALIDATION_INFO_DIR", "/tmp/models/validation_info"
 )
@@ -99,6 +100,16 @@
 if isinstance(common_max_new_tokens, str):
     common_max_new_tokens = [int(mnt) for mnt in common_max_new_tokens.split(",")]
 
+# pass metrics to skip as a comma separated list (ce,mean_diff)
+if isinstance(skip_assertions, str):
+    _skip_assertions = []
+    for metric in skip_assertions.split(","):
+        metric = metric.lower()
+        if metric not in {"ce", "mean_diff"}:
+            pytest.fail("FMS_TEST_SHAPES_SKIP_ASSERTIONS can only accept metrics ce and mean_diff")
+        _skip_assertions.append(metric)
+    skip_assertions = set(_skip_assertions)
+
 common_shapes = list(
     itertools.product(
         common_model_paths,
@@ -521,12 +532,14 @@ def _metric_calculator(r: torch.Tensor, t: torch.Tensor):
         ce_failure_rate = len(ce_fail_responses_list) / total_tokens
         dprint(f"mean diff failure rate: {diff_failure_rate}")
         dprint(f"cross entropy loss failure rate: {ce_failure_rate}")
-        assert diff_failure_rate < failure_rate_threshold, (
-            f"failure rate for mean diff was too high: {diff_failure_rate}"
-        )
-        assert ce_failure_rate < failure_rate_threshold, (
-            f"failure rate for cross entropy loss was too high: {ce_failure_rate}"
-        )
+        if "mean_diff" not in skip_assertions:
+            assert diff_failure_rate < failure_rate_threshold, (
+                f"failure rate for mean diff was too high: {diff_failure_rate}"
+            )
+        if "ce" not in skip_assertions:
+            assert ce_failure_rate < failure_rate_threshold, (
+                f"failure rate for cross entropy loss was too high: {ce_failure_rate}"
+            )
 
         print("passed validation level 1")
     else: