Coverage for maze_dataset/tokenization/modular/all_tokenizers.py: 57%

1"""Contains `get_all_tokenizers()` and supporting limited-use functions. 

2 

3# `get_all_tokenizers()` 

4returns a comprehensive collection of all valid `MazeTokenizerModular` objects. 

5This is an overwhelming majority subset of the set of all possible `MazeTokenizerModular` objects. 

6Other tokenizers not contained in `get_all_tokenizers()` may be possible to construct, but they are untested and not guaranteed to work. 

7This collection is in a separate module since it is expensive to compute and will grow more expensive as features are added to `MazeTokenizerModular`. 

8 

9## Use Cases 

10In general, uses for this module are limited to development of the library and specific research studying many tokenization behaviors. 

11- Unit testing: 

12 - Tokenizers to use in unit tests are sampled from `get_all_tokenizers()` 

13- Large-scale tokenizer research: 

14 - Specific research training models on many tokenization behaviors can use `get_all_tokenizers()` as the maximally inclusive collection 

15 - `get_all_tokenizers()` may be subsequently filtered using `MazeTokenizerModular.has_element` 

16For other uses, it's likely that the computational expense can be avoided by using 

17- `maze_tokenizer.get_all_tokenizer_hashes()` for membership checks 

18- `utils.all_instances` for generating smaller subsets of `MazeTokenizerModular` or `_TokenizerElement` objects 

19 

20# `EVERY_TEST_TOKENIZERS` 

21A collection of the tokenizers which should always be included in unit tests when test fuzzing is used. 

22This collection should be expanded as specific tokenizers become canonical or popular. 

23""" 

24 

25import functools 

26import multiprocessing 

27import random 

28from functools import cache 

29from pathlib import Path 

30from typing import Callable 

31 

32import frozendict 

33import numpy as np 

34from muutils.spinner import NoOpContextManager, SpinnerContext 

35from tqdm import tqdm 

36 

37from maze_dataset.tokenization import ( 

38 CoordTokenizers, 

39 MazeTokenizerModular, 

40 PromptSequencers, 

41 StepTokenizers, 

42 _TokenizerElement, 

43) 

44from maze_dataset.tokenization.modular.all_instances import FiniteValued, all_instances 

45from maze_dataset.tokenization.modular.hashing import ( 

46 AllTokenizersHashBitLength, 

47 AllTokenizersHashDtype, 

48 AllTokenizersHashesArray, 

49) 

50 

51# Always include this as the first item in the dict `validation_funcs` whenever using `all_instances` with `MazeTokenizerModular` 

52# TYPING: error: Type variable "maze_dataset.utils.FiniteValued" is unbound [valid-type] 

53# note: (Hint: Use "Generic[FiniteValued]" or "Protocol[FiniteValued]" base class to bind "FiniteValued" inside a class) 

54# note: (Hint: Use "FiniteValued" in function signature to bind "FiniteValued" inside a function) 

55MAZE_TOKENIZER_MODULAR_DEFAULT_VALIDATION_FUNCS: frozendict.frozendict[ 

56 type[FiniteValued], 

57 Callable[[FiniteValued], bool], 

58] = frozendict.frozendict( 

59 { 

60 # TYPING: Item "bool" of the upper bound "bool | IsDataclass | Enum" of type variable "FiniteValued" has no attribute "is_valid" [union-attr] 

61 _TokenizerElement: lambda x: x.is_valid(), 

62 # Currently no need for `MazeTokenizerModular.is_valid` since that method contains no special cases not already covered by `_TokenizerElement.is_valid` 

63 # MazeTokenizerModular: lambda x: x.is_valid(), 

64 # TYPING: error: No overload variant of "set" matches argument type "FiniteValued" [call-overload] 

65 # note: Possible overload variants: 

66 # note: def [_T] set(self) -> set[_T] 

67 # note: def [_T] set(self, Iterable[_T], /) -> set[_T] 

68 # TYPING: error: Argument 1 to "len" has incompatible type "FiniteValued"; expected "Sized" [arg-type] 

69 StepTokenizers.StepTokenizerPermutation: lambda x: len(set(x)) == len(x) 

70 and x != (StepTokenizers.Distance(),), 

71 }, 

72) 

73 

74DOWNLOAD_URL: str = "https://raw.githubusercontent.com/understanding-search/maze-dataset/main/maze_dataset/tokenization/MazeTokenizerModular_hashes.npz" 

75 

76 

77@cache 

78def get_all_tokenizers() -> list[MazeTokenizerModular]: 

79 """Computes a complete list of all valid tokenizers. 

80 

81 Warning: This is an expensive function. 

82 """ 

83 return list( 

84 all_instances( 

85 MazeTokenizerModular, 

86 validation_funcs=MAZE_TOKENIZER_MODULAR_DEFAULT_VALIDATION_FUNCS, 

87 ), 

88 ) 

89 

90 

91@cache 

92def get_all_tokenizers_names() -> list[str]: 

93 """computes the sorted list of names of all tokenizers""" 

94 return sorted([tokenizer.name for tokenizer in get_all_tokenizers()]) 

95 

96 

97EVERY_TEST_TOKENIZERS: list[MazeTokenizerModular] = [ 

98 MazeTokenizerModular(), 

99 MazeTokenizerModular( 

100 prompt_sequencer=PromptSequencers.AOTP(coord_tokenizer=CoordTokenizers.CTT()), 

101 ), 

102 # TODO: add more here as specific tokenizers become canonical and frequently used 

103] 

