Address comments

Author: Lewis300

python/src/coreai_models/export/compression.py

@@ -166,13 +166,16 @@ def quantize_pytorch_model(
 
         reset_positions = set(state_indices)
         for pos in reset_positions:
-            assert pos >= 2, (
-                "States cannot occupy the first two input positions. "
-                "Those must be reserved for input_ids and position_ids"
-            )
-            assert pos < len(inputs), (
-                f"State index out of bounds, got {pos}, while the number of inputs is {len(inputs)}"
-            )
+            if pos < 2:
+                raise ValueError(
+                    "States cannot occupy the first two input positions. "
+                    "Those must be reserved for input_ids and position_ids"
+                )
+            if pos >= len(inputs):
+                raise IndexError(
+                    f"State index out of bounds, got {pos}, while the number of inputs is "
+                    f"{len(inputs)}"
+                )
 
         # Match the caller's declared bound: position_ids.shape[1] <= cache_seq_len - 1
         # (the `seq_pos` Dim in `BaseForCausalLM.build_dynamic_shapes`).

python/src/coreai_models/models/base.py

@@ -69,11 +69,12 @@ def __post_init__(self) -> None:
                 f"max_context_length={self.max_context_length} is too small to trace: "
                 f"it must be at least query_len + 2 = {self.query_len + 2}."
             )
-        assert self.cache_seq_len <= self.max_context_length, (
-            "cache_seq_len must not be greater than max_context_length. Received "
-            f"cache_seq_len = {self.cache_seq_len}, "
-            f"max_context_length = {self.max_context_length}"
-        )
+        if self.cache_seq_len > self.max_context_length:
+            raise ValueError(
+                "cache_seq_len must not be greater than max_context_length. Received "
+                f"cache_seq_len = {self.cache_seq_len}, "
+                f"max_context_length = {self.max_context_length}"
+            )
 
     @property
     def caches_are_static(self) -> bool: