Spaces:
Sleeping
Sleeping
File size: 1,920 Bytes
e04eb0b | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 | #!/usr/bin/env python3
"""Offline unit tests for extract_limen_trajectory.py."""
import tempfile
import unittest
from pathlib import Path
import numpy as np
import torch
from extract_limen_trajectory import (
choose_device,
choose_dtype,
sha256_file,
validate_arrays,
)
class TestValidation(unittest.TestCase):
def test_valid_arrays(self):
hidden = np.zeros((4, 6, 8), dtype=np.float32)
logits = np.zeros((4, 32), dtype=np.float32)
validate_arrays(hidden, logits)
def test_hidden_rank_rejected(self):
with self.assertRaises(ValueError):
validate_arrays(np.zeros((4, 8)), None)
def test_nonfinite_hidden_rejected(self):
hidden = np.zeros((4, 6, 8))
hidden[0, 0, 0] = np.nan
with self.assertRaises(ValueError):
validate_arrays(hidden, None)
def test_token_misalignment_rejected(self):
hidden = np.zeros((4, 6, 8))
logits = np.zeros((3, 32))
with self.assertRaises(ValueError):
validate_arrays(hidden, logits)
class TestRuntimeChoices(unittest.TestCase):
def test_cpu_auto_dtype(self):
self.assertEqual(choose_dtype(torch.device("cpu"), "auto"), torch.float32)
def test_explicit_dtype(self):
self.assertEqual(
choose_dtype(torch.device("cpu"), "bfloat16"),
torch.bfloat16,
)
def test_explicit_cpu(self):
self.assertEqual(choose_device("cpu").type, "cpu")
class TestHashing(unittest.TestCase):
def test_sha256_file(self):
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "value.bin"
path.write_bytes(b"limen")
self.assertEqual(
sha256_file(path),
"96ad72da603bfdffc6d04b3b6a22f9f90b546d1ced158ec3671d46e672457255",
)
if __name__ == "__main__":
unittest.main(verbosity=2)
|