104 

105 

106@cache 

107def all_tokenizers_set() -> set[MazeTokenizerModular]: 

108 """Casts `get_all_tokenizers()` to a set.""" 

109 return set(get_all_tokenizers()) 

110 

111 

112@cache 

113def _all_tokenizers_except_every_test_tokenizers() -> list[MazeTokenizerModular]: 

114 """Returns""" 

115 return list(all_tokenizers_set().difference(EVERY_TEST_TOKENIZERS)) 

116 

117 

118def sample_all_tokenizers(n: int) -> list[MazeTokenizerModular]: 

119 """Samples `n` tokenizers from `get_all_tokenizers()`.""" 

120 return random.sample(get_all_tokenizers(), n) 

121 

122 

123def sample_tokenizers_for_test(n: int | None) -> list[MazeTokenizerModular]: 

124 """Returns a sample of size `n` of unique elements from `get_all_tokenizers()`, 

125 

126 always including every element in `EVERY_TEST_TOKENIZERS`. 

127 """ 

128 if n is None: 

129 return get_all_tokenizers() 

130 

131 if n < len(EVERY_TEST_TOKENIZERS): 

132 err_msg: str = f"`n` must be at least {len(EVERY_TEST_TOKENIZERS) = } such that the sample can contain `EVERY_TEST_TOKENIZERS`." 

133 raise ValueError( 

134 err_msg, 

135 ) 

136 sample: list[MazeTokenizerModular] = random.sample( 

137 _all_tokenizers_except_every_test_tokenizers(), 

138 n - len(EVERY_TEST_TOKENIZERS), 

139 ) 

140 sample.extend(EVERY_TEST_TOKENIZERS) 

141 return sample 

142 

143 

144def save_hashes( 

145 path: Path | None = None, 

146 verbose: bool = False, 

147 parallelize: bool | int = False, 

148) -> AllTokenizersHashesArray: 

149 """Computes, sorts, and saves the hashes of every member of `get_all_tokenizers()`.""" 

150 spinner = ( 

151 functools.partial(SpinnerContext, spinner_chars="square_dot") 

152 if verbose 

153 else NoOpContextManager 

154 ) 

155 

156 # get all tokenizers 

157 with spinner(initial_value="getting all tokenizers...", update_interval=2.0): 

158 all_tokenizers = get_all_tokenizers() 

159 

160 # compute hashes 

161 hashes_array_np64: AllTokenizersHashesArray 

162 if parallelize: 

163 n_cpus: int = ( 

164 parallelize if int(parallelize) > 1 else multiprocessing.cpu_count() 

165 ) 

166 with spinner( # noqa: SIM117 

167 initial_value=f"using {n_cpus} processes to compute {len(all_tokenizers)} tokenizer hashes...", 

168 update_interval=2.0, 

169 ): 

170 with multiprocessing.Pool(processes=n_cpus) as pool: 

171 hashes_list: list[int] = list(pool.map(hash, all_tokenizers)) 

172 

173 with spinner(initial_value="converting hashes to numpy array..."): 

174 hashes_array_np64 = np.array(hashes_list, dtype=np.int64) 

175 else: 

176 with spinner( 

177 initial_value=f"computing {len(all_tokenizers)} tokenizer hashes...", 

178 ): 

179 hashes_array_np64 = np.array( 

180 [ 

181 hash(obj) # uses stable hash 

182 for obj in tqdm(all_tokenizers, disable=not verbose) 

183 ], 

184 dtype=np.int64, 

185 ) 

186 

187 # convert to correct dtype 

188 hashes_array: AllTokenizersHashesArray = ( 

189 hashes_array_np64 % (1 << AllTokenizersHashBitLength) 

190 if AllTokenizersHashBitLength < 64 # noqa: PLR2004 

191 else hashes_array_np64 

192 ).astype(AllTokenizersHashDtype) 

193 

194 # make sure there are no dupes 

195 with spinner(initial_value="sorting and checking for hash collisions..."): 

196 sorted_hashes, counts = np.unique(hashes_array, return_counts=True) 

197 if sorted_hashes.shape[0] != hashes_array.shape[0]: 

198 collisions: np.array = sorted_hashes[counts > 1] 

199 n_collisions: int = hashes_array.shape[0] - sorted_hashes.shape[0] 

200 err_msg: str = ( 

201 f"{n_collisions} tokenizer hash collisions: {collisions}\n" 

202 "Report error to the developer to increase the hash size or otherwise update the tokenizer hashing size:\n" 

203 f"https://github.com/understanding-search/maze-dataset/issues/new?labels=bug,tokenization&title=Tokenizer+hash+collision+error&body={n_collisions}+collisions+out+of+{hashes_array.shape[0]}+total+hashes", 

204 ) 

205 

206 raise ValueError( 

207 err_msg, 

208 ) 

209 

210 # save and return 

211 with spinner(initial_value="saving hashes...", update_interval=0.5): 

212 if path is None: 

213 path = Path(__file__).parent / "MazeTokenizerModular_hashes.npz" 

214 np.savez_compressed( 

215 path, 

216 hashes=sorted_hashes, 

217 ) 

218 

219 return sorted_hashes