26 lines
921 B
Python
26 lines
921 B
Python
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"))
|