prain: expert id ranges

2026-10-27 13:52:05 +01:00
parent c8d5706956
commit db1736c42d
2 changed files with 84 additions and 0 deletions
@@ -0,0 +1,25 @@
+import os
+import unittest
+
+import prain.config as mc
+from tests.helpers import TmpCase
+
+
+class ParseIds(unittest.TestCase):
+    def test_ranges_lists_and_star(self):
+        self.assertEqual(mc.parse_ids("0-3,7", 8), (0, 1, 2, 3, 7))
+        self.assertEqual(mc.parse_ids([5, 1, 5], 8), (1, 5))
+        self.assertEqual(mc.parse_ids("*", 3), (0, 1, 2))
+
+    def test_rejects_bad_specs(self):
+        for spec in ("3-1", "a", "1,", "0-8", [1.5]):
+            with self.assertRaises(ValueError, msg=spec):
+                mc.parse_ids(spec, 8)
+
+    def test_format_is_the_inverse(self):
+        self.assertEqual(mc.format_ids([7, 0, 1, 2, 3, 9, 10]), "0-3,7,9-10")
+        self.assertEqual(mc.parse_ids(mc.format_ids([4, 2, 3]), 8), (2, 3, 4))
+
+    def test_layer_index(self):
+        self.assertEqual(mc.layer_index("language_model.model.layers.12.experts"), 12)
+        self.assertIsNone(mc.layer_index("model.experts"))