Class with LongLlama's memory cache Args: keys (`torch.FloatTensor` of shape `(batch_size, num_heads, mem_length, embed_size_per_hea

A convenient stateful Jax RNG wrapper. Can be used to wrap RNG inside pure function.

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LongLlamaDecoderLayer`] Args: config: LlamaCo

Pipeline that assigns documents to the indexes of the batch. To be more precise, each document from token_source is assigned an index in the

Huggingface dataset, where the dataset is loaded using the huggingface datasets.load_dataset() function.

JSON dataset, where each line of the data file contains a JSON dictionary with text fields.

r""" This is the configuration class to store the configuration of a [`~LLaMAModel`]. It is used to instantiate an LLaMA model according to th

Construct a LLaMA tokenizer. Based on byte-level Byte-Pair-Encoding. Args: vocab_file (`str`): Path to the vocabulary fil

Multi-headed attention from 'Attention Is All You Need' paper with FoT modifications

Class for configuring memory caches for LongLlama model. Args: positionals (`boolean`) Whether to use positional embeddi

Based on BaseModelOutputWithPast Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`)

For handling a single dataset. data_args should be separated using separate_data_args.

Custom msgpack checkpointer that saves large train states by serializing and saving tensors one by one in a streaming fashion. Avoids running

Example processor that converts a dictionary of texts into tokens.

AdamW optimizer with cosine schedule.

Base class for the data pipeline. token_source should generate tuples consisting of (tokens_from_doc, loss_mask, name_of_the_data_source)

Datset builder class.

An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models.

JIT traceable version of FlaxLogitsWarper that performs temperature scaling.

Utility class for initializing JAX distributed.

Pipeline that assigns multiple (k) indexes of the batch to a single doc. It achieves this by using DocAwareDataPipeline with k times smaller

Pipeline that concatenates docs/examples sequentially to create the batch.

r""" This is the configuration class to store the configuration of a [`LongLlamaModel`]. It is used to instantiate an LongLLaMA model accordin

Configurable optax optimizer factory.

PaLM optimizer factory. This optimizer implements the optimizer described in the PaLM paper: https://arxiv.org/abs/2204.02311

Filters out examples that have less than min_example_length tokens. Collects statistics about token_source (averaged over num_stat_samples).