File size: 1,942 Bytes
58904bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import os
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch
import numpy as np
from dropout_decay.data import load_cached_splits
from dropout_decay.experiments.device import assert_mps_only
class DataAndDeviceTests(unittest.TestCase):
def test_load_cached_splits_uses_expected_train_val_partition(self):
with tempfile.TemporaryDirectory() as tmp:
cache_dir = Path(tmp)
(cache_dir / "tokenizer-v16.json").write_text("{}", encoding="utf-8")
np.save(cache_dir / "tokens-v16-uint16.npy", np.arange(100, dtype=np.uint16))
tokenizer, splits = load_cached_splits(
cache_dir=cache_dir,
vocab_size=16,
max_required_train_tokens=80,
val_tokens=10,
allow_short_corpus=False,
)
self.assertEqual(tokenizer.vocab_size, 16)
self.assertEqual(len(splits.train), 90)
self.assertEqual(len(splits.val), 10)
self.assertEqual(splits.tokenizer_path, cache_dir / "tokenizer-v16.json")
def test_load_cached_splits_rejects_missing_cache(self):
with tempfile.TemporaryDirectory() as tmp:
with self.assertRaises(FileNotFoundError):
load_cached_splits(
cache_dir=Path(tmp),
vocab_size=16,
max_required_train_tokens=10,
val_tokens=10,
allow_short_corpus=False,
)
def test_mps_guard_rejects_fallback_before_backend_checks(self):
env = dict(os.environ)
env["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
with patch.dict(os.environ, env, clear=True):
with self.assertRaises(SystemExit) as caught:
assert_mps_only()
self.assertIn("MPS_FALLBACK", str(caught.exception))
if __name__ == "__main__":
unittest.main()
|