1"""a whole bunch of utilities for tokenization"""
  2
  3import re
  4import typing
  5import warnings
  6from collections import Counter
  7from typing import Callable, Literal, overload
  8
  9import numpy as np
 10from jaxtyping import Bool, Float, Int, Int8
 11from muutils.errormode import ErrorMode
 12from muutils.misc import list_join
 13from muutils.misc.sequence import WhenMissing
 14
 15from maze_dataset.constants import (
 16	CARDINAL_MAP,
 17	SPECIAL_TOKENS,
 18	VOCAB,
 19	ConnectionArray,
 20	ConnectionList,
 21	CoordTup,
 22)
 23
 24# filtering things from a prompt or generated text
 25# ==================================================
 26
 27
 28def remove_padding_from_token_str(token_str: str) -> str:
 29	"""remove padding tokens from a joined token string"""
 30	token_str = token_str.replace(f"{SPECIAL_TOKENS.PADDING} ", "")
 31	token_str = token_str.replace(f"{SPECIAL_TOKENS.PADDING}", "")
 32	return token_str  # noqa: RET504
 33
 34
 35def tokens_between(
 36	tokens: list[str],
 37	start_value: str,
 38	end_value: str,
 39	include_start: bool = False,
 40	include_end: bool = False,
 41	except_when_tokens_not_unique: bool = False,
 42) -> list[str]:
 43	"""given a list `tokens`, get the tokens between `start_value` and `end_value`
 44
 45	_extended_summary_
 46
 47	# Parameters:
 48	- `tokens : list[str]`
 49	- `start_value : str`
 50	- `end_value : str`
 51	- `include_start : bool`
 52		(defaults to `False`)
 53	- `include_end : bool`
 54		(defaults to `False`)
 55	- `except_when_tokens_not_unique : bool`
 56		when `True`, raise an error if `start_value` or `end_value` are not unique in the input tokens
 57		(defaults to `False`)
 58
 59	# Returns:
 60	- `list[str]`
 61
 62	# Raises:
 63	- `ValueError` : if `start_value` and `end_value` are the same
 64	- `ValueError` : if `except_when_tokens_not_unique` is `True` and `start_value` or `end_value` are not unique in the input tokens
 65	- `ValueError` : if `start_value` or `end_value` are not present in the input tokens
 66	"""
 67	if start_value == end_value:
 68		err_msg: str = f"start_value and end_value cannot be the same: {start_value = } {end_value = }"
 69		raise ValueError(
 70			err_msg,
 71		)
 72	if except_when_tokens_not_unique:
 73		if (tokens.count(start_value) != 1) or (tokens.count(end_value) != 1):
 74			err_msg: str = (
 75				"start_value or end_value is not unique in the input tokens:"
 76				f"\n{tokens.count(start_value) = } {tokens.count(end_value) = }"
 77				f"\n{start_value = } {end_value = }"
 78				f"\n{tokens = }"
 79			)
 80			raise ValueError(err_msg)
 81	else:
 82		if (tokens.count(start_value) < 1) or (tokens.count(end_value) < 1):
 83			err_msg: str = (
 84				"start_value or end_value is not present in the input tokens:"
 85				f"\n{tokens.count(start_value) = } {tokens.count(end_value) = }"
 86				f"\n{start_value = } {end_value = }"
 87				f"\n{tokens = }"
 88			)
 89			raise ValueError(err_msg)
 90
 91	start_idx: int = tokens.index(start_value) + int(not include_start)
 92	end_idx: int = tokens.index(end_value) + int(include_end)
 93
 94	assert start_idx < end_idx, "Start must come before end"
 95
 96	return tokens[start_idx:end_idx]
 97
 98
 99def get_adj_list_tokens(tokens: list[str]) -> list[str]:
100	"get tokens between ADJLIST_START and ADJLIST_END, without the special tokens themselves"
101	return tokens_between(
102		tokens,
103		SPECIAL_TOKENS.ADJLIST_START,
104		SPECIAL_TOKENS.ADJLIST_END,
105	)
106
107
108def get_path_tokens(tokens: list[str], trim_end: bool = False) -> list[str]:
109	"""The path is considered everything from the first path coord to the path_end token, if it exists."""
110	if SPECIAL_TOKENS.PATH_START not in tokens:
111		err_msg: str = f"Path start token {SPECIAL_TOKENS.PATH_START} not found in tokens:\n{tokens}"
112		raise ValueError(
113			err_msg,
114		)
115	start_idx: int = tokens.index(SPECIAL_TOKENS.PATH_START) + int(trim_end)
116	end_idx: int | None = None
117	if trim_end and (SPECIAL_TOKENS.PATH_END in tokens):
118		end_idx = tokens.index(SPECIAL_TOKENS.PATH_END)
119	return tokens[start_idx:end_idx]
120
121
122def get_context_tokens(tokens: list[str]) -> list[str]:
123	"get tokens between ADJLIST_START and PATH_START"
124	return tokens_between(
125		tokens,
126		SPECIAL_TOKENS.ADJLIST_START,
127		SPECIAL_TOKENS.PATH_START,
128		include_start=True,
129		include_end=True,
130	)
131
132
133def get_origin_tokens(tokens: list[str]) -> list[str]:
134	"get tokens_between ORIGIN_START and ORIGIN_END"
135	return tokens_between(
136		tokens,
137		SPECIAL_TOKENS.ORIGIN_START,
138		SPECIAL_TOKENS.ORIGIN_END,
139		include_start=False,
140		include_end=False,
141	)
142
143
144def get_target_tokens(tokens: list[str]) -> list[str]:
145	"get tokens_between TARGET_START and TARGET_END"
146	return tokens_between(
147		tokens,
148		SPECIAL_TOKENS.TARGET_START,
149		SPECIAL_TOKENS.TARGET_END,
150		include_start=False,
151		include_end=False,
152	)
153
154
155def get_cardinal_direction(coords: Int[np.ndarray, "start_end=2 row_col=2"]) -> str:
156	"""Returns the cardinal direction token corresponding to traveling from `coords[0]` to `coords[1]`."""
157	return CARDINAL_MAP[tuple(coords[1] - coords[0])]
158
159
160def get_relative_direction(coords: Int[np.ndarray, "prev_cur_next=3 row_col=2"]) -> str:
161	"""Returns the relative first-person direction token corresponding to traveling from `coords[1]` to `coords[2]`.
162
163	# Parameters
164	- `coords`: Contains 3 Coords, each of which must neighbor the previous Coord.
165		- `coords[0]`: The previous location, used to determine the current absolute direction that the "agent" is facing.
166		- `coords[1]`: The current location
167		- `coords[2]`: The next location. May be equal to the current location.
168	"""
169	if coords.shape != (3, 2):
170		err_msg: str = f"`coords` must have shape (3,2). Got {coords.shape} instead."
171		raise ValueError(err_msg)
172	directions = coords[1:] - coords[:-1]
173	if not np.all(np.linalg.norm(directions, axis=1) <= np.array([1.1, 1.1])):
174		# Use floats as constant since `np.linalg.norm` returns float array
175		err_msg: str = f"Adjacent `coords` must be neighboring or equivalent. Got {coords} instead."
176		raise ValueError(
177			err_msg,
178		)
179	if np.array_equal(coords[1], coords[2]):
180		return VOCAB.PATH_STAY
181	if np.array_equal(coords[0], coords[2]):
182		return VOCAB.PATH_BACKWARD
183	if np.array_equal(coords[0], coords[1]):
184		err_msg: str = f"Previous first-person direction indeterminate from {coords=}."
185		raise ValueError(
186			err_msg,
187		)
188	if np.array_equal(directions[0], directions[1]):
189		return VOCAB.PATH_FORWARD
190	directions = np.append(
191		directions,
192		[[0], [0]],
193		axis=1,
194	)  # Augment to represent unit basis vectors in 3D
195	match np.cross(directions[0], directions[1])[-1]:
196		case 1:
197			return VOCAB.PATH_LEFT
198		case -1:
199			return VOCAB.PATH_RIGHT
200
201
202class TokenizerPendingDeprecationWarning(PendingDeprecationWarning):
203	"""Pending deprecation warnings related to the `MazeTokenizerModular` upgrade."""
204
205	pass
206
207
208def str_is_coord(coord_str: str, allow_whitespace: bool = True) -> bool:
209	"""return True if the string represents a coordinate, False otherwise"""
210	warnings.warn(
211		"`util.str_is_coord` only supports legacy UT strings. Function will be replaced with a generalized version in a future release.",
212		TokenizerPendingDeprecationWarning,
213	)
214	strip_func: Callable[[str], str] = lambda x: x.strip() if allow_whitespace else x  # noqa: E731
215
216	coord_str = strip_func(coord_str)
217
218	return all(
219		[
220			coord_str.startswith("("),
221			coord_str.endswith(")"),
222			"," in coord_str,
223			all(
224				strip_func(x).isdigit()
225				for x in strip_func(coord_str.lstrip("(").rstrip(")")).split(",")
226			),
227		],
228	)
229
230
231class TokenizerDeprecationWarning(DeprecationWarning):
232	"""Deprecation warnings related to the `MazeTokenizerModular` upgrade."""
233
234	pass
235
236
237# coordinate to strings
238# ==================================================
239
240
241def _coord_to_strings_UT(coord: typing.Sequence[int]) -> list[str]:
242	"""convert a coordinate to a string: `(i,j)`->"(i,j)"
243
244	always returns a list of length 1
245	"""
246	return [f"({','.join(str(c) for c in coord)})"]
247
248
249def _coord_to_strings_indexed(coord: typing.Sequence[int]) -> list[str]:
250	"""convert a coordinate to a list of indexed strings: `(i,j)`->"(", "i", ",", "j", ")"
251
252	always returns a list of length 5
253	"""
254	return [
255		"(",
256		*list_join([str(c) for c in coord], lambda: ","),
257		")",
258	]
259
260
261def coord_str_to_tuple(
262	coord_str: str,
263	allow_whitespace: bool = True,
264) -> tuple[int, ...]:
265	"""convert a coordinate string to a tuple"""
266	strip_func: Callable[[str], str] = lambda x: x.strip() if allow_whitespace else x  # noqa: E731
267	coord_str = strip_func(coord_str)
268	stripped: str = strip_func(coord_str.lstrip("(").rstrip(")"))
269	return tuple(int(strip_func(x)) for x in stripped.split(","))
270
271
272def coord_str_to_coord_np(coord_str: str, allow_whitespace: bool = True) -> np.ndarray:
273	"""convert a coordinate string to a numpy array"""
274	return np.array(coord_str_to_tuple(coord_str, allow_whitespace=allow_whitespace))
275
276
277def coord_str_to_tuple_noneable(coord_str: str) -> CoordTup | None:
278	"""convert a coordinate string to a tuple, or None if the string is not a coordinate string"""
279	if not str_is_coord(coord_str):
280		return None
281	return coord_str_to_tuple(coord_str)
282
283
284def coords_string_split_UT(coords: str) -> list[str]:
285	"""Splits a string of tokens into a list containing the UT tokens for each coordinate.
286
287	Not capable of producing indexed tokens ("(", "1", ",", "2", ")"), only unique tokens ("(1,2)").
288	Non-whitespace portions of the input string not matched are preserved in the same list:
289	"(1,2) <SPECIAL_TOKEN> (5,6)" -> ["(1,2)", "<SPECIAL_TOKEN>", "(5,6)"]
290	"""
291	# ty gpt4
292	return re.findall(r"\([^)]*\)|\S+", coords)
293
294
295# back and forth in wrapped form
296# ==================================================
297@overload
298def strings_to_coords(
299	text: str | list[str],
300	when_noncoord: Literal["skip"] = "skip",
301) -> list[CoordTup]: ...
302@overload
303def strings_to_coords(
304	text: str | list[str],
305	when_noncoord: Literal["error"] = "error",
306) -> list[CoordTup]: ...
307@overload
308def strings_to_coords(
309	text: str | list[str],
310	when_noncoord: Literal["include"] = "include",
311) -> list[str | CoordTup]: ...
312def strings_to_coords(
313	text: str | list[str],
314	when_noncoord: WhenMissing = "skip",
315) -> list[str | CoordTup]:
316	"""converts a list of tokens to a list of coordinates
317
318	returns list[CoordTup] if `when_noncoord` is "skip" or "error"
319	returns list[str | CoordTup] if `when_noncoord` is "include"
320	"""
321	warnings.warn(
322		"`util.strings_to_coords` only supports legacy UT strings. Function will be replaced with a generalized version in a future release.",
323		TokenizerPendingDeprecationWarning,
324	)
325	tokens_joined: str = text if isinstance(text, str) else " ".join(text)
326	tokens_processed: list[str] = coords_string_split_UT(tokens_joined)
327	result: list[str] = list()
328	for token in tokens_processed:
329		coord: CoordTup | None = coord_str_to_tuple_noneable(token)
330		if coord is None:
331			if when_noncoord == "skip":
332				continue
333			if when_noncoord == "error":
334				err_msg: str = (
335					f"Invalid non-coordinate token '{token}' in text: '{text}'"
336				)
337				raise ValueError(
338					err_msg,
339				)
340			if when_noncoord == "include":
341				result.append(token)
342			else:
343				err_msg: str = f"Invalid when_noncoord value '{when_noncoord}'"
344				raise ValueError(err_msg)
345		else:
346			result.append(coord)
347	return result
348
349
350@overload
351def coords_to_strings(
352	coords: list[str | CoordTup],
353	coord_to_strings_func: Callable[[CoordTup], list[str]],
354	when_noncoord: Literal["include", "skip"] = "skip",
355) -> list[str]: ...
356@overload
357def coords_to_strings(
358	coords: list[CoordTup],
359	coord_to_strings_func: Callable[[CoordTup], list[str]],
360	when_noncoord: Literal["error"] = "error",
361) -> list[str]: ...
362def coords_to_strings(
363	coords: list[str | CoordTup],
364	coord_to_strings_func: Callable[[CoordTup], list[str]],
365	when_noncoord: WhenMissing = "skip",
366) -> list[str]:
367	"""converts a list of coordinates to a list of strings (tokens)
368
369	expects list[CoordTup] if `when_noncoord` is "error"
370	expects list[str | CoordTup] if `when_noncoord` is "include" or "skip"
371	"""
372	result: list[str] = list()
373	for coord in coords:
374		if isinstance(coord, str):
375			if when_noncoord == "skip":
376				continue
377			if when_noncoord == "error":
378				err_msg: str = (
379					f"Invalid non-coordinate '{coord}' in list of coords: '{coords}'"
380				)
381				raise ValueError(
382					err_msg,
383				)
384			if when_noncoord == "include":
385				result.append(coord)
386			else:
387				err_msg: str = f"Invalid when_noncoord value '{when_noncoord}'"
388				raise ValueError(err_msg)
389		else:
390			result.extend(coord_to_strings_func(coord))
391	return result
392
393
394def get_token_regions(toks: list[str]) -> tuple[list[str], list[str]]:
395	"""Splits a list of tokens into adjacency list tokens and non-adjacency list tokens."""
396	adj_list_start, adj_list_end = (
397		toks.index("<ADJLIST_START>") + 1,
398		toks.index("<ADJLIST_END>"),
399	)
400	adj_list = toks[adj_list_start:adj_list_end]
401	non_adj_list = toks[:adj_list_start] + toks[adj_list_end:]
402	return adj_list, non_adj_list
403
404
405def equal_except_adj_list_sequence(  # noqa: C901
406	rollout1: list[str],
407	rollout2: list[str],
408	do_except: bool = False,
409	when_counter_mismatch: ErrorMode = ErrorMode.EXCEPT,
410	when_len_mismatch: ErrorMode = ErrorMode.EXCEPT,
411) -> bool:
412	"""Returns if the rollout strings are equal, allowing for differently sequenced adjacency lists.
413
414	<ADJLIST_START> and <ADJLIST_END> tokens must be in the rollouts.
415	Intended ONLY for determining if two tokenization schemes are the same for rollouts generated from the same maze.
416	This function should NOT be used to determine if two rollouts encode the same `LatticeMaze` object.
417
418	# Warning: CTT False Positives
419	This function is not robustly correct for some corner cases using `CoordTokenizers.CTT`.
420	If rollouts are passed for identical tokenizers processing two slightly different mazes, a false positive is possible.
421	More specifically, some cases of zero-sum adding and removing of connections in a maze within square regions along the diagonal will produce a false positive.
422	"""
423	if len(rollout1) != len(rollout2):
424		if do_except:
425			when_len_mismatch.process(
426				f"Rollouts are not the same length: {len(rollout1)} != {len(rollout2)}",
427			)
428		return False
429	if ("<ADJLIST_START>" in rollout1) ^ ("<ADJLIST_START>" in rollout2):
430		if do_except:
431			err_msg: str = f"Rollouts do not have the same <ADJLIST_START> token: `{'<ADJLIST_START>' in rollout1 = }` != `{'<ADJLIST_START>' in rollout2 = }`"
432			raise ValueError(
433				err_msg,
434			)
435		return False
436	if ("<ADJLIST_END>" in rollout1) ^ ("<ADJLIST_END>" in rollout2):
437		if do_except:
438			err_msg: str = f"Rollouts do not have the same <ADJLIST_END> token: `{'<ADJLIST_END>' in rollout1 = }` != `{'<ADJLIST_END>' in rollout2 = }`"
439			raise ValueError(
440				err_msg,
441			)
442		return False
443
444	adj_list1, non_adj_list1 = get_token_regions(rollout1)
445	adj_list2, non_adj_list2 = get_token_regions(rollout2)
446	if non_adj_list1 != non_adj_list2:
447		if do_except:
448			when_len_mismatch.process(
449				f"Non-adjacency list tokens are not the same:\n{non_adj_list1}\n!=\n{non_adj_list2}",
450			)
451			err_msg: str = f"Non-adjacency list tokens are not the same:\n{non_adj_list1}\n!=\n{non_adj_list2}"
452			raise ValueError(
453				err_msg,
454			)
455		return False
456	counter1: Counter = Counter(adj_list1)
457	counter2: Counter = Counter(adj_list2)
458	counters_eq: bool = counter1 == counter2
459	if not counters_eq:
460		if do_except:
461			when_counter_mismatch.process(
462				f"Adjacency list counters are not the same:\n{counter1}\n!=\n{counter2}\n{counter1 - counter2 = }",
463			)
464		return False
465
466	return True
467
468
469def connection_list_to_adj_list(
470	conn_list: ConnectionList,
471	shuffle_d0: bool = True,
472	shuffle_d1: bool = True,
473) -> Int8[np.ndarray, "conn start_end=2 coord=2"]:
474	"""converts a `ConnectionList` (special lattice format) to a shuffled adjacency list
475
476	# Parameters:
477	- `conn_list: ConnectionList`
478		special internal format for graphs which are subgraphs of a lattice
479	- `shuffle_d0: bool`
480		shuffle the adjacency list along the 0th axis (order of pairs)
481	- `shuffle_d1: bool`
482		shuffle the adjacency list along the 1st axis (order of coordinates in each pair).
483		If `False`, all pairs have the smaller coord first.
484
485
486	# Returns:
487	- `Int8[np.ndarray, "conn start_end=2 coord=2"]`
488		adjacency list in the shape `(n_connections, 2, 2)`
489	"""
490	n_connections: int = conn_list.sum()
491	adj_list: Int8[np.ndarray, "conn start_end=2 coord=2"] = np.full(
492		(n_connections, 2, 2),
493		-1,
494		dtype=np.int8,
495	)
496
497	if shuffle_d1:
498		flip_d1: Float[np.ndarray, " conn"] = np.random.rand(n_connections)
499
500	# loop over all nonzero elements of the connection list
501	i: int = 0
502	for d, x, y in np.ndindex(conn_list.shape):
503		if conn_list[d, x, y]:
504			c_start: CoordTup = (x, y)
505			c_end: CoordTup = (
506				x + (1 if d == 0 else 0),
507				y + (1 if d == 1 else 0),
508			)
509			adj_list[i, 0] = np.array(c_start, dtype=np.int8)
510			adj_list[i, 1] = np.array(c_end, dtype=np.int8)
511
512			# flip if shuffling
513			# magic value is fine here
514			if shuffle_d1 and (flip_d1[i] > 0.5):  # noqa: PLR2004
515				c_s, c_e = adj_list[i, 0].copy(), adj_list[i, 1].copy()
516				adj_list[i, 0] = c_e
517				adj_list[i, 1] = c_s
518
519			i += 1
520
521	if shuffle_d0:
522		np.random.shuffle(adj_list)
523
524	return adj_list
525
526
527def is_connection(
528	edges: ConnectionArray,
529	connection_list: ConnectionList,
530) -> Bool[np.ndarray, "is_connection=edges"]:
531	"""Returns if each edge in `edges` is a connection (`True`) or wall (`False`) in `connection_list`."""
532	sorted_edges = np.sort(edges, axis=1)
533	edge_direction = (
534		(sorted_edges[:, 1, :] - sorted_edges[:, 0, :])[:, 0] == 0
535	).astype(np.int8)
536	return connection_list[edge_direction, sorted_edges[:, 0, 0], sorted_edges[:, 0, 1]]
537
538
539# string to coordinate representation
540# ==================================================