maxtext.input_pipeline.grain_tokenizer module

Tokenize Op used by Grain

class maxtext.input_pipeline.grain_tokenizer.TokenizerTransformBase(feature_names, sequence_length, tokenizer)

Bases: object

Base class for tokenizer transforms with common functionality.

Parameters:

• feature_names (str | Sequence[str])

• sequence_length (int | Sequence[int])

• tokenizer (SentencePieceTokenizer | HFTokenizer | TikTokenTokenizer)

feature_names: str | Sequence[str]

sequence_length: int | Sequence[int]

tokenizer: SentencePieceTokenizer | HFTokenizer | TikTokenTokenizer

class maxtext.input_pipeline.grain_tokenizer.TokenizeAndTrim(*args, **kwargs)

Bases: TokenizerTransformBase, MapTransform

Tokenize and trim features to sequence length.

Parameters:

• feature_names (str | Sequence[str])

• sequence_length (int | Sequence[int])

• tokenizer (SentencePieceTokenizer | HFTokenizer | TikTokenTokenizer)

map(element)

Maps to each element.

Parameters:

element (dict[str, Any])

Return type:

dict[str, Any]

class maxtext.input_pipeline.grain_tokenizer.TokenizeAndChunk(*args, **kwargs)

Bases: TokenizerTransformBase, FlatMapTransform

Tokenize and chunk features into multiple examples of sequence length.

Parameters:

• feature_names (str | Sequence[str])

• sequence_length (int | Sequence[int])

• tokenizer (SentencePieceTokenizer | HFTokenizer | TikTokenTokenizer)

• max_fan_out (int)

max_fan_out: int = 2048

flat_map(element)

Tokenize and chunk text into multiple examples of sequence length.

Parameters:

element (dict[str, Any])

Return type:

list[dict[str, Any]]