A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions. It contains a batch (TensorDict

A container class used to store the whole dataset. Inside the class, each data sample will be stored in `Item` class. The properties of the datase

Encoder class for encoding queries using a specified model. Attributes: model_name (str): The name of the model. model_path

Sharding manager to support data resharding when using FSDP + Ulysses

Used to count mfu during training loop Example: flops_counter = FlopsCounter(config) flops_achieved, flops_promised = flops_

A container class used to store and manipulate a sample within a dataset. Information related to this sample during training/inference will be sto

A stopping criteria that halts the text generation process if any specified stop word is encountered. Inspired by https://discuss.huggingfac

ClipEncoder class for encoding queries using CLIP.

DataProtoFuture aims to eliminate actual data fetching on driver. By doing so, the driver doesn't have to wait for data so that asynchronous

r"""Dense retriever based on pre-built faiss index.

Evaluator is used to summarize the results of all metrics.

A checkpoint manager that saves and loads - model - optimizer - lr_scheduler - extra_states in a SPMD way. We save

A checkpoint manager that saves and loads - model - optimizer - lr_scheduler - extra_states in a SPMD way. We save

A memory buffer is a contiguous torch tensor that may combine multiple tensors sharing with the underlying memory. It must have a unique type

STEncoder class for encoding queries using SentenceTransformers. Attributes: model_name (str): The name of the model. model_

Adaptive KL controller described in the paper: https://arxiv.org/pdf/1909.08593.pdf

r"""BM25 retriever based on pre-built pyserini index.

Wrapper replacing None outputs by empty tensors, which allows the use of checkpointing.

Fixed KL controller.

Creates a HF model that inherits from REPLUG_Generation class

r"""A tool class used to build an index used in retrieval.

LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla

LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev

r"""Multi-modal retriever based on pre-built faiss index.

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

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

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

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] This model definition supports pip

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

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

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

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`] This model definition supports pip

PromptCompressor is designed for compressing prompts based on a given language model. This class initializes with the language model and its

Merge logits of different docs in one batch. Reference: fastRAG

Note that this trainer runs on the driver process on a single CPU/GPU node.

Define a resource pool specification. Resource pool will be initialized first. Mapping