Coverage for maze_dataset/tokenization/modular/hashing.py: 45%

1"""legacy system for checking a `ModularMazeTokenizer` is valid -- compare its hash to a table of known hashes 

2 

3this has been superseded by the fst system 

4 

5""" 

6 

7import hashlib 

8from pathlib import Path 

9 

10import numpy as np 

11from jaxtyping import UInt32 

12 

13# NOTE: these all need to match! 

14 

15AllTokenizersHashBitLength = 32 

16"bit length of the hashes of all tokenizers, must match `AllTokenizersHashDtype` and `AllTokenizersHashesArray`" 

17 

18AllTokenizersHashDtype = np.uint32 

19"numpy data type of the hashes of all tokenizers, must match `AllTokenizersHashBitLength` and `AllTokenizersHashesArray`" 

20 

21AllTokenizersHashesArray = UInt32[np.ndarray, " n_tokens"] 

22"jaxtyping type of the hashes of all tokenizers, must match `AllTokenizersHashBitLength` and `AllTokenizersHashDtype`" 

23 

24 

25def _hash_tokenizer_name(s: str) -> int: 

26 h64: int = int.from_bytes( 

27 hashlib.shake_256(s.encode("utf-8")).digest(64), 

28 byteorder="big", 

29 ) 

30 return (h64 >> 32) ^ (h64 & 0xFFFFFFFF) 

31 

32 

33_ALL_TOKENIZER_HASHES: AllTokenizersHashesArray 

34"private array of all tokenizer hashes" 

35_TOKENIZER_HASHES_PATH: Path = Path(__file__).parent / "MazeTokenizerModular_hashes.npz" 

36"path to where we expect the hashes file -- in the same dir as this file, by default. change with `set_tokenizer_hashes_path`" 

37 

38 

39def set_tokenizer_hashes_path(path: Path) -> None: 

40 """set path to tokenizer hashes, and reload the hashes if needed 

41 

42 the hashes are expected to be stored in and read from `_TOKENIZER_HASHES_PATH`, 

43 which by default is `Path(__file__).parent / "MazeTokenizerModular_hashes.npz"` or in this file's directory. 

44 

45 However, this might not always work, so we provide a way to change this. 

46 """ 

47 global _TOKENIZER_HASHES_PATH, _ALL_TOKENIZER_HASHES # noqa: PLW0603 

48 

49 path = Path(path) 

50 if path.is_dir(): 

51 path = path / "MazeTokenizerModular_hashes.npz" 

52 

53 if not path.is_file(): 

54 err_msg: str = f"could not find maze tokenizer hashes file at: {path}" 

55 raise FileNotFoundError(err_msg) 

56 

57 if _TOKENIZER_HASHES_PATH.absolute() != path.absolute(): 

58 # reload if they aren't equal 

59 _TOKENIZER_HASHES_PATH = path 

60 _ALL_TOKENIZER_HASHES = _load_tokenizer_hashes() 

61 else: 

62 # always set to new path 

63 _TOKENIZER_HASHES_PATH = path 

64 

65 

66def _load_tokenizer_hashes() -> AllTokenizersHashesArray: 

67 """Loads the sorted list of `all_tokenizers.get_all_tokenizers()` hashes from disk.""" 

68 global _TOKENIZER_HASHES_PATH # noqa: PLW0602 

69 try: 

70 path: Path = _TOKENIZER_HASHES_PATH 

71 return np.load(path)["hashes"] 

72 except FileNotFoundError as e: 

73 err_msg: str = ( 

74 "Tokenizers hashes cannot be loaded. To fix this, run" 

75 "\n`python -m maze-dataset.tokenization.save_hashes` which will save the hashes to" 

76 "\n`data/MazeTokenizerModular_hashes.npz`" 

77 "\nrelative to the current working directory -- this is where the code looks for them." 

78 ) 

79 raise FileNotFoundError(err_msg) from e 

80 

81 

82def get_all_tokenizer_hashes() -> AllTokenizersHashesArray: 

83 """returns all the tokenizer hashes in an `AllTokenizersHashesDtype` array, setting global variable if needed""" 

84 # naughty use of globals 

85 global _ALL_TOKENIZER_HASHES # noqa: PLW0603 

86 try: 

87 got_tokenizers: bool = len(_ALL_TOKENIZER_HASHES) > 0 

88 if got_tokenizers: 

89 return _ALL_TOKENIZER_HASHES 

90 else: 

91 _ALL_TOKENIZER_HASHES = _load_tokenizer_hashes() 

92 except NameError: 

93 _ALL_TOKENIZER_HASHES = _load_tokenizer_hashes() 

94 

95 return _ALL_TOKENIZER_HASHES