# embedding_metadata

key: chroma:document
string_value:
CODE:

def amplitude_to_db(
    spectrogram: np.ndarray,
    reference: float = 1.0,
    min_value: float = 1e-5,
    db_range: Optional[float] = None,
) -> np.ndarray:
    """
    Converts an amplitude spectrogram to the decibel scale. This computes `20 * log10(spectrogram / reference)`, using
    basic logarithm properties for numerical stability.

    The motivation behind applying the log function on the (mel) spectrogram is that humans do not hear loudness on a
    linear scale. Generally to double the perceived volume of a sound we need to put 8 times as much energy into it.
    This means that large variations in energy may not sound all that different if the sound is loud to begin with.
    This compression operation makes the (mel) spectrogram features match more closely what humans actually hear.

    Args:
        spectrogram (`np.ndarray`):
            The input amplitude (mel) spectrogram.
        reference (`float`, *optional*, defaults to 1.0):
            Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set
            the loudest part to 0 dB. Must be greater than zero.
        min_value (`float`, *optional*, defaults to `1e-5`):
            The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking
            `log(0)`. The default of `1e-5` corresponds to a minimum of -100 dB. Must be greater than zero.
        db_range (`float`, *optional*):
            Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the
            peak value and the smallest value will never be more than 80 dB. Must be greater than zero.

    Returns:
        `np.ndarray`: the spectrogram in decibels
    """
    if reference <= 0.0:
        raise ValueError("reference must be greater than zero")
    if min_value <= 0.0:
        raise ValueError("min_value must be greater than zero")

    reference = max(min_value, reference)

    spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None)
    spectrogram = 20.0 * (np.log10(spectrogram) - np.log10(reference))

    if db_range is not None:
        if db_range <= 0.0:
            raise ValueError("db_range must be greater than zero")
        spectrogram = np.clip(spectrogram, a_min=spectrogram.max() - db_range, a_max=None)



SUMMARY:
**Summary**

This function converts an amplitude spectrogram to the decibel scale using the formula `20 * log10(spectrogram / reference)`, with support for clipping and dynamic range limits.

**Main Purpose and Functionality**

*Converts amplitude spectrogram to decibels*

**Implementation Details**

* Computes `20 * log10(spectrogram / reference)`
* Clips values to minimum and maximum bounds:
	+ `np.clip(spectrogram, a_min=min_value, a_max=None)` (avoiding `log(0)`)
	+ `np.clip(spectrogram, a_min=spectrogram.max() - db_range, a_max=None)` with dynamic range limit (`db_range`)

**Notable Dependencies/Requirements**

* `np.ndarray` for input spectrogram
* `Optional[float]` for minimum value and dynamic range limit (`db_range`)
* Referenced functions:
	+ `np.log10`
	+ `np.max`

**Interface:**

* Input: `spectrogram`, `reference`, `min_value`, `db_range`
* Output: `np.ndarray` (spectrogram in decibels)


key: chroma:document
string_value:
CODE:

def amplitude_to_db_batch(
    spectrogram: np.ndarray, reference: float = 1.0, min_value: float = 1e-5, db_range: Optional[float] = None
) -> np.ndarray:
    """
    Converts a batch of amplitude spectrograms to the decibel scale. This computes `20 * log10(spectrogram / reference)`,
    using basic logarithm properties for numerical stability.

    The function supports batch processing, where each item in the batch is an individual amplitude (mel) spectrogram.

    Args:
        spectrogram (`np.ndarray`):
            The input batch of amplitude (mel) spectrograms. Expected shape is (batch_size, *spectrogram_shape).
        reference (`float`, *optional*, defaults to 1.0):
            Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set
            the loudest part to 0 dB. Must be greater than zero.
        min_value (`float`, *optional*, defaults to `1e-5`):
            The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking
            `log(0)`. The default of `1e-5` corresponds to a minimum of -100 dB. Must be greater than zero.
        db_range (`float`, *optional*):
            Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the
            peak value and the smallest value will never be more than 80 dB. Must be greater than zero.

    Returns:
        `np.ndarray`: the batch of spectrograms in decibels
    """
    if reference <= 0.0:
        raise ValueError("reference must be greater than zero")
    if min_value <= 0.0:
        raise ValueError("min_value must be greater than zero")

    reference = max(min_value, reference)

    spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None)
    spectrogram = 20.0 * (np.log10(spectrogram) - np.log10(reference))

    if db_range is not None:
        if db_range <= 0.0:
            raise ValueError("db_range must be greater than zero")
        # Apply db_range clipping per batch item
        max_values = spectrogram.max(axis=(1, 2), keepdims=True)
        spectrogram = np.clip(spectrogram, a_min=max_values - db_range, a_max=None)



SUMMARY:
Here is a concise summary of the provided function:

**Function Name:** `amplitude_to_db_batch`

**Main Purpose:** Converts a batch of amplitude spectrograms to decibel scale.

**Key Implementation Details:**

* Computes `20 * log10(spectrogram / reference)` for each spectrogram in batch.
* Clips values to prevent taking `log(0)` at minimum value (`min_value`).
* Applicably clamps the dynamic range to `db_range` if specified.

**Notable Digits and Dependencies:**

* Input shape: (batch_size, *spectrogram_shape)
* Reference value must be greater than zero.
* Minimum value constraint ensures against taking log(0) for numerical stability.
* Optional parameter `db_range` allows custom dynamic range clipping.
* External dependencies:
	+ NumPy (`np`) for array operations.


key: chroma:document
string_value:
CODE:

def get_mel_filter_banks(
    nb_frequency_bins: int,
    nb_mel_filters: int,
    frequency_min: float,
    frequency_max: float,
    sample_rate: int,
    norm: Optional[str] = None,
    mel_scale: str = "htk",
) -> np.array:
    warnings.warn(
        "The function `get_mel_filter_banks` is deprecated and will be removed in version 4.31.0 of Transformers",
        FutureWarning,
    )
    return mel_filter_bank(
        num_frequency_bins=nb_frequency_bins,
        num_mel_filters=nb_mel_filters,
        min_frequency=frequency_min,
        max_frequency=frequency_max,
        sampling_rate=sample_rate,
        norm=norm,
        mel_scale=mel_scale,


SUMMARY:
**Summary**

`get_mel_filter_banks` calculates Mel filter banks for audio processing. It takes parameters such as frequency bins, mel filters, and sampling rate to generate a matrix of Mel coefficients. The function uses `mel_filter_bank` internal implementation.

**Key Details**

* Parameters:
	+ `num_frequency_bins`: int (number of frequency bins)
	+ `num_mel_filters`: int (number of mel filters)
	+ `min_frequency`, `max_frequency`: float (frequency limits)
	+ `sampling_rate`: int (audio sampling rate)
	+ `norm`: str, optional (normalization type, default: None)
	+ `mel_scale`: str (Mel scale, default: "htk")
* Returns: np.array (Mel filter bank matrix)
* Dependencies:
	+ `mel_filter_bank` internal implementation
	+ `Transformers` library version 4.31.0 and later ( deprecated warning)

**Context**: This function is marked as deprecated, indicating that it will be removed in a future version of the Transformers library.


key: chroma:document
string_value:
CODE:

def fram_wave(waveform: np.array, hop_length: int = 160, fft_window_size: int = 400, center: bool = True):
    """
    In order to compute the short time fourier transform, the waveform needs to be split in overlapping windowed
    segments called `frames`.

    The window length (window_length) defines how much of the signal is contained in each frame, while the hop length
    defines the step between the beginning of each new frame.


    Args:
        waveform (`np.array` of shape `(sample_length,)`):
            The raw waveform which will be split into smaller chunks.
        hop_length (`int`, *optional*, defaults to 160):
            Step between each window of the waveform.
        fft_window_size (`int`, *optional*, defaults to 400):
            Defines the size of the window.
        center (`bool`, defaults to `True`):
            Whether or not to center each frame around the middle of the frame. Centering is done by reflecting the
            waveform on the left and on the right.

    Return:
        framed_waveform (`np.array` of shape `(waveform.shape // hop_length , fft_window_size)`):
            The framed waveforms that can be fed to `np.fft`.
    """
    warnings.warn(
        "The function `fram_wave` is deprecated and will be removed in version 4.31.0 of Transformers",
        FutureWarning,
    )
    frames = []
    for i in range(0, waveform.shape[0] + 1, hop_length):
        if center:
            half_window = (fft_window_size - 1) // 2 + 1
            start = i - half_window if i > half_window else 0
            end = i + half_window if i < waveform.shape[0] - half_window else waveform.shape[0]
            frame = waveform[start:end]
            if start == 0:
                padd_width = (-i + half_window, 0)
                frame = np.pad(frame, pad_width=padd_width, mode="reflect")

            elif end == waveform.shape[0]:
                padd_width = (0, (i - waveform.shape[0] + half_window))
                frame = np.pad(frame, pad_width=padd_width, mode="reflect")

        else:
            frame = waveform[i : i + fft_window_size]
            frame_width = frame.shape[0]
            if frame_width < waveform.shape[0]:
                frame = np.lib.pad(
                    frame, pad_width=(0, fft_window_size - frame_width), mode="constant", constant_values=0
                )
        frames.append(frame)

    frames = np.stack(frames, 0)


SUMMARY:
**Summary: fram_wave Function**

This function computes the short-time Fourier transform (STFT) of a given waveform by splitting it into overlapping windowed segments called "frames".

**Primary Purpose:** Computs STFT of a waveform.

**Key Implementation Details:**

* Splits waveform into frames with specified hop length and FFT window size.
* Frames are centered around the middle of each frame using reflection.
* Function can be used to extract features for audio processing or analysis tasks.

**Notable Dependencies:**

* `np` (NumPy) library for numerical computations.
* `np.fft` module for Fourier transform calculations.


key: chroma:document
string_value:
CODE:

def stft(frames: np.array, windowing_function: np.array, fft_window_size: int = None):
    """
    Calculates the complex Short-Time Fourier Transform (STFT) of the given framed signal. Should give the same results
    as `torch.stft`.

    Args:
        frames (`np.array` of dimension `(num_frames, fft_window_size)`):
            A framed audio signal obtained using `audio_utils.fram_wav`.
        windowing_function (`np.array` of dimension `(nb_frequency_bins, nb_mel_filters)`:
            A array representing the function that will be used to reduces the amplitude of the discontinuities at the
            boundaries of each frame when computing the STFT. Each frame will be multiplied by the windowing_function.
            For more information on the discontinuities, called *Spectral leakage*, refer to [this
            tutorial]https://download.ni.com/evaluation/pxi/Understanding%20FFTs%20and%20Windowing.pdf
        fft_window_size (`int`, *optional*):
            Size of the window om which the Fourier transform is applied. This controls the frequency resolution of the
            spectrogram. 400 means that the fourrier transform is computed on windows of 400 samples. The number of
            frequency bins (`nb_frequency_bins`) used to divide the window into equal strips is equal to
            `(1+fft_window_size)//2`. An increase of the fft_window_size slows the calculus time proportionnally.

    Example:

    ```python
    >>> from transformers.audio_utils import stft, fram_wave
    >>> import numpy as np

    >>> audio = np.random.rand(50)
    >>> fft_window_size = 10
    >>> hop_length = 2
    >>> framed_audio = fram_wave(audio, hop_length, fft_window_size)
    >>> spectrogram = stft(framed_audio, np.hanning(fft_window_size + 1))
    ```

    Returns:
        spectrogram (`np.ndarray`):
            A spectrogram of shape `(num_frames, nb_frequency_bins)` obtained using the STFT algorithm
    """
    warnings.warn(
        "The function `stft` is deprecated and will be removed in version 4.31.0 of Transformers",
        FutureWarning,
    )
    frame_size = frames.shape[1]

    if fft_window_size is None:
        fft_window_size = frame_size

    if fft_window_size < frame_size:
        raise ValueError("FFT size must greater or equal the frame size")
    # number of FFT bins to store
    nb_frequency_bins = (fft_window_size >> 1) + 1

    spectrogram = np.empty((len(frames), nb_frequency_bins), dtype=np.complex64)
    fft_signal = np.zeros(fft_window_size)

    for f, frame in enumerate(frames):
        if windowing_function is not None:
            np.multiply(frame, windowing_function, out=fft_signal[:frame_size])
        else:
            fft_signal[:frame_size] = frame
        spectrogram[f] = np.fft.fft(fft_signal, axis=0)[:nb_frequency_bins]


SUMMARY:
**Summary: Time Domain Short-TFT Function**

Calculates the complex Short-Time Fourier Transform (STFT) of input frames using FFT.

**Functionality:**

1. Applies windowing function to each frame for spectral leakage reduction
2. Computes STFT on windowed frames, storing result in a spectrogram array

**Code Focus:**

* `np.fft.fft` for frequency transform
* Windowing function application and spectral leakage reduction
* Iterative computation of individual spectrogram lines

**Input Dependencies:**

* `frames`: input audio signal framed with specific hop length and fft window size
* `windowing_function`: optional, reduces amplitude of discontinuities at frame boundaries (`np.hanning` used as an example)

**Output:**
A 2D spectrogram array of shape `(num_frames, nb_frequency_bins)`, representing the complex STFT output.

Note: The function is deprecated and will be removed in version 4.31.0 of Transformers.


key: chroma:document
string_value:
CODE:

def get_dataloader_sampler(dataloader):
    if hasattr(dataloader, "batch_sampler") and dataloader.batch_sampler is not None:
        return get_dataloader_sampler(dataloader.batch_sampler)
    elif hasattr(dataloader, "sampler"):


SUMMARY:
**Summary:**

The `get_dataloader_sampler` function determines the sampler for a PyTorch DataLoader. It recursively checks for a `batch_sampler` attribute and if present, returns its connected sampler. If not, it directly returns the `sampler` of the DataLoader.

**Key Implementation Details:**

* Recursively uses `batch_sampler` to find a nested sampler if present.
* Directly returns `sampler` when no `batch_sampler` is found.
* Both `batch_sampler` and `sampler` are key attributes of PyTorch DataLoaders.

No notable dependencies or requirements.


key: chroma:document
string_value:
CODE:

def atleast_1d(tensor_or_array: Union[torch.Tensor, np.ndarray]):
    if isinstance(tensor_or_array, torch.Tensor):
        if hasattr(torch, "atleast_1d"):
            tensor_or_array = torch.atleast_1d(tensor_or_array)
        elif tensor_or_array.ndim < 1:
            tensor_or_array = tensor_or_array[None]
    else:
        tensor_or_array = np.atleast_1d(tensor_or_array)


SUMMARY:
**Summary: `atleast_1d` Function**

**Purpose:** Upscales a 0D or lower-dimensional tensor/array to at least 1D.

**Primary Purpose:**
Ensures input is a 1D tensor or array for further processing.

**Key Implementation Details:**

* Checks if input `tensor_or_array` is a PyTorch tensor (<code>`torch.Tensor</code>).
* If PyTorch, attempts to use built-in `atleast_1d` function (<code>`torch.atleast_1d`).
* If not, adds singleton dimension (`None`) if necessary.
* If input is NumPy array, uses same scaling logic as PyTorch.

**Notable Dependencies/Requirements:**

*: Both `torch` and `numpy` libraries are required for this function to work.


key: chroma:document
string_value:
CODE:

def torch_pad_and_concatenate(tensor1, tensor2, padding_index=-100):
    """Concatenates `tensor1` and `tensor2` on first axis, applying padding on the second if necessary."""
    tensor1 = atleast_1d(tensor1)
    tensor2 = atleast_1d(tensor2)

    if len(tensor1.shape) == 1 or tensor1.shape[1] == tensor2.shape[1]:
        return torch.cat((tensor1, tensor2), dim=0)

    # Let's figure out the new shape
    new_shape = (tensor1.shape[0] + tensor2.shape[0], max(tensor1.shape[1], tensor2.shape[1])) + tensor1.shape[2:]

    # Now let's fill the result tensor
    result = tensor1.new_full(new_shape, padding_index)
    result[: tensor1.shape[0], : tensor1.shape[1]] = tensor1
    result[tensor1.shape[0] :, : tensor2.shape[1]] = tensor2


SUMMARY:
**Summary:**

### Purpose and Functionality:

The `torch_pad_and_concatenate` function concatenates two tensors (`tensor1` and `tensor2`) on the first axis, applying padding on the second axis if necessary.

### Key Implementation Details:

* Uses `torch.cat` to concatenate tensors along the first axis.
* Calculates new shape based on tensor dimensions.
* Creates a result tensor with new shape using `tensor.new_full`.
* Fills result tensor with elements from `tensor1` and `tensor2`.

### Notable Dependencies or Requirements:

* `atleast_1d`: ensures input tensors have at least one dimension.
* `torch.cat`: concatenates tensors along the specified axis.
* `torch.tensor`: creates a new PyTorch tensor.

**RAG Entry (searchable and contextual):**

`torch_pad_and_concatenate(tensor1, tensor2, padding_index=-100)`
`Concatenates tensors on first axis, applying padding on second if necessary.`


key: chroma:document
string_value:
CODE:

def numpy_pad_and_concatenate(array1, array2, padding_index=-100):
    """Concatenates `array1` and `array2` on first axis, applying padding on the second if necessary."""
    array1 = atleast_1d(array1)
    array2 = atleast_1d(array2)

    if len(array1.shape) == 1 or array1.shape[1] == array2.shape[1]:
        return np.concatenate((array1, array2), axis=0)

    # Let's figure out the new shape
    new_shape = (array1.shape[0] + array2.shape[0], max(array1.shape[1], array2.shape[1])) + array1.shape[2:]

    # Now let's fill the result tensor
    result = np.full_like(array1, padding_index, shape=new_shape)
    result[: array1.shape[0], : array1.shape[1]] = array1
    result[array1.shape[0] :, : array2.shape[1]] = array2


SUMMARY:
Here is a concise summary of the provided code chunk:

**Function Summary: `numpy_pad_and_concatenate`**

**Purpose:** Concatenates two 4D arrays (`array1` and `array2`) on the first axis, padding the second array if necessary.

**Key Implementation Details:**

* Verifies input dimensions and pads `array2` to match `array1`'s shape.
* Fills a new tensor with padded values using NumPy's full_like function.
* Copies data from `array1` and `array2` into the resulting tensor, ensuring alignment along axis 1.

**Notable Dependencies:**

* Numpy (`np`)
* `atleast_1d` function
* `full_like` function

**Functions Called:** `np.concatenate`, `np.full_like`

This function appears to be part of a larger library or module providing NumPy-specific functionality for array manipulation.


key: chroma:document
string_value:
CODE:

def nested_concat(tensors, new_tensors, padding_index=-100):
    """
    Concat the `new_tensors` to `tensors` on the first dim and pad them on the second if needed. Works for tensors or
    nested list/tuples/dict of tensors.
    """
    if not (isinstance(tensors, torch.Tensor) and isinstance(new_tensors, torch.Tensor)):
        assert (
            type(tensors) is type(new_tensors)
        ), f"Expected `tensors` and `new_tensors` to have the same type but found {type(tensors)} and {type(new_tensors)}."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_concat(t, n, padding_index=padding_index) for t, n in zip(tensors, new_tensors))
    elif isinstance(tensors, torch.Tensor):
        return torch_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index)
    elif isinstance(tensors, Mapping):
        return type(tensors)(
            {k: nested_concat(t, new_tensors[k], padding_index=padding_index) for k, t in tensors.items()}
        )
    elif isinstance(tensors, np.ndarray):
        return numpy_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index)
    else:


SUMMARY:
**Function Summary**

`nested_concat` functions concatanates and pads `new_tensors` to `tensors` in a nested manner, supporting tensors, lists/tuples, dictionaries, and numpy arrays.

**Key Implementation Details**

*   Checks the types of `tensors` and `new_tensors`, asserting they must be the same.
*   Supports concatenation on first dimension (dim=0) using `torch_pad_and_concatenate`.
*   Recursively handles nested tensors on all supported data structures (`list`, tuple, dict, np.ndarray).
*   Raises a `TypeError` for unsupported data types.

**Notable Dependencies/Requirements**

*   `torch.Tensor`
*   `torch.pad_and_concatenate`
*   `type(tensors)`
*   `type(new_tensors)`
*   `zip(tensors, new_tensors)`


key: chroma:document
string_value:
CODE:

def find_batch_size(tensors):
    """
    Find the first dimension of a tensor in a nested list/tuple/dict of tensors.
    """
    if isinstance(tensors, (list, tuple)):
        for t in tensors:
            result = find_batch_size(t)
            if result is not None:
                return result
    elif isinstance(tensors, Mapping):
        for key, value in tensors.items():
            result = find_batch_size(value)
            if result is not None:
                return result
    elif isinstance(tensors, torch.Tensor):
        return tensors.shape[0] if len(tensors.shape) >= 1 else None
    elif isinstance(tensors, np.ndarray):


SUMMARY:
**Summary:**

`find_batch_size` function recursively traverses a nested data structure (list, tuple, dict, PyTorch tensor, or NumPy array) to find the first dimension's index.

**Key Implementation Details:**

- Recursively checks list and tuple elements
- Recursively checks dictionary values
- Returns tensor shape[0] if available

**Notable Dependencies/Requirements:**

- `torch` (torch.Tensor type checking)
- `numpy` (np.ndarray type checking)

This function aims to extract the batch size from a data structure, allowing for handling of different tensor representations. The implementation utilizes Python's built-in data structures and tensors' shape attributes.


key: chroma:document
string_value:
CODE:

def nested_numpify(tensors):
    "Numpify `tensors` (even if it's a nested list/tuple/dict of tensors)."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_numpify(t) for t in tensors)
    if isinstance(tensors, Mapping):
        return type(tensors)({k: nested_numpify(t) for k, t in tensors.items()})

    t = tensors.cpu()
    if t.dtype == torch.bfloat16:
        # As of Numpy 1.21.4, NumPy does not support bfloat16 (see
        # https://github.com/numpy/numpy/blob/a47ecdea856986cd60eabbd53265c2ca5916ad5d/doc/source/user/basics.types.rst ).
        # Until Numpy adds bfloat16, we must convert float32.
        t = t.to(torch.float32)


SUMMARY:
**numpify Function Summary**
=====================================

### Purpose and Functionality

The `nested_numpify` function recursively converts PyTorch tensors and their nested structures into NumPy arrays, handling types including lists, tuples, dictionaries, and tensor data types.

### Key Implementation Details

*   Recursively traverses input tensors (lists, tuples, dictionaries) and conversions are performed for each element or key-value pair.
*   Tensors with `torch.bfloat16` data type are converted to `torch.float32`.
*   The resulting NumPy array is retrieved from the converted PyTorch tensor (`t.numpy()`).

### Notable Dependencies and Requirements

*   References: None
*   External Dependencies:
    *   PyTorch library
    *   NumPy library (specifically, version 1.21.4 for `bfloat16` support)
*   Implemented Interfaces:
    *   Converts tensors to NumPy arrays
    *   Handles recursive traversal of tensor structures


key: chroma:document
string_value:
CODE:

def nested_detach(tensors):
    "Detach `tensors` (even if it's a nested list/tuple/dict of tensors)."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_detach(t) for t in tensors)
    elif isinstance(tensors, Mapping):
        return type(tensors)({k: nested_detach(t) for k, t in tensors.items()})


SUMMARY:
**Function Summary:**

`nested_detach`: Detaches nested tensors from gradients, handling multiple data types (lists/tuples/dicts of tensors).

**Key Details:**

* Uses recursion to handle nested structures
* Handles three data types:
	+ Lists/Tuples: recurse on individual elements using `nested_detach(t) for t in tensors`
	+ Dictionaries: use dictionary comprehension with `k: nested_detach(t)` for key-value pairs
	+ PyTorch Tensors: use `detach()` method directly
* Returns detached tensor(s)

**Notable Dependencies/Interfaces:**

* `torch.Tensor` interface
* `list`, `tuple`, and `Mapping` data types


key: chroma:document
string_value:
CODE:

def nested_xla_mesh_reduce(tensors, name):
    if is_torch_xla_available():
        import torch_xla.core.xla_model as xm

        if isinstance(tensors, (list, tuple)):
            return type(tensors)(nested_xla_mesh_reduce(t, f"{name}_{i}") for i, t in enumerate(tensors))
        if isinstance(tensors, Mapping):
            return type(tensors)(
                {k: nested_xla_mesh_reduce(t, f"{name}_{i}") for i, (k, t) in enumerate(tensors.items())}
            )

        tensors = atleast_1d(tensors)
        return xm.mesh_reduce(name, tensors, torch.cat)
    else:


SUMMARY:
**Summary:**

The `nested_xla_mesh_reduce` function reduces a nested tensor by leveraging PyTorch's XLA acceleration under the hood (`torch_xla.core.xla_model`). It recursively applies mesh reduction to tensors in various data structures (lists, tuples, and dictionaries), utilizing `torch.cat` for concatenation and `atleast_1d` for conversion to 1D tensors.

**Key Implementation Details:**

* Recursively reduces nested tensors using XLA's `mesh_reduce`
* Utilizes `torch_cat` for concatenation
* Ensures input tensors have shape compatible with mesh reduction

**Notable Dependencies/Requirements:**

* PyTorch (`torch_xla.core.xla_model`)
* `torch.cat` function


key: chroma:document
string_value:
CODE:

def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any:
    try:
        if isinstance(tensor, (tuple, list)):
            return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor)
        if isinstance(tensor, Mapping):
            return type(tensor)({k: distributed_concat(t, num_total_examples) for k, t in tensor.items()})
        tensor = atleast_1d(tensor).contiguous()
        output_tensors = [tensor.clone() for _ in range(dist.get_world_size())]
        dist.all_gather(output_tensors, tensor)
        concat = torch.cat(output_tensors, dim=0)

        # truncate the dummy elements added by SequentialDistributedSampler
        if num_total_examples is not None:
            concat = concat[:num_total_examples]
        return concat
    except AssertionError:


SUMMARY:
**Function Summary: Distributed Concatenation**

* **Purpose**: Concatenates tensors in a distributed manner, handling various input types (tensor, tuple, list, and mapping).
* **Implementation Details**:
	+ Checks input type and applies transformation accordingly.
	+ Splits tensor into individual ranks using `dist.get_world_size()` and collects them via `dist.all_gather`.
	+ Concatenates tensors along dimension 0 using `torch.cat`.
	+ Optionally truncates concatenated output based on `num_total_examples`.
* **Notable Dependencies**:
	+ `dist` (distributed training framework)
	+ ` atleast_1d`, `contiguous`, and `Tensor` (PyTorch tensor manipulation)
	+ `Mapping` (Python dictionary)

Note: The function raises an AssertionError if not currently using distributed training.


key: chroma:document
string_value:
CODE:

def distributed_broadcast_scalars(
    scalars: List[Union[int, float]],
    num_total_examples: Optional[int] = None,
    device: Optional[torch.device] = torch.device("cuda"),
) -> torch.Tensor:
    try:
        tensorized_scalar = torch.tensor(scalars).to(device)
        output_tensors = [tensorized_scalar.clone() for _ in range(dist.get_world_size())]
        dist.all_gather(output_tensors, tensorized_scalar)
        concat = torch.cat(output_tensors, dim=0)

        # truncate the dummy elements added by SequentialDistributedSampler
        if num_total_examples is not None:
            concat = concat[:num_total_examples]
        return concat
    except AssertionError:


SUMMARY:
**Summary: Distributed Broadcast Scalar Function**

**Purpose and Functionality:**
`distributed_broadcast_scalars`: Broadcasting scalars to all GPUs in a distributed PyTorch environment.

**Key Implementation Details:**

* Creates a tensor from input `scalars` using `torch.tensor()` and moves it to the specified device (`cuda` by default).
* Synchronizes data across all processes using `dist.all_gather()` and creates a broadcasted tensor with size equal to the total number of examples.
* Optionally truncates the output tensor based on the provided `num_total_examples`.

**Notable Dependencies and Requirements:**

* `torch`: PyTorch library for distributed training.
* `dist`: Distributed PyTorch module for parallelization.
* `torch.device`: PyTorch device manager.

**Return Value:**
A broadcasted tensor containing the input scalars.


key: chroma:document
string_value:
CODE:

def reissue_pt_warnings(caught_warnings):
    # Reissue warnings
    if len(caught_warnings) > 1:
        for w in caught_warnings:
            if w.category is not UserWarning:


SUMMARY:
**Summary:**

`reissue_pt_warnings` function reissues PTLiveTemplateWarning instances that are not `UserWarnings`. It iterates through the input list `caught_warnings`, skipping `UserWarning` instances.

**Key Implementation Details:**

* Iterates over `caught_warnings` to reissue non-`UserWarning` warnings
* Uses `warnings.warn` with warning category and message from the input warning instance

**Notable Dependencies/Requirements:**

* `warnings` module (specifically, `warnings.warn` function)
* `PTLiveTemplateWarning` class (input parameter is an instance of this class)


key: chroma:document
string_value:
CODE:

class DistributedSamplerWithLoop(DistributedSampler):
    """
    Like a torch.utils.data.distributed.DistributedSampler` but loops at the end back to the beginning of the shuffled
    samples to make each process have a round multiple of batch_size samples.

    Args:
        dataset (`torch.utils.data.Dataset`):
            Dataset used for sampling.
        batch_size (`int`):
            The batch size used with this sampler
        kwargs (`Dict[str, Any]`, *optional*):
            All other keyword arguments passed to `DistributedSampler`.
    """

    def __init__(self, dataset, batch_size, **kwargs):
        super().__init__(dataset, **kwargs)
        self.batch_size = batch_size

    def __iter__(self):
        indices = list(super().__iter__())
        remainder = 0 if len(indices) % self.batch_size == 0 else self.batch_size - len(indices) % self.batch_size
        # DistributedSampler already added samples from the beginning to make the number of samples a round multiple
        # of the world size, so we skip those.
        start_remainder = 1 if self.rank < len(self.dataset) % self.num_replicas else 0
        indices += indices[start_remainder : start_remainder + remainder]


SUMMARY:
**DistributedSamplerWithLoop Class Summary**

### Overview

A custom `DistributedSampler` class that loops at the end back to the beginning of shuffled samples, ensuring each process receives a round multiple of batch_size samples.

### Key Functionality

- Iterates over dataset with looped sampling
- Adjusts indices to make number of samples a round multiple of batch size for each process
- Skips initial batch with remaining indices for distributed sampler's "start" offset

### Notable Implementation Details

- Uses `super().__iter__()` to inherit original DistributedSampler iteration behavior
- Calculates adjusted remainder and adds it to the end of the iterator's sample indices

### Dependencies and Requirements

* `torch.utils.data.Dataset`
* `train` - `self.dataset` for dataset usage in `super().__init__()
* Inherits from `torch.utils.data.distributed.DistributedSampler`

### Notable Signature
```python
def __iter__(self):
    # ...
```
Note: `Indices` calculated and returned via iterator object.


key: chroma:document
string_value:
CODE:

class EvalLoopContainer:
    """
    Container to store intermediate results of evaluation loop

    Args:
        do_nested_concat (`bool`, *optional*, defaults to `True`):
            If set to `True`, each iteration will recursively concatenate a new object containing tensors to
            the existing stored tensors, provided that the structure of the existing object and the new one
            are identical. If set to `False`, all newly added tensors will be stored in a list.
        padding_index (`int`, *optional*, defaults to -100):
            Value used to pad tensors of different shapes when `do_nested_concat=True`.
    """

    def __init__(self, do_nested_concat: bool = True, padding_index: int = -100):
        self.do_nested_concat = do_nested_concat
        self.padding_index = padding_index
        self.tensors = None
        self.arrays = None

    def add(self, tensors) -> None:
        """Add tensors to the stored objects. If `do_nested_concat=True`, the tensors will be concatenated recursively."""
        if self.tensors is None:
            self.tensors = tensors if self.do_nested_concat else [tensors]
        elif self.do_nested_concat:
            self.tensors = nested_concat(self.tensors, tensors, padding_index=self.padding_index)
        else:
            self.tensors.append(tensors)

    def to_cpu_and_numpy(self) -> None:
        """Move tensors in stored objects to CPU and convert them to numpy arrays."""

        # Check if we have something to add, if not just return
        if self.tensors is None:
            return

        new_arrays = nested_numpify(self.tensors)
        if self.arrays is None:
            self.arrays = new_arrays
        elif self.do_nested_concat:
            self.arrays = nested_concat(self.arrays, new_arrays, padding_index=self.padding_index)
        else:
            self.arrays.extend(new_arrays)

        # reset device tensors after adding to cpu
        self.tensors = None

    def get_arrays(self):
        """Returns the numpified and moved to CPU stored objects."""
        self.to_cpu_and_numpy()


SUMMARY:
**EvalLoopContainer Class Summary**

The `EvalLoopContainer` class is designed to store intermediate results of evaluation loops. It provides methods for adding tensors, handling nested concatenation, and moving data to CPU.

**Key Features:**

* Storages tensors in either a single object or a list (optional)
* Recursively concatenates new tensors when `do_nested_concat=True`
* Moves tensors to CPU and converts them to numpy arrays (`to_cpu_and_numpy`)
* Provides a method to retrieve numpified and moved-to-CPU stored objects (`get_arrays`)

**Implementation Details:**

* Uses the `nested_concat` function for recursive concatenation and `nested_numpify` function for converting tensors to numpy arrays.
* Stores tensors in either `tensors` or `arrays` attribute, depending on the `do_nested_concat` flag.

**Notable Dependencies/Interfaces:**

* `nested_concat`: function for recursive tensor concatenation
* `nested_numpify`: function for converting tensors to numpy arrays

Example usage:
```python
container = EvalLoopContainer(do_nested_concat=True)
container.add(tensors1)
container.add(tensors2)
result = container.get_arrays()
```
Note: The exact implementation of `nested concat` and `nested numpify` functions are not included in this summary, as they are assumed to be separate modules or functions.


key: chroma:document
string_value:
CODE:

class SequentialDistributedSampler(Sampler):
    """
    Distributed Sampler that subsamples indices sequentially, making it easier to collate all results at the end.

    Even though we only use this sampler for eval and predict (no training), which means that the model params won't
    have to be synced (i.e. will not hang for synchronization even if varied number of forward passes), we still add
    extra samples to the sampler to make it evenly divisible (like in `DistributedSampler`) to make it easy to `gather`
    or `reduce` resulting tensors at the end of the loop.
    """

    def __init__(self, dataset, num_replicas=None, rank=None, batch_size=None):
        warnings.warn(
            "SequentialDistributedSampler is deprecated and will be removed in v5 of Transformers.",
            FutureWarning,
        )
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        num_samples = len(self.dataset)
        # Add extra samples to make num_samples a multiple of batch_size if passed
        if batch_size is not None:
            self.num_samples = int(math.ceil(num_samples / (batch_size * num_replicas))) * batch_size
        else:
            self.num_samples = int(math.ceil(num_samples / num_replicas))
        self.total_size = self.num_samples * self.num_replicas
        self.batch_size = batch_size

    def __iter__(self):
        indices = list(range(len(self.dataset)))

        # add extra samples to make it evenly divisible
        indices += indices[: (self.total_size - len(indices))]
        assert (
            len(indices) == self.total_size
        ), f"Indices length {len(indices)} and total size {self.total_size} mismatched"

        # subsample
        indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples]
        assert (
            len(indices) == self.num_samples
        ), f"Indices length {len(indices)} and sample number {self.num_samples} mismatched"

        return iter(indices)

    def __len__(self):


SUMMARY:
**SequentialDistributedSampler**

A distributed sampler for subsampling indices sequentially, ensuring evenly divisible samples for easy collation.

### Primary Purpose

* Subsample indices from a dataset to facilitate even distribution and easier collation of results during evaluation.

### Key Implementation Details

* Distributes indices among `num_replicas` using `rank`-based ordering.
* Ensures `num_samples` is a multiple of batch size (`batch_size`) when specified.

### Notable Dependencies or Requirements

* Distributed package (e.g., PyTorch, DDP)
* Available dimensions:
	+ `dataset`: input dataset object
	+ `dist`: distributed package interface (PyTorch, torch.distributed)

### Initializer Details

* Takes in:
	+ `dataset`: input dataset object
	+ `num_replicas` (optional): number of replicas to distribute indices among
	+ `rank` (optional): current rank for indexing distribution
	+ `batch_size` (optional): batch size for subsampling, affecting total sample count

### Sampler Operations

* Returns an iterator producing indices with a rank-ordered sequence.
* Employs list concatenation and slicing to add extra samples for even distribution.


key: chroma:document
string_value:
CODE:

def get_tpu_sampler(dataset: torch.utils.data.Dataset, batch_size: int):
    if xm.xrt_world_size() <= 1:
        return RandomSampler(dataset)


SUMMARY:
**get_tpu_sampler Function Summary**

**Purpose and Functionality**
This function returns a sampler for generating batches from a given `torch.utils.data.Dataset`. The sampler is optimized for multi-gpu (TPU) distributed training.

**Key Implementation Details:**

* Returns either a `RandomSampler` or a `DistributedSampler` depending on the number of workers (`xrt_world_size()`).
	+ If one worker, returns `RandomSampler`.
	+ Otherwise, returns `DistributedSampler` with customized number of replicas and rank.

**Notable Dependencies and Requirements:**

* `torch.utils.data.Dataset`
* `torch.distributed` (specifically, `xm` module)
* External dependency:
	+ PyTorch library (not explicitly listed in the code snippet)

**Notes:** The function assumes access to TPU (Tensor Processing Unit) resources for distributed training.


key: chroma:document
string_value:
CODE:

def nested_new_like(arrays, num_samples, padding_index=-100):
    """Create the same nested structure as `arrays` with a first dimension always at `num_samples`."""
    if isinstance(arrays, (list, tuple)):
        return type(arrays)(nested_new_like(x, num_samples) for x in arrays)


SUMMARY:
**Summary**

The `nested_new_like` function creates a new array with the same nested structure as `arrays`, but with the first dimension always of length `num_samples`. It handles recursive nesting and inserts padding (`-100`) if necessary.

**Key Details:**

*   Creates a nested array by applying the function to each element in `arrays`.
*   Fills any non-full arrays (1st dimension < num\_samples) with padding `-100`.

**Implementation:**

*   Returns:
    *   A new array with the same shape as `arrays`, but with the first dimension expanded to length `num_samples`.
    *   If `arrays` is not a nested structure, it fills only the requested dimensions (`shape[1:]`) with padding.


key: chroma:document
string_value:
CODE:

def expand_like(arrays, new_seq_length, padding_index=-100):
    """Expand the `arrays` so that the second dimension grows to `new_seq_length`. Uses `padding_index` for padding."""
    result = np.full_like(arrays, padding_index, shape=(arrays.shape[0], new_seq_length) + arrays.shape[2:])
    result[:, : arrays.shape[1]] = arrays


SUMMARY:
**Summary:**

`expand_like` function expands the second dimension of input `arrays` to `new_seq_length`. 
It pads with `-100` and fills first columns with original data.

**Key Implementation Details:**

- Uses `padding_index` for padding
- Creates a new array with shape `(shape[0], new_seq_length, ...)` filled with `padding_index`
- Fills first `arrays.shape[1]` columns with original data from `arrays`

**Dependencies and Requirements:**

- Import of `numpy` library for `np.full_like` function
- Input `arrays` is expected to have shape `(shape[0], ...)` where `...` represents additional dimensions


key: chroma:document
string_value:
CODE:

def nested_truncate(tensors, limit):
    "Truncate `tensors` at `limit` (even if it's a nested list/tuple/dict of tensors)."
    if isinstance(tensors, (list, tuple)):
        return type(tensors)(nested_truncate(t, limit) for t in tensors)
    if isinstance(tensors, Mapping):
        return type(tensors)({k: nested_truncate(t, limit) for k, t in tensors.items()})



SUMMARY:
**Code Summary**

`nested_truncate` function truncates input `tensors` (list, tuple, dict, or list/dict elements) to a specified `limit`. Recursively applies truncation to nested data structures.

**Implementation Details**

*   Truncates elements of type list, tuple, and Mapping
*   Uses slicing (`tensors[:limit]`) for flat lists/tuples

**Key Methods/Ideals**

*   Used with tensors, which can be recursively defined (e.g., `type(tensors)`)
*   Inherited/implemented interfaces: None


key: chroma:document
string_value:
CODE:

class DistributedTensorGatherer:
    """
    A class responsible for properly gathering tensors (or nested list/tuple of tensors) on the CPU by chunks.

    If our dataset has 16 samples with a batch size of 2 on 3 processes and we gather then transfer on CPU at every
    step, our sampler will generate the following indices:

        `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1]`

    to get something of size a multiple of 3 (so that each process gets the same dataset length). Then process 0, 1 and
    2 will be responsible of making predictions for the following samples:

        - P0: `[0, 1, 2, 3, 4, 5]`
        - P1: `[6, 7, 8, 9, 10, 11]`
        - P2: `[12, 13, 14, 15, 0, 1]`

    The first batch treated on each process will be

        - P0: `[0, 1]`
        - P1: `[6, 7]`
        - P2: `[12, 13]`

    So if we gather at the end of the first batch, we will get a tensor (nested list/tuple of tensor) corresponding to
    the following indices:

        `[0, 1, 6, 7, 12, 13]`

    If we directly concatenate our results without taking any precautions, the user will then get the predictions for
    the indices in this order at the end of the prediction loop:

        `[0, 1, 6, 7, 12, 13, 2, 3, 8, 9, 14, 15, 4, 5, 10, 11, 0, 1]`

    For some reason, that's not going to roll their boat. This class is there to solve that problem.

    Args:
        world_size (`int`):
            The number of processes used in the distributed training.
        num_samples (`int`):
            The number of samples in our dataset.
        make_multiple_of (`int`, *optional*):
            If passed, the class assumes the datasets passed to each process are made to be a multiple of this argument
            (by adding samples).
        padding_index (`int`, *optional*, defaults to -100):
            The padding index to use if the arrays don't all have the same sequence length.
    """

    def __init__(self, world_size, num_samples, make_multiple_of=None, padding_index=-100):
        warnings.warn(
            "DistributedTensorGatherer is deprecated and will be removed in v5 of Transformers.",
            FutureWarning,
        )
        self.world_size = world_size
        self.num_samples = num_samples
        total_size = world_size if make_multiple_of is None else world_size * make_multiple_of
        self.total_samples = int(np.ceil(num_samples / total_size)) * total_size
        self.process_length = self.total_samples // world_size
        self._storage = None
        self._offsets = None
        self.padding_index = padding_index

    def add_arrays(self, arrays):
        """
        Add `arrays` to the internal storage, Will initialize the storage to the full size at the first arrays passed
        so that if we're bound to get an OOM, it happens at the beginning.
        """
        if arrays is None:
            return
        if self._storage is None:
            self._storage = nested_new_like(arrays, self.total_samples, padding_index=self.padding_index)
            self._offsets = list(range(0, self.total_samples, self.process_length))

        slice_len, self._storage = self._nested_set_tensors(self._storage, arrays)
        for i in range(self.world_size):
            self._offsets[i] += slice_len

    def _nested_set_tensors(self, storage, arrays):
        if isinstance(arrays, (list, tuple)):
            result = [self._nested_set_tensors(x, y) for x, y in zip(storage, arrays)]
            return result[0][0], type(arrays)(r[1] for r in result)
        assert (
            arrays.shape[0] % self.world_size == 0
        ), f"Arrays passed should all have a first dimension multiple of {self.world_size}, found {arrays.shape[0]}."

        slice_len = arrays.shape[0] // self.world_size
        for i in range(self.world_size):
            if len(arrays.shape) == 1:
                storage[self._offsets[i] : self._offsets[i] + slice_len] = arrays[i * slice_len : (i + 1) * slice_len]
            else:
                # Expand the array on the fly if needed.
                if len(storage.shape) > 1 and storage.shape[1] < arrays.shape[1]:
                    storage = expand_like(storage, arrays.shape[1], padding_index=self.padding_index)
                storage[self._offsets[i] : self._offsets[i] + slice_len, : arrays.shape[1]] = arrays[
                    i * slice_len : (i + 1) * slice_len
                ]
        return slice_len, storage

    def finalize(self):
        """
        Return the properly gathered arrays and truncate to the number of samples (since the sampler added some extras
        to get each process a dataset of the same length).
        """
        if self._storage is None:
            return
        if self._offsets[0] != self.process_length:
            logger.warning("Not all data has been set. Are you sure you passed all values?")


SUMMARY:
**DistributedTensorGatherer Class Summary**

### Main Purpose and Functionality

The `DistributedTensorGatherer` class is designed to efficiently gather tensors on the CPU by chunks in a distributed environment. It ensures that each process receives data of the same length, avoiding interleaving predictions from different samples.

### Key Implementation Details

* The class initializes internal storage for gathering arrays using `nested_new_like`.
* The `_nested_set_tensors` method distributes arrays across processes, handling single-dimension and multi-dimensional arrays.
* The `finalize` method truncates the gathered data to the correct number of samples after adding padding indices.

### Notable Dependencies or Requirements

* The class uses NumPy (`np`) for tensor operations and warnings logging (`logger`).
* No external dependencies are required.
* The class implements interfaces based on custom `nested_new_like`, `nested_truncate`, and `expand_like` functions, which are assumed to be defined elsewhere in the codebase.

### Implementation Signature

```markdown
class DistributedTensorGatherer:
    def __init__(self, world_size, num_samples, make_multiple_of=None, padding_index=-100):
        pass
    
    def add_arrays(self, arrays):
        pass
    
    def _nested_set_tensors(self, storage, arrays):
        pass
    
    def finalize(self):
        pass
```


key: chroma:document
string_value:
CODE:

def get_length_grouped_indices(lengths, batch_size, mega_batch_mult=None, generator=None):
    """
    Return a list of indices so that each slice of `batch_size` consecutive indices correspond to elements of similar
    lengths. To do this, the indices are:

    - randomly permuted
    - grouped in mega-batches of size `mega_batch_mult * batch_size`
    - sorted by length in each mega-batch

    The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of
    maximum length placed first, so that an OOM happens sooner rather than later.
    """
    # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller.
    if mega_batch_mult is None:
        mega_batch_mult = min(len(lengths) // (batch_size * 4), 50)
        # Just in case, for tiny datasets
        if mega_batch_mult == 0:
            mega_batch_mult = 1

    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
    indices = torch.randperm(len(lengths), generator=generator)
    megabatch_size = mega_batch_mult * batch_size
    megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)]
    megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]

    # The rest is to get the biggest batch first.
    # Since each megabatch is sorted by descending length, the longest element is the first
    megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches]
    max_idx = torch.argmax(torch.tensor(megabatch_maximums)).item()
    # Switch to put the longest element in first position
    megabatches[0][0], megabatches[max_idx][0] = megabatches[max_idx][0], megabatches[0][0]



SUMMARY:
Here is a concise summary of the provided function:

**Function:** `get_length_grouped_indices`

**Purpose:** Generate a list of indices by randomly permuting and grouping `batch_size` consecutive indices into mega-batches based on element lengths.

**Key Implementation:**

* Randomly permutes indices using `torch.randperm`
* Groups indices into mega-batches based on `mega_batch_mult * batch_size`
* Sorts each mega-batch by descending length
* Prioritizes the longest element in the first position of each mega-batch

**Notable Dependencies and Requirements:**

* pytorch (`torch`); requires a GPU-based distributed sampler for random seed generation.
* `generator` (optional): allows external seed management.


key: chroma:document
string_value:
CODE:

class LengthGroupedSampler(Sampler):
    r"""
    Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
    keeping a bit of randomness.
    """

    def __init__(
        self,
        batch_size: int,
        dataset: Optional[Dataset] = None,
        lengths: Optional[List[int]] = None,
        model_input_name: Optional[str] = None,
        generator=None,
    ):
        if dataset is None and lengths is None:
            raise ValueError("One of dataset and lengths must be provided.")

        self.batch_size = batch_size
        if lengths is None:
            model_input_name = model_input_name if model_input_name is not None else "input_ids"
            if (
                not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding))
                or model_input_name not in dataset[0]
            ):
                raise ValueError(
                    "Can only automatically infer lengths for datasets whose items are dictionaries with an "
                    f"'{model_input_name}' key."
                )
            lengths = [len(feature[model_input_name]) for feature in dataset]
        elif isinstance(lengths, torch.Tensor):
            logger.info(
                "If lengths is a torch.Tensor, LengthGroupedSampler will be slow. Converting lengths to List[int]..."
            )
            lengths = lengths.tolist()

        self.lengths = lengths
        self.generator = generator

    def __len__(self):
        return len(self.lengths)

    def __iter__(self):
        indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator)


SUMMARY:
**LengthGroupedSampler Class Summary**

**Main Purpose and Functionality:**
The `LengthGroupedSampler` class is a sampler that groups indices from a dataset based on feature lengths while maintaining randomization.

**Key Implementation Details:**

* Employs the `get_length_grouped_indices` function to generate indices.
* Infers feature lengths automatically for datasets with dictionary-like structures and specified input name.
* Raises errors if either `dataset` or `lengths` are missing.
* Includes an optional generator parameter.

**Notable Dependencies and Requirements:**

* `torch` library (version 1.9+)
* `get_length_grouped_indices` function (not provided in the chunk)

**Method Signatures and Interfaces:**

* `__init__(batch_size, dataset, lengths, model_input_name, generator)`
* `__len__()`
* `__iter__()`

Note that the `get_length_grouped_indices` function is not included in this summary as its implementation details are not provided.


key: chroma:document
string_value:
CODE:

class DistributedLengthGroupedSampler(DistributedSampler):
    r"""
    Distributed Sampler that samples indices in a way that groups together features of the dataset of roughly the same
    length while keeping a bit of randomness.
    """

    # Copied and adapted from PyTorch DistributedSampler.
    def __init__(
        self,
        batch_size: int,
        dataset: Optional[Dataset] = None,
        num_replicas: Optional[int] = None,
        rank: Optional[int] = None,
        seed: int = 0,
        drop_last: bool = False,
        lengths: Optional[List[int]] = None,
        model_input_name: Optional[str] = None,
    ):
        if dataset is None and lengths is None:
            raise ValueError("One of dataset and lengths must be provided.")
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()

        self.batch_size = batch_size
        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last

        if lengths is None:
            model_input_name = model_input_name if model_input_name is not None else "input_ids"
            if (
                not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding))
                or model_input_name not in dataset[0]
            ):
                raise ValueError(
                    "Can only automatically infer lengths for datasets whose items are dictionaries with an "
                    f"'{model_input_name}' key."
                )
            lengths = [len(feature[model_input_name]) for feature in dataset]
        elif isinstance(lengths, torch.Tensor):
            logger.info(
                "If lengths is a torch.Tensor, DistributedLengthGroupedSampler will be slow. Converting lengths to"
                " List[int]..."
            )
            lengths = lengths.tolist()

        self.lengths = lengths

        # If the dataset length is evenly divisible by # of replicas, then there
        # is no need to drop any data, since the dataset will be split equally.
        if self.drop_last and len(self.lengths) % self.num_replicas != 0:
            # Split to nearest available length that is evenly divisible.
            # This is to ensure each rank receives the same amount of data when
            # using this Sampler.
            self.num_samples = math.ceil((len(self.lengths) - self.num_replicas) / self.num_replicas)
        else:
            self.num_samples = math.ceil(len(self.lengths) / self.num_replicas)
        self.total_size = self.num_samples * self.num_replicas
        self.seed = seed

    def __iter__(self) -> Iterator:
        # Deterministically shuffle based on epoch and seed
        g = torch.Generator()
        g.manual_seed(self.seed + self.epoch)
        indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g)

        if not self.drop_last:
            # add extra samples to make it evenly divisible
            indices += indices[: (self.total_size - len(indices))]
        else:
            # remove tail of data to make it evenly divisible.
            indices = indices[: self.total_size]
        assert len(indices) == self.total_size

        # subsample
        indices = indices[self.rank : self.total_size : self.num_replicas]
        assert len(indices) == self.num_samples



SUMMARY:
**Distributed Length Grouped Sampler Class Summary**

**Main Purpose and Functionality:**
This class is a customized `DistributedSampler` that groups together features of the dataset with roughly similar lengths while maintaining some randomness.

**Key Implementation Details:**

* Sampling indices using `get_length_grouped_indices` with epoch, seed, batch size, and data lengths.
* Deterministically shuffles based on epoch and seed.
* Groups samples into batches using a random permutation.
* Adjusts batch sizes to make the total number of samples evenly divisible by the number of replicas.

**Notable Dependencies or Requirements:**

* `torch` (for tensor operations)
* `dist` (for distributed package and replica indices)

**Interfaces and Inherited Dependencies:**

* `Dataset` interface (for automatically inferring data lengths)
* `BatchEncoding` class (for checking dataset item format)


key: chroma:document
string_value:
CODE:

class ShardSampler(Sampler):
    """
    Sampler that shards batches between several processes. Dispatches indices batch by batch: on 2 processes with batch
    size 4, the first two batches are `[0, 1, 2, 3, 4, 5, 6, 7]` and `[8, 9, 10, 11, 12, 13, 14, 15]`, which shard into
    `[0, 1, 2, 3]` and `[8, 9, 10, 11]` for GPU-0 and `[4, 5, 6, 7]` and `[12, 13, 14, 15]` for GPU-1.

    The sampler thus yields `[0, 1, 2, 3, 8, 9, 10, 11]` on GPU-0 and `[4, 5, 6, 7, 12, 13, 14, 15]` on GPU-1.
    """

    def __init__(
        self,
        dataset: Dataset,
        batch_size: int = 1,
        drop_last: bool = False,
        num_processes: int = 1,
        process_index: int = 0,
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.num_processes = num_processes
        self.process_index = process_index

        self.total_batch_size = total_batch_size = batch_size * num_processes

        num_batches = len(dataset) // total_batch_size if drop_last else math.ceil(len(dataset) / total_batch_size)
        self.total_num_samples = num_batches * total_batch_size

    def __iter__(self):
        indices = list(range(len(self.dataset)))

        # Add extra samples to make it evenly divisible. While loop is there in the edge case we have a tiny dataset
        # and it needs to be done several times.
        while len(indices) < self.total_num_samples:
            indices += indices[: (self.total_num_samples - len(indices))]

        result = []
        for batch_start in range(self.batch_size * self.process_index, self.total_num_samples, self.total_batch_size):
            result += indices[batch_start : batch_start + self.batch_size]

        return iter(result)

    def __len__(self):
        # Each shard only sees a fraction of total_num_samples.


SUMMARY:
**ShardSampler Class Summary**

*   **Purpose**: Shard batches between several processes, dispatching indices in parallel to multiple GPUs.

### Key Implementation Details

*   Shards dataset into `num_processes` number of queues and assigns them to `num_processes` individual processes (GPU-0 to GPU-`process_index`)
*   Dispatches batch by batch; first `batch_size * process_index` batches go to the current process.
*   Calculates the total number of samples per queue based on the dataset size and batch size.
*   Iterates over each batch, assigning indices from the main list, repeated as needed to fill the queue.

### Notable Dependencies/Requirements

*   `Sampler`: Base class for sharding algorithm
*   `Dataset`: Provides data for sharding; inherits from this class to access dataset properties (e.g., `__len__`)
*   External dependencies:
    *   `math` (for ceil function)
    *   Python's built-in `range` and list concatenation functions (`+`).


key: chroma:document
string_value:
CODE:

class IterableDatasetShard(IterableDataset):
    """
    Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will
    always yield a number of samples that is a round multiple of the actual batch size (which is `batch_size x
    num_processes`). Depending on the value of the `drop_last` attribute, it will either stop the iteration at the
    first batch that would be too small or loop with indices from the beginning.

    On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]` with a batch size of
    2:

    - the shard on process 0 will yield `[0, 1, 4, 5, 8, 9]` so will see batches `[0, 1]`, `[4, 5]`, `[8, 9]`
    - the shard on process 1 will yield `[2, 3, 6, 7, 10, 11]` so will see batches `[2, 3]`, `[6, 7]`, `[10, 11]`

    <Tip warning={true}>

        If your IterableDataset implements some randomization that needs to be applied the same way on all processes
        (for instance, a shuffling), you should use a `torch.Generator` in a `generator` attribute of the `dataset` to
        generate your random numbers and call the [`~trainer_pt_utils.IterableDatasetShard.set_epoch`] method of this
        object. It will set the seed of this `generator` to `seed + epoch` on all processes before starting the
        iteration. Alternatively, you can also implement a `set_epoch()` method in your iterable dataset to deal with
        this.

    </Tip>

    Args:
        dataset (`torch.utils.data.IterableDataset`):
            The batch sampler to split in several shards.
        batch_size (`int`, *optional*, defaults to 1):
            The size of the batches per shard.
        drop_last (`bool`, *optional*, defaults to `False`):
            Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the
            beginning.
        num_processes (`int`, *optional*, defaults to 1):
            The number of processes running concurrently.
        process_index (`int`, *optional*, defaults to 0):
            The index of the current process.
        seed (`int`, *optional*, defaults to 0):
            A random seed that will be used for the random number generation in
            [`~trainer_pt_utils.IterableDatasetShard.set_epoch`].
    """

    def __init__(
        self,
        dataset: IterableDataset,
        batch_size: int = 1,
        drop_last: bool = False,
        num_processes: int = 1,
        process_index: int = 0,
        seed: int = 0,
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.drop_last = drop_last
        self.num_processes = num_processes
        self.process_index = process_index
        self.seed = seed
        self.epoch = 0
        self.num_examples = 0

    def set_epoch(self, epoch):
        self.epoch = epoch
        if hasattr(self.dataset, "set_epoch"):
            self.dataset.set_epoch(epoch)

    def __iter__(self):
        self.num_examples = 0
        if (
            not hasattr(self.dataset, "set_epoch")
            and hasattr(self.dataset, "generator")
            and isinstance(self.dataset.generator, torch.Generator)
        ):
            self.dataset.generator.manual_seed(self.seed + self.epoch)
        real_batch_size = self.batch_size * self.num_processes
        process_slice = range(self.process_index * self.batch_size, (self.process_index + 1) * self.batch_size)

        first_batch = None
        current_batch = []
        for element in self.dataset:
            self.num_examples += 1
            current_batch.append(element)
            # Wait to have a full batch before yielding elements.
            if len(current_batch) == real_batch_size:
                for i in process_slice:
                    yield current_batch[i]
                if first_batch is None:
                    first_batch = current_batch.copy()
                current_batch = []

        # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning.
        if not self.drop_last and len(current_batch) > 0:
            if first_batch is None:
                first_batch = current_batch.copy()
            while len(current_batch) < real_batch_size:
                current_batch += first_batch
            for i in process_slice:
                yield current_batch[i]

    def __len__(self):
        # Will raise an error if the underlying dataset is not sized.
        if self.drop_last:
            return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size
        else:


SUMMARY:
**Summary**

`IterableDatasetShard` class: Wraps an `IterableDataset` to yield samples for a specific process. 

*   Purpose:
    *   Split an iterable dataset into shards based on batch size and processes.
    *   Yield batches in a round multiple of the actual batch size, depending on the `drop_last` attribute.
*   Key Implementation Details:
    *   Uses PyTorch `torch.Generator` for randomization if available.
    *   Sets seed for generator based on process index and epoch number.
    *   Yields samples from the dataset in chunks of `batch_size x num_processes`.
*   Notable Dependencies:
    *   `torch.utils.data.IterableDataset`
    *   `torch.Generator`
*   Interfaces:
    *   Implements PyTorch interface for `IterableDataset`

**Methods**

*   `__iter__()`: Iterates over the dataset, yielding samples in chunks of `batch_size x num_processes`.
*   `set_epoch(epoch)`: Sets the epoch number and updates the seed for randomization if applicable.
*   `__len__()`: Returns the size of the dataset based on the batch size and processes.

**Attributes**

*   `dataset`: The underlying iterable dataset.
*   `batch_size`: The size of batches per shard.
*   `drop_last`: Whether to drop the last incomplete batch or complete it with samples from the beginning.
*   `num_processes`: The number of processes running concurrently.
*   `process_index`: The index of the current process.
*   `seed`: A random seed used for randomization.
*   `epoch`: The epoch number.
*   `num_examples`: The number of examples yielded so far.


key: docstring
string_value:
# In order to keep `trainer.py` compact and easy to understand, place any secondary PT Trainer
# helper methods here


key: chroma:document
string_value:
CODE:

def _get_learning_rate(self):
    if self.is_deepspeed_enabled:
        # with deepspeed's fp16 and dynamic loss scale enabled the optimizer/scheduler steps may
        # not run for the first few dozen steps while loss scale is too large, and thus during
        # that time `get_last_lr` will fail if called during that warm up stage, so work around it:
        try:
            last_lr = self.lr_scheduler.get_last_lr()[0]
        except AssertionError as e:
            if "need to call step" in str(e):
                logger.warning("tried to get lr value before scheduler/optimizer started stepping, returning lr=0")
                last_lr = 0
            else:
                raise
    else:
        if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
            last_lr = self.optimizer.param_groups[0]["lr"]
        else:
            last_lr = self.lr_scheduler.get_last_lr()[0]
        if torch.is_tensor(last_lr):
            last_lr = last_lr.item()


SUMMARY:
**Function Summary: `_get_learning_rate`**

### Purpose and Functionality

This function retrieves the last learning rate value from an instance of the class. If DeepSpeed is enabled, it calculates the adjusted learning rate to avoid potential errors during warm-up stages. Otherwise, it extracts the latest learning rate from either the optimizer or the lr_scheduler.

**Key Implementation Details:**

* Uses `try`-`except` block to handle potential issues with `get_last_lr` when DeepSpeed is enabled.
* Checks if `lr_scheduler` is a `ReduceLROnPlateau` instance and extracts the current learning rate accordingly.
* Converts `last_lr` to a Python scalar value if it's a tensor.

**Notable Dependencies:**

* `torch.optim.lr_scheduler.ReduceLROnPlateau`
* External dependencies not specified in this function snippet.


key: chroma:document
string_value:
CODE:

def _secs2timedelta(secs):
    """
    convert seconds to hh:mm:ss.msec, msecs rounded to 2 decimals
    """

    msec = int(abs(secs - int(secs)) * 100)


SUMMARY:
**Function Summary**

### `_secs2timedelta`

Converts seconds to `hh:mm:ss.msec` format.

#### Key Implementation Details

* Uses `datetime.timedelta` to convert seconds to timedelta object
* Rounds milliseconds to 2 decimal places using integer division and string formatting

#### Notable Dependencies/Requirements

* `datetime` module (specifically, `timedelta` class)


key: chroma:document
string_value:
CODE:

def metrics_format(self, metrics: Dict[str, float]) -> Dict[str, float]:
    """
    Reformat Trainer metrics values to a human-readable format

    Args:
        metrics (`Dict[str, float]`):
            The metrics returned from train/evaluate/predict

    Returns:
        metrics (`Dict[str, float]`): The reformatted metrics
    """

    metrics_copy = metrics.copy()
    for k, v in metrics_copy.items():
        if "_mem_" in k:
            metrics_copy[k] = f"{ v >> 20 }MB"
        elif "_runtime" in k:
            metrics_copy[k] = _secs2timedelta(v)
        elif k == "total_flos":
            metrics_copy[k] = f"{ int(v) >> 30 }GF"
        elif isinstance(metrics_copy[k], float):
            metrics_copy[k] = round(v, 4)



SUMMARY:
**Summary:**

Reformats Trainer metrics values to a human-readable format.

```python
def metrics_format(self, metrics: Dict[str, float]) -> Dict[str, float]:
    """
    Reformat Trainer metrics values to a human-readable format
    
    Args:
        metrics (Dict[str, float]):
            Metrics returned from train/evaluate/predict
    
    Returns:
        metrics (Dict[str, float]): Reformatted metrics
    """

    reformatted_metrics = metrics.copy()
    
    # Convert memory usage and total FLOPs to human-readable format
    for key in ["_mem_", "_runtime", "total_flops"]:
        if key in metrics_copy:
            value = metrics_copy.pop(key)
            if key == "_mem_":
                reformatted_metrics[f"{key}{value}"] = f"{ value >> 20 }MB"
            elif key == "_runtime":
                reformatted_metrics[k] = _secs2timedelta(value)
            elif key == "total_flops":
                reformatted_metrics[k] = f"{ int(value) >> 30 }GF"
    
    # Round all other values to 4 decimal places
    for key, value in metrics_copy.items():
        if not isinstance(value, (int, float)):
            continue
        reformatted_metrics[key] = round(value, 4)
        
    return metrics_copy
```

**Key Implementation Details:**

- Uses bitwise right shift (`>>`) to reformat memory usage and FLOPs values.
- Applies `_secs2timedelta` function to format runtime value.
- Rounds all other float values to 4 decimal places.

**Notable Dependencies or Requirements:**

- `_secs2timedelta` function (not shown in this snippet).
- `float`, `int`, and `dict` types.


key: chroma:document
string_value:
CODE:

def log_metrics(self, split, metrics):
    """
    Log metrics in a specially formatted way

    Under distributed environment this is done only for a process with rank 0.

    Args:
        split (`str`):
            Mode/split name: one of `train`, `eval`, `test`
        metrics (`Dict[str, float]`):
            The metrics returned from train/evaluate/predictmetrics: metrics dict

    Notes on memory reports:

    In order to get memory usage report you need to install `psutil`. You can do that with `pip install psutil`.

    Now when this method is run, you will see a report that will include: :

    ```
    init_mem_cpu_alloc_delta   =     1301MB
    init_mem_cpu_peaked_delta  =      154MB
    init_mem_gpu_alloc_delta   =      230MB
    init_mem_gpu_peaked_delta  =        0MB
    train_mem_cpu_alloc_delta  =     1345MB
    train_mem_cpu_peaked_delta =        0MB
    train_mem_gpu_alloc_delta  =      693MB
    train_mem_gpu_peaked_delta =        7MB
    ```

    **Understanding the reports:**

    - the first segment, e.g., `train__`, tells you which stage the metrics are for. Reports starting with `init_`
        will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the
        `__init__` will be reported along with the `eval_` metrics.
    - the third segment, is either `cpu` or `gpu`, tells you whether it's the general RAM or the gpu0 memory
        metric.
    - `*_alloc_delta` - is the difference in the used/allocated memory counter between the end and the start of the
        stage - it can be negative if a function released more memory than it allocated.
    - `*_peaked_delta` - is any extra memory that was consumed and then freed - relative to the current allocated
        memory counter - it is never negative. When you look at the metrics of any stage you add up `alloc_delta` +
        `peaked_delta` and you know how much memory was needed to complete that stage.

    The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the
    main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may
    use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more
    memory than the rest since it stores the gradient and optimizer states for all participating GPUS. Perhaps in the
    future these reports will evolve to measure those too.

    The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the
    memory shared with other processes. It is important to note that it does not include swapped out memory, so the
    reports could be imprecise.

    The CPU peak memory is measured using a sampling thread. Due to python's GIL it may miss some of the peak memory if
    that thread didn't get a chance to run when the highest memory was used. Therefore this report can be less than
    reality. Using `tracemalloc` would have reported the exact peak memory, but it doesn't report memory allocations
    outside of python. So if some C++ CUDA extension allocated its own memory it won't be reported. And therefore it
    was dropped in favor of the memory sampling approach, which reads the current process memory usage.

    The GPU allocated and peak memory reporting is done with `torch.cuda.memory_allocated()` and
    `torch.cuda.max_memory_allocated()`. This metric reports only "deltas" for pytorch-specific allocations, as
    `torch.cuda` memory management system doesn't track any memory allocated outside of pytorch. For example, the very
    first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory.

    Note that this tracker doesn't account for memory allocations outside of [`Trainer`]'s `__init__`, `train`,
    `evaluate` and `predict` calls.

    Because `evaluation` calls may happen during `train`, we can't handle nested invocations because
    `torch.cuda.max_memory_allocated` is a single counter, so if it gets reset by a nested eval call, `train`'s tracker
    will report incorrect info. If this [pytorch issue](https://github.com/pytorch/pytorch/issues/16266) gets resolved
    it will be possible to change this class to be re-entrant. Until then we will only track the outer level of
    `train`, `evaluate` and `predict` methods. Which means that if `eval` is called during `train`, it's the latter
    that will account for its memory usage and that of the former.

    This also means that if any other tool that is used along the [`Trainer`] calls
    `torch.cuda.reset_peak_memory_stats`, the gpu peak memory stats could be invalid. And the [`Trainer`] will disrupt
    the normal behavior of any such tools that rely on calling `torch.cuda.reset_peak_memory_stats` themselves.

    For best performance you may want to consider turning the memory profiling off for production runs.
    """
    if not self.is_world_process_zero():
        return

    print(f"***** {split} metrics *****")
    metrics_formatted = self.metrics_format(metrics)
    k_width = max(len(str(x)) for x in metrics_formatted.keys())
    v_width = max(len(str(x)) for x in metrics_formatted.values())
    for key in sorted(metrics_formatted.keys()):


SUMMARY:
**Function Summary: `log_metrics`**

**Purpose and Functionality**
The `log_metrics` function logs metrics in a specially formatted way, providing detailed memory usage reports for CPU and GPU.

**Key Implementation Details**

*   Logs memory reports for each stage (init, train, eval, test) under a distributed environment, where the process rank 0 is considered.
*   Uses `psutil` to calculate memory usage.
*   Calculates memory allocations and releases using Python's GIL (`psutil` or `tracemalloc`) for CPU and PyTorch-specific CUDA calls for GPU.

**Notable Dependencies or Requirements**

*   `torch` library (depending on GPU support)
*   `psutil` library for memory usage tracking
*   `pytorch` issue [16266](https://github.com/pytorch/pytorch/issues/16266) regarding memory allocation

**Interface and Parameters**

*   `split` (`str`) - Mode/split name: one of `train`, `eval`, `test`
*   `metrics` (`Dict[str, float]`) - The metrics returned from train/evaluate/predict, which are formatted by the method.

**Design Choice and Context**
This function is used primarily for process rank 0 under a distributed environment. In case of model parallel or DataParallel use cases, GPU memory may need to be considered separately in these configurations.


key: chroma:document
string_value:
CODE:

def save_metrics(self, split, metrics, combined=True):
    """
    Save metrics into a json file for that split, e.g. `train_results.json`.

    Under distributed environment this is done only for a process with rank 0.

    Args:
        split (`str`):
            Mode/split name: one of `train`, `eval`, `test`, `all`
        metrics (`Dict[str, float]`):
            The metrics returned from train/evaluate/predict
        combined (`bool`, *optional*, defaults to `True`):
            Creates combined metrics by updating `all_results.json` with metrics of this call

    To understand the metrics please read the docstring of [`~Trainer.log_metrics`]. The only difference is that raw
    unformatted numbers are saved in the current method.

    """
    if not self.is_world_process_zero():
        return

    path = os.path.join(self.args.output_dir, f"{split}_results.json")
    with open(path, "w") as f:
        json.dump(metrics, f, indent=4, sort_keys=True)

    if combined:
        path = os.path.join(self.args.output_dir, "all_results.json")
        if os.path.exists(path):
            with open(path, "r") as f:
                all_metrics = json.load(f)
        else:
            all_metrics = {}

        all_metrics.update(metrics)
        with open(path, "w") as f:


SUMMARY:
**Function Summary: `save_metrics`**

*   **Purpose**: Saves metrics to a JSON file for a specified split (`train`, `eval`, `test`, or `all`) and optionally updates the combined metrics in `all_results.json`.

**Key Implementation Details:**

*   The function uses `json.dump` to serialize the `metrics` dictionary to the output directory.
*   In distributed environments, only processes with rank 0 commit to the JSON file.

**Notable Dependencies and Requirements:**

*   `os`: used for interacting with the operating system (output directory construction).
*   `json`: used for serializing/deserializing the metrics data.

**Method Signatures and Interfaces:**

*   The function does not inherit from any interfaces but uses the `Trainer` class's methods (`is_world_process_zero`, `args.output_dir`).


key: chroma:document
string_value:
CODE:

def save_state(self):
    """
    Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model

    Under distributed environment this is done only for a process with rank 0.
    """
    if not self.is_world_process_zero():
        return

    path = os.path.join(self.args.output_dir, "trainer_state.json")


SUMMARY:
**Summary**

The `save_state` function saves the trainer's state to a JSON file when running in a distributed environment. It is only executed on process rank 0, preserving the model's tokenizer when saved.

**Key Implementation Details**

* Saves trainer state to local file at `self.args.output_dir/trainer_state.json`
* Excludes writes for non-rank-0 processes

**Notable Dependencies/Requirements**

* `os.path.join` for constructing output path
* `self.state` (implies a `TrainerState` class or similar) for saving state to JSON
* `Trainer.save_model` (presumably on same model, not included due to context)
* Distributed environment setup utilizing process ranks (e.g., `multiprocessing` module is implied, though not explicitly mentioned)


key: chroma:document
string_value:
CODE:

def get_model_param_count(model, trainable_only=False):
    """
    Calculate model's total param count. If trainable_only is True then count only those requiring grads
    """
    if is_deepspeed_zero3_enabled():

        def numel(p):
            return p.ds_numel if hasattr(p, "ds_numel") else p.numel()

    else:

        def numel(p):
            return p.numel()



SUMMARY:
**Function Summary:**

`get_model_param_count(model, trainable_only=False)` calculates the total parameter count of the input `model`. If `trainable_only` is `True`, only parameters requiring gradients are counted.

**Key Implementation Details:**

- Uses DeepSpeed 3 (DS3) enabled models by default (`numel(p) = p.ds_numel if hasattr(p, "ds_numel") else p.numel()`).
- Checks for trainable parameters with `models.requires_grad`.

**Notable Dependencies/Requirements:**

- `model` parameter is assumed to be an instance of PyTorch model.
- DeepSpeed 3 (DS3) related classes/functions (`is_deepspeed_zero3_enabled()`, `ds_numel`) may be external dependencies.


key: chroma:document
string_value:
CODE:

def get_parameter_names(model, forbidden_layer_types):
    """
    Returns the names of the model parameters that are not inside a forbidden layer.
    """
    result = []
    for name, child in model.named_children():
        result += [
            f"{name}.{n}"
            for n in get_parameter_names(child, forbidden_layer_types)
            if not isinstance(child, tuple(forbidden_layer_types))
        ]
    # Add model specific parameters (defined with nn.Parameter) since they are not in any child.
    result += list(model._parameters.keys())


SUMMARY:
**Summary: `get_parameter_names` Function**

* **Purpose:** Returns the names of model parameters outside forbidden layer types.

**Key Implementation Details:**

* Recursively traverses the model's child nodes using `named_children()` method.
* Filters out parameters within specified forbidden layer types (`isinstance(child, tuple(forbidden_layer_types))`).
* Includes global model parameters not contained in any subtree by looking at `_parameters.keys()`.

**Notable Dependencies/Requirements:**

* Uses PyTorch's `nn.Parameter` and its related metadata (`_parameters`) internally.
* Does not utilize external dependencies.


key: chroma:document
string_value:
CODE:

def get_module_class_from_name(module, name):
    """
    Gets a class from a module by its name.

    Args:
        module (`torch.nn.Module`): The module to get the class from.
        name (`str`): The name of the class.
    """
    modules_children = list(module.children())
    if module.__class__.__name__ == name:
        return module.__class__
    elif len(modules_children) == 0:
        return
    else:
        for child_module in modules_children:
            module_class = get_module_class_from_name(child_module, name)
            if module_class is not None:


SUMMARY:
**Function Summary: `get_module_class_from_name`**

Gets a class from a PyTorch neural network module by its name.

* Primary purpose: Recursively search for a specific class within the given module and its children.
* Key implementation details:
	+ Iterates through the module's children using `module.children()`.
	+ Checks if the current module matches the target name, returning itself if so.
	+ Recursively calls itself on each child module until the target class is found or no more children exist.
* Notable dependencies and requirements: None
* Interface:
```
get_module_class_from_name(module: torch.nn.Module, name: str)
```

**Key Takeaways:**

* This function assumes that the input `module` is a valid PyTorch neural network module.
* The function returns the class of the matched module or `None` if no match is found.
* Recursive approach used to traverse the module's hierarchy.


key: chroma:document
string_value:
CODE:

def remove_dummy_checkpoint(is_main_process, output_dir, filenames):
    if is_main_process:
        for filename in filenames:
            file = os.path.join(output_dir, filename)
            if os.path.isfile(file):


SUMMARY:
**Concise Summary**

`remove_dummy_checkpoint` function removes dummy checkpoint files from an output directory when running as the main process.

**Detailed Analysis**

* **Purpose**: Removes unwanted files based on user-provided `filenames`.
* **Key Implementation Details**:
	+ Iterates over a list of file names using `os.path.join`.
	+ Checks if each generated file exists using `os.path.isfile` and removes it using `os.remove` when in main process.
* **Notable Dependencies or Requirements**: None.

**Summary Classifications**

|  | Main Purpose | Key Implementation Details |
| --- | --- | --- |
| 1 | Removes dummy checkpoint files | Iterates over file names, checks existence, and deletes them. |
| 2 | User-provided input required | Uses `os.path.join` to construct file paths. |


key: chroma:document
string_value:
CODE:

class LayerWiseDummyOptimizer(torch.optim.Optimizer):
    """
    For Layer-wise optimizers such as GaLoRE optimizer, the optimization
    step is already done through the post gradient hooks. Therefore
    the trick is to create a dummy optimizer that can take arbitrary
    args and kwargs and return a no-op during training.

    Initial idea from @hiyouga in LLaMA-Factory:
    https://github.com/hiyouga/LLaMA-Factory/commit/8664262cde3919e10eaecbd66e8c5d356856362e#diff-ebe08ab14496dfb9e06075f0fdd36799ef6d1535cc4dd4715b74c4e3e06fe3ba
    """

    def __init__(self, optimizer_dict=None, *args, **kwargs):
        dummy_tensor = torch.randn(1, 1)
        self.optimizer_dict = optimizer_dict
        super().__init__([dummy_tensor], {"lr": kwargs.get("lr", 1e-03)})

    def zero_grad(self, set_to_none: bool = True) -> None:
        pass

    def step(self, closure=None) -> Optional[float]:


SUMMARY:
**LayerWiseDummyOptimizer Class Summary**

The `LayerWiseDummyOptimizer` class is a custom optimization wrapper designed for Layer-wise optimizers like GaLoRE optimizer. Its primary purpose is to create a dummy optimizer that takes arbitrary arguments and returns no-op during training.

**Key Implementation Details:**

* Initializes with an optional `optimizer_dict`, a random tensor (`torch.randn(1, 1)`) and sets the learning rate using `kwargs.get("lr", 1e-03)`.
* Provides dummy implementations for `zero_grad` and `step` methods.

**Notable Dependencies:**

* Uses PyTorch library (`torch.optim.Optimizer`)
* References the `torch.randn` function
* Note that it inherits from `torch.optim.Optimizer`, which implements standard optimization algorithms, but uses a dummy implementation to achieve its specific purpose.


key: chroma:document
string_value:
CODE:

class LayerWiseDummyScheduler(LRScheduler):
    """
    For Layer-wise optimizers such as GaLoRE optimizer, the optimization and scheduling step
    are already done through the post gradient hooks. Therefore
    the trick is to create a dummy scheduler that can take arbitrary
    args and kwargs and return a no-op during training.
    """

    def __init__(self, *args, **kwargs):
        self.default_lr = kwargs["lr"]
        optimizer = LayerWiseDummyOptimizer(**kwargs)
        last_epoch = -1
        verbose = False
        super().__init__(optimizer, last_epoch, verbose)

    def get_lr(self):
        # default value
        lrs = [self.default_lr]

        # we take each lr in the parameters if they exist, assumes the optimizer to be the `LayerWiseDummyOptimizer`
        if self.optimizer is not None:
            param_wise_lrs = [
                [group["lr"] for group in optim.param_groups] for optim in self.optimizer.optimizer_dict.values()
            ]
            lrs = list(chain(*param_wise_lrs))

        return lrs

    def _get_closed_form_lr(self):


SUMMARY:
**LayerWiseDummyScheduler Class Summary**

*   **Purpose and Functionality:** The LayerWiseDummyScheduler is a mock scheduler that implements the LRScheduler interface. Its primary function is to return no-op operations during training. It acts as a dummy scheduler for optimizers like GaLoRE.
*   **Key Implementation Details:**
    *   `get_lr` method:
        *   Returns the initial learning rate (`self.default_lr`)
        *   If `self.optimizer` exists, it retrieves parameter-wise learning rates from the optimizer
    *   `_get_closed_form_lr` method (not implemented):
        *   Not available (returns `None` or an error message)

**Notable Dependencies:**

*   `LayerWiseDummyOptimizer`: The specific optimizer instance that extends this scheduler.

**External Dependencies:**

*   None

**Inherited/Implemented Interfaces:**

*   LRScheduler interface


key: chroma:document
string_value:
CODE:

class TransposeType(ExplicitEnum):
    """
    Possible ...
    """

    NO = "no"
    SIMPLE = "simple"
    CONV1D = "conv1d"


SUMMARY:
**Summary:**

The `TransposeType` class defines an enumeration with four possible values:

* `NO`: No transpose operation
* `SIMPLE`: Simple transpose (likely referring to spatial or batch dimension swapping)
* `CONV1D`: Transpose for 1D convolutional layers
* `CONV2D`: Transpose for 2D convolutional layers

The primary purpose of this enumeration is to provide a way to specify the type of transpose operation when working with neural networks, likely for model architecture or layer configurations.

**Key Implementation Details:**

* No specific algorithmic implementation details are present in this code snippet.
* The emphasis is on providing a clear and concise representation of possible transpose types.

**Notable Dependencies or Requirements:**

* None explicitly mentioned; however, the context suggests an environment with `ExplicitEnum` support, likely in a .NET or C#-based framework.


key: chroma:document
string_value:
CODE:

def convert_tf_weight_name_to_pt_weight_name(
    tf_name, start_prefix_to_remove="", tf_weight_shape=None, name_scope=None
):
    """
    Convert a TF 2.0 model variable name in a pytorch model weight name.

    Conventions for TF2.0 scopes -> PyTorch attribute names conversions:

        - '$1___$2' is replaced by $2 (can be used to duplicate or remove layers in TF2.0 vs PyTorch)
        - '_._' is replaced by a new level separation (can be used to convert TF2.0 lists in PyTorch nn.ModulesList)

    return tuple with:

        - pytorch model weight name
        - transpose: `TransposeType` member indicating whether and how TF2.0 and PyTorch weights matrices should be
          transposed with regards to each other
    """
    if name_scope is not None:
        if not tf_name.startswith(name_scope) and "final_logits_bias" not in tf_name:
            raise ValueError(
                f"Weight name {tf_name} does not start with name_scope {name_scope}. This is an internal error "
                "in Transformers, so (unless you were doing something really evil) please open an issue to report it!"
            )
        tf_name = tf_name[len(name_scope) :]
        tf_name = tf_name.lstrip("/")
    tf_name = tf_name.replace(":0", "")  # device ids
    tf_name = re.sub(
        r"/[^/]*___([^/]*)/", r"/\1/", tf_name
    )  # '$1___$2' is replaced by $2 (can be used to duplicate or remove layers in TF2.0 vs PyTorch)
    tf_name = tf_name.replace(
        "_._", "/"
    )  # '_._' is replaced by a level separation (can be used to convert TF2.0 lists in PyTorch nn.ModulesList)
    tf_name = re.sub(r"//+", "/", tf_name)  # Remove empty levels at the end
    tf_name = tf_name.split("/")  # Convert from TF2.0 '/' separators to PyTorch '.' separators
    # Some weights have a single name without "/" such as final_logits_bias in BART
    if len(tf_name) > 1:
        tf_name = tf_name[1:]  # Remove level zero

    tf_weight_shape = list(tf_weight_shape)

    # When should we transpose the weights
    if tf_name[-1] == "kernel" and tf_weight_shape is not None and len(tf_weight_shape) == 4:
        transpose = TransposeType.CONV2D
    elif tf_name[-1] == "kernel" and tf_weight_shape is not None and len(tf_weight_shape) == 3:
        transpose = TransposeType.CONV1D
    elif bool(
        tf_name[-1] in ["kernel", "pointwise_kernel", "depthwise_kernel"]
        or "emb_projs" in tf_name
        or "out_projs" in tf_name
    ):
        transpose = TransposeType.SIMPLE
    else:
        transpose = TransposeType.NO

    # Convert standard TF2.0 names in PyTorch names
    if tf_name[-1] == "kernel" or tf_name[-1] == "embeddings" or tf_name[-1] == "gamma":
        tf_name[-1] = "weight"
    if tf_name[-1] == "beta":
        tf_name[-1] = "bias"

    # The SeparableConv1D TF layer contains two weights that are translated to PyTorch Conv1D here
    if tf_name[-1] == "pointwise_kernel" or tf_name[-1] == "depthwise_kernel":
        tf_name[-1] = tf_name[-1].replace("_kernel", ".weight")

    # Remove prefix if needed
    tf_name = ".".join(tf_name)
    if start_prefix_to_remove:
        tf_name = tf_name.replace(start_prefix_to_remove, "", 1)



SUMMARY:
**Function Summary: `convert_tf_weight_name_to_pt_weight_name`**

This function converts a TF 2.0 model variable name to a PyTorch model weight name.

**Main Purpose and Functionality:**
The function takes a TF name, an optional parameter to remove from the start of the name, and optional parameters for weight shape and scope.
It returns a tuple containing:
1. The converted PyTorch model weight name
2. A `TransposeType` member indicating whether weights should be transposed between TF 2.0 and PyTorch

**Key Implementation Details:**

* Replaces special names like '$1___$2' with $2
* Converts '/' separators to '.'
* Transposes kernel weights if necessary, based on weight shape
* Standardizes weights as 'weight', biases as 'bias'

**Notable Dependencies or Requirements:**

* `re` module for regular expression replacements
* `TransposeType` member for transpose operation

* `name_scope` parameter assumes an internal error handling mechanism


key: chroma:document
string_value:
CODE:

def apply_transpose(transpose: TransposeType, weight, match_shape=None, pt_to_tf=True):
    """
    Apply a transpose to some weight then tries to reshape the weight to the same shape as a given shape, all in a
    framework agnostic way.
    """
    if transpose is TransposeType.CONV2D:
        # Conv2D weight:
        #    PT: (num_out_channel, num_in_channel, kernel[0], kernel[1])
        # -> TF: (kernel[0], kernel[1], num_in_channel, num_out_channel)
        axes = (2, 3, 1, 0) if pt_to_tf else (3, 2, 0, 1)
        weight = transpose_func(weight, axes=axes)
    elif transpose is TransposeType.CONV1D:
        # Conv1D weight:
        #    PT: (num_out_channel, num_in_channel, kernel)
        # -> TF: (kernel, num_in_channel, num_out_channel)
        weight = transpose_func(weight, axes=(2, 1, 0))
    elif transpose is TransposeType.SIMPLE:
        weight = transpose_func(weight)

    if match_shape is None:
        return weight

    if len(match_shape) < len(weight.shape):
        weight = squeeze(weight)
    elif len(match_shape) > len(weight.shape):
        weight = expand_dims(weight, axis=0)

    if list(match_shape) != list(weight.shape):
        try:
            weight = reshape(weight, match_shape)
        except AssertionError as e:
            e.args += (match_shape, match_shape)
            raise e



SUMMARY:
**Function Summary: `apply_transpose`**

**Purpose:** Apply transpose to a weight tensor and optionally reshape it to conform to a given target shape.

**Functionality:**

1. Applies transpose to the input weight based on the specified `transpose` type (`CONV2D`, `CONV1D`, or `SIMPLE`).
2. If a match shape is provided, resizes the weight tensor using:
	* `squeeze` if the source and target shapes have fewer dimensions.
	* `expand_dims` if the source and target shapes have more dimensions.
3. Attempts to reshape the weighted tensor to conform to the target shape using `reshape`. Fails with an assertion error if the shapes do not match.

**Notable Details:**

* Utilizes `transpose_func` for transpose operations (Conv2D) and a generic version `(weight, axes=axes)` for simple transposes.
* Applies transpose type-specific transformations:
	+ Conv2D: flips dimensions (PT -> TF).
	+ Conv1D: rearranges dimensions (PT -> TF).
* Supports match-shaped resizing using `squeeze` and `expand_dims`.
* Raises an assertion error if the target shape does not conform to the source tensor.

**Dependencies/Requirements:**

* `TransposeType`: enumeration of transpose types (`CONV2D`, `CONV1D`, or `SIMPLE`)
* `transpose_func`: a function implementing transpose operations
* `squeeze` and `expand_dims` functions for resizing tensors
* `reshape` function with support for conditional exception handling

**External Interfaces:**

* `TransposeType`
* `transpose_func`
* `squeeze`
* `expand_dims`
* `reshape`


key: chroma:document
string_value:
CODE:

def load_pytorch_checkpoint_in_tf2_model(
    tf_model,
    pytorch_checkpoint_path,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
):
    """Load pytorch checkpoints in a TF 2.0 model"""
    try:
        import tensorflow as tf  # noqa: F401
        import torch  # noqa: F401
        from safetensors.torch import load_file as safe_load_file  # noqa: F401
    except ImportError:
        logger.error(
            "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    # Treats a single file as a collection of shards with 1 shard.
    if isinstance(pytorch_checkpoint_path, str):
        pytorch_checkpoint_path = [pytorch_checkpoint_path]

    # Loads all shards into a single state dictionary
    pt_state_dict = {}
    for path in pytorch_checkpoint_path:
        pt_path = os.path.abspath(path)
        logger.info(f"Loading PyTorch weights from {pt_path}")
        if pt_path.endswith(".safetensors"):
            state_dict = safe_load_file(pt_path)
        else:
            weights_only_kwarg = {"weights_only": True}
            state_dict = torch.load(pt_path, map_location="cpu", **weights_only_kwarg)

        pt_state_dict.update(state_dict)

    logger.info(f"PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters")

    return load_pytorch_weights_in_tf2_model(
        tf_model,
        pt_state_dict,
        tf_inputs=tf_inputs,
        allow_missing_keys=allow_missing_keys,
        output_loading_info=output_loading_info,
        _prefix=_prefix,
        tf_to_pt_weight_rename=tf_to_pt_weight_rename,


SUMMARY:
**Summary**

loads PyTorch checkpoints into a TensorFlow 2.0 model, handling state dictionary merging and weight loading.

**Main Purpose and Functionality**
```python
def load_pytorch_checkpoint_in_tf2_model(
    tf_model,
    pytorch_checkpoint_path,
    ...
):
```
This function loads a PyTorch checkpoint into a TF 2.0 model, treating a single file as a collection of shards with one shard. It handles loading weights from individual checkpoint files using `safe_load_file` and `torch.load`, and merges them into a unified state dictionary.

**Key Implementation Details**

* Loads `pytorch_checkpoint_path` using `os.path.abspath` to get the absolute path.
* Uses `safe_load_file` for `.safetensors` checkpoints and `torch.load` for other checkpoint files, with `weights_only=True` for optimal performance.
* Merges loaded state dictionaries into a single dictionary using `pt_state_dict.update(state_dict)`.

**Dependencies and Requirements**

* TensorFlow
* PyTorch
* Safetensors (library)
* Python logger (`logger.error`, `logger.info`)


key: chroma:document
string_value:
CODE:

def load_pytorch_model_in_tf2_model(tf_model, pt_model, tf_inputs=None, allow_missing_keys=False):
    """Load pytorch checkpoints in a TF 2.0 model"""
    pt_state_dict = pt_model.state_dict()

    return load_pytorch_weights_in_tf2_model(
        tf_model, pt_state_dict, tf_inputs=tf_inputs, allow_missing_keys=allow_missing_keys


SUMMARY:
**Summary:**

`load_pytorch_model_in_tf2_model` Function:

Loads PyTorch models into TF 2.0 models.

### Main Purpose and Functionality:

loads PT_state_dict into TF_model using `load_pytorch_weights_in_tf2_model`.

### Key Implementation Details:

- Loads PyTorch state dictionary (`pt_state_dict`) from `tf_model`.
- Calls `load_pytorch_weights_in_tf2_model` with `tf_model`, `pt_state_dict`, and optional inputs.
  
### Notable Dependencies or Requirements:

- `load_pytorch_weights_in_tf2_model`
- `tf_model` (TF 2.0 model)
- `pt_model` (PyTorch model)
- External dependencies on specific methods and classes implied by the function signature.


key: chroma:document
string_value:
CODE:

def load_pytorch_weights_in_tf2_model(
    tf_model,
    pt_state_dict,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
):
    """Load pytorch state_dict in a TF 2.0 model."""
    try:
        import tensorflow as tf  # noqa: F401
        import torch  # noqa: F401
    except ImportError:
        logger.error(
            "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    # Numpy doesn't understand bfloat16, so upcast to a dtype that doesn't lose precision
    pt_state_dict = {
        k: v.numpy() if v.dtype != torch.bfloat16 else v.float().numpy() for k, v in pt_state_dict.items()
    }
    return load_pytorch_state_dict_in_tf2_model(
        tf_model,
        pt_state_dict,
        tf_inputs=tf_inputs,
        allow_missing_keys=allow_missing_keys,
        output_loading_info=output_loading_info,
        _prefix=_prefix,
        tf_to_pt_weight_rename=tf_to_pt_weight_rename,


SUMMARY:
**Summary**

`load_pytorch_weights_in_tf2_model` loads PyTorch state dictionary into a TensorFlow 2.0 model.

**Key Implementation Details**

* Converts PyTorch weights to compatible tensors using `pt_state_dict.item()`.
* Calls `load_pytorch_state_dict_in_tf2_model` with adapted input parameters.
* Handles potential missing keys in the PyTorch state dictionary with configurable behavior.

**Notable Dependencies or Requirements**

* `tensorflow` (TF 2.0) - loaded using `import tensorflow as tf`
* `torch` - loaded using `import torch`
* `numpy` (implied by tensor upcasting)

**Primary Purpose**

Load and integrate PyTorch state dictionary into a TensorFlow 2.0 model, handling potential data type compatibility issues.


key: chroma:document
string_value:
CODE:

def _log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name):
    if len(unexpected_keys) > 0:
        logger.warning(
            "Some weights of the PyTorch model were not used when initializing the TF 2.0 model"
            f" {class_name}: {unexpected_keys}\n- This IS expected if you are initializing"
            f" {class_name} from a PyTorch model trained on another task or with another architecture"
            " (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n- This IS"
            f" NOT expected if you are initializing {class_name} from a PyTorch model that you expect"
            " to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a"
            " BertForSequenceClassification model)."
        )
    else:
        logger.warning(f"All PyTorch model weights were used when initializing {class_name}.\n")
    if len(missing_keys) > 0:
        logger.warning(
            f"Some weights or buffers of the TF 2.0 model {class_name} were not initialized from the"
            f" PyTorch model and are newly initialized: {missing_keys}\nYou should probably TRAIN this model on a"
            " down-stream task to be able to use it for predictions and inference."
        )
    else:
        logger.warning(
            f"All the weights of {class_name} were initialized from the PyTorch model.\n"
            "If your task is similar to the task the model of the checkpoint was trained on, "
            f"you can already use {class_name} for predictions without further training."
        )

    if len(mismatched_keys) > 0:
        mismatched_warning = "\n".join(
            [
                f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
                for key, shape1, shape2 in mismatched_keys
            ]
        )
        logger.warning(
            f"Some weights of {class_name} were not initialized from the model checkpoint"
            f" are newly initialized because the shapes did not"
            f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able"
            " to use it for predictions and inference."


SUMMARY:
**Function Summary: `_log_key_warnings`**

### Purpose and Functionality

This function logs warnings related to PyTorch model initialization in the TF 2.0 framework. It checks for:

* Unused weights (unexpected keys) when initializing from a PyTorch model trained on another task or architecture.
* Missing weights or buffers not initialized from the PyTorch model.
* Mismatched key shapes between the checkpoint and model instantiated.

### Key Implementation Details

* Logs warnings using the `logger` object with contextual messages for each warning type (unused, missing, mismatched keys).
* Uses list comprehensions to generate formatted strings for display in the warning logs.

### Notable Dependencies and Requirements

* `pytorch`: The PyTorch library, which is assumed to be used as a checkpoint or model source.
* `tf`: The TensorFlow 2.0 framework, which is used to initialize models.
* `logger` object: A logging mechanism for displaying warnings about key mismatches.

### Method Signatures and Interfaces

* `_log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name)`:
	+ Parameters:
	+ `missing_keys`: List of missing weights or buffers.
	+ `unexpected_keys`: List of unused weights.
	+ `mismatched_keys`: List of mismatched key shapes between the checkpoint and model instantiated.
	+ `class_name`: The name of the class that is being initialized.

### Interface Inheritance

This function does not inherit any interfaces, but it uses built-in Python lists to store data.


key: chroma:document
string_value:
CODE:

def load_pytorch_state_dict_in_tf2_model(
    tf_model,
    pt_state_dict,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
    ignore_mismatched_sizes=False,
    skip_logger_warnings=False,
):
    """Load a pytorch state_dict in a TF 2.0 model. pt_state_dict can be either an actual dict or a lazy-loading
    safetensors archive created with the safe_open() function."""
    import tensorflow as tf

    if tf_inputs is None:
        tf_inputs = tf_model.dummy_inputs

    if _prefix is None:
        _prefix = ""
    if tf_inputs:
        with tf.name_scope(_prefix):
            tf_model(tf_inputs, training=False)  # Make sure model is built
    # Convert old format to new format if needed from a PyTorch state_dict
    tf_keys_to_pt_keys = {}
    for key in pt_state_dict.keys():
        new_key = None
        if "gamma" in key:
            new_key = key.replace("gamma", "weight")
        if "beta" in key:
            new_key = key.replace("beta", "bias")
        if "running_var" in key:
            new_key = key.replace("running_var", "moving_variance")
        if "running_mean" in key:
            new_key = key.replace("running_mean", "moving_mean")

        # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
        key_components = key.split(".")
        name = None
        if key_components[-3::2] == ["parametrizations", "original0"]:
            name = key_components[-2] + "_g"
        elif key_components[-3::2] == ["parametrizations", "original1"]:
            name = key_components[-2] + "_v"
        if name is not None:
            key_components = key_components[:-3] + [name]
            new_key = ".".join(key_components)

        if new_key is None:
            new_key = key
        tf_keys_to_pt_keys[new_key] = key

    # Matt: All TF models store the actual model stem in a MainLayer class, including the base model.
    # In PT, the derived models (with heads) use the base model class as the stem instead,
    # and there is no MainLayer class. This means that TF base classes have one
    # extra layer in their weight names, corresponding to the MainLayer class. This code block compensates for that.
    start_prefix_to_remove = ""
    if not any(s.startswith(tf_model.base_model_prefix) for s in tf_keys_to_pt_keys.keys()):
        start_prefix_to_remove = tf_model.base_model_prefix + "."

    symbolic_weights = tf_model.trainable_weights + tf_model.non_trainable_weights
    tf_loaded_numel = 0
    all_pytorch_weights = set(tf_keys_to_pt_keys.keys())
    missing_keys = []
    mismatched_keys = []
    is_safetensor_archive = hasattr(pt_state_dict, "get_tensor")
    for symbolic_weight in symbolic_weights:
        sw_name = symbolic_weight.name
        name, transpose = convert_tf_weight_name_to_pt_weight_name(
            sw_name,
            start_prefix_to_remove=start_prefix_to_remove,
            tf_weight_shape=symbolic_weight.shape,
            name_scope=_prefix,
        )
        if tf_to_pt_weight_rename is not None:
            aliases = tf_to_pt_weight_rename(name)  # Is a tuple to account for possible name aliasing
            for alias in aliases:  # The aliases are in priority order, take the first one that matches
                if alias in tf_keys_to_pt_keys:
                    name = alias
                    break
            else:
                # If none of the aliases match, just use the first one (it'll be reported as missing)
                name = aliases[0]

        # Find associated numpy array in pytorch model state dict
        if name not in tf_keys_to_pt_keys:
            if allow_missing_keys:
                missing_keys.append(name)
                continue
            elif tf_model._keys_to_ignore_on_load_missing is not None:
                # authorized missing keys don't have to be loaded
                if any(re.search(pat, name) is not None for pat in tf_model._keys_to_ignore_on_load_missing):
                    continue
            raise AttributeError(f"{name} not found in PyTorch model")
        state_dict_name = tf_keys_to_pt_keys[name]
        if is_safetensor_archive:
            array = pt_state_dict.get_tensor(state_dict_name)
        else:
            array = pt_state_dict[state_dict_name]
        try:
            array = apply_transpose(transpose, array, symbolic_weight.shape)
        except tf.errors.InvalidArgumentError as e:
            if not ignore_mismatched_sizes:
                error_msg = str(e)
                error_msg += (
                    "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method."
                )
                raise tf.errors.InvalidArgumentError(error_msg)
            else:
                mismatched_keys.append((name, array.shape, symbolic_weight.shape))
                continue

        tf_loaded_numel += tensor_size(array)

        symbolic_weight.assign(tf.cast(array, symbolic_weight.dtype))
        del array  # Immediately free memory to keep peak usage as low as possible
        all_pytorch_weights.discard(name)

    logger.info(f"Loaded {tf_loaded_numel:,} parameters in the TF 2.0 model.")

    unexpected_keys = list(all_pytorch_weights)

    if tf_model._keys_to_ignore_on_load_missing is not None:
        for pat in tf_model._keys_to_ignore_on_load_missing:
            missing_keys = [k for k in missing_keys if re.search(pat, k) is None]
    if tf_model._keys_to_ignore_on_load_unexpected is not None:
        for pat in tf_model._keys_to_ignore_on_load_unexpected:
            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
    if not skip_logger_warnings:
        _log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name=tf_model.__class__.__name__)

    if output_loading_info:
        loading_info = {
            "missing_keys": missing_keys,
            "unexpected_keys": unexpected_keys,
            "mismatched_keys": mismatched_keys,
        }
        return tf_model, loading_info



SUMMARY:
**Summary: Load PyTorch State Dict in TF 2.0 Model**

This function loads a PyTorch state dictionary into a TensorFlow 2.0 model.

**Purpose and Functionality**

The primary purpose is to map PyTorch weights to their corresponding TensorFlow counterparts, handling prefixes, aliases, and missing keys. The function converts old format to new format if needed.

**Implementation Details**

1. ** renaming keys**: The code maps TF keys to PT keys using predefined renaming rules.
2. **convert_tf_weight_name_to_pt_weightName**: This function transforms weight names from TF to PT format.
3. **load_numel**: Calculates the total number of loaded parameters.
4. **assign_weights**: Assigns PyTorch weights to corresponding TensorFlow weights.

**Key Implementation Details**

* `tf_inputs`: Dummy inputs for the model.
* `_prefix`: Name scope for logging and naming consistency.
* `tf_keys_to_pt_keys`: A dictionary mapping TF keys to PT keys.
* `apply_transpose`: Transposes the array based on the specified transpose order.
* `tensor_size`: Calculates the size of a tensor.

**Dependencies**

* TensorFlow 2.0
* PyTorch state dictionary (safetensors archive or regular dict)

**Requirements**

* `.trainable_weights` and `.non_trainable_weights` attributes of the model.
* `_keys_to_ignore_on_load_missing` and `_keys_to_ignore_on_load_unexpected` attributes of the model.

Note that specific implementation details, such as the naming rules and aliasing mechanism, are not included in this summary.


key: chroma:document
string_value:
CODE:

def load_sharded_pytorch_safetensors_in_tf2_model(
    tf_model,
    safetensors_shards,
    tf_inputs=None,
    allow_missing_keys=False,
    output_loading_info=False,
    _prefix=None,
    tf_to_pt_weight_rename=None,
    ignore_mismatched_sizes=False,
):
    all_loading_infos = []
    for shard in safetensors_shards:
        with safe_open(shard, framework="tf") as safetensors_archive:
            tf_model, loading_info = load_pytorch_state_dict_in_tf2_model(
                tf_model,
                safetensors_archive,
                tf_inputs=tf_inputs,
                allow_missing_keys=allow_missing_keys,
                output_loading_info=True,
                _prefix=_prefix,
                tf_to_pt_weight_rename=tf_to_pt_weight_rename,
                ignore_mismatched_sizes=ignore_mismatched_sizes,
                skip_logger_warnings=True,  # We will emit merged warnings at the end
            )
        all_loading_infos.append(loading_info)
    # Now we just need to merge the loading info
    # Keys are missing only if they're missing in *every* shard
    missing_keys = sorted(set.intersection(*[set(info["missing_keys"]) for info in all_loading_infos]))
    # Keys are unexpected/mismatched if they're unexpected/mismatched in *any* shard
    unexpected_keys = sum([info["unexpected_keys"] for info in all_loading_infos], [])
    mismatched_keys = sum([info["mismatched_keys"] for info in all_loading_infos], [])

    _log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name=tf_model.__class__.__name__)

    if output_loading_info:
        loading_info = {
            "missing_keys": missing_keys,
            "unexpected_keys": unexpected_keys,
            "mismatched_keys": mismatched_keys,
        }
        return tf_model, loading_info



SUMMARY:
**Summary**

 loads PyTorch Safetensors in TF2 model and merges loading information.

**Main Functionality:**

* Loads Safetensors archives for each shard in a sharded dataset
* Merges loading information from all shards to identify missing, unexpected, and mismatched keys.


key: chroma:document
string_value:
CODE:

def load_tf2_checkpoint_in_pytorch_model(
    pt_model, tf_checkpoint_path, tf_inputs=None, allow_missing_keys=False, output_loading_info=False
):
    """
    Load TF 2.0 HDF5 checkpoint in a PyTorch model We use HDF5 to easily do transfer learning (see
    https://github.com/tensorflow/tensorflow/blob/ee16fcac960ae660e0e4496658a366e2f745e1f0/tensorflow/python/keras/engine/network.py#L1352-L1357).
    """
    try:
        import tensorflow as tf  # noqa: F401
        import torch  # noqa: F401
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    import transformers

    from .modeling_tf_utils import load_tf_weights

    logger.info(f"Loading TensorFlow weights from {tf_checkpoint_path}")

    # Instantiate and load the associated TF 2.0 model
    tf_model_class_name = "TF" + pt_model.__class__.__name__  # Add "TF" at the beginning
    tf_model_class = getattr(transformers, tf_model_class_name)
    tf_model = tf_model_class(pt_model.config)

    if tf_inputs is None:
        tf_inputs = tf_model.dummy_inputs

    if tf_inputs is not None:
        tf_model(tf_inputs, training=False)  # Make sure model is built

    load_tf_weights(tf_model, tf_checkpoint_path)

    return load_tf2_model_in_pytorch_model(
        pt_model, tf_model, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info


SUMMARY:
**Summary**
------------

`load_tf2_checkpoint_in_pytorch_model`: Load a TensorFlow 2.0 checkpoint into a PyTorch model.

### Purpose and Functionality
Loads the specified TensorFlow 2.0 checkpoint into a PyTorch model, utilizing transfer learning through HDF5 files.

### Implementation Details

*   Uses `transformers` library to instantiate the corresponding TF 2.0 model class.
*   Loads the checkpoint into the instantiated model using `load_tf_weights`.
*   Optionally loads pre-defined inputs for the model by calling `tf_model(tf_inputs, training=False)`.

### Notable Dependencies and Requirements

*   TensorFlow (`tf`)
*   PyTorch
*   Transformers library (imported)
*   `logger` object (assumed external dependency)

**Key Signatures:**

*   `load_tf_weights(tf_model, tf_checkpoint_path)`
*   `tf_model_class.pt_model.config`
*   `load_tf2_model_in_pytorch_model(pt_model, tf_model, ..., allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info)`


key: chroma:document
string_value:
CODE:

def load_tf2_model_in_pytorch_model(pt_model, tf_model, allow_missing_keys=False, output_loading_info=False):
    """Load TF 2.0 model in a pytorch model"""
    weights = tf_model.weights

    return load_tf2_weights_in_pytorch_model(
        pt_model, weights, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info


SUMMARY:
**Summary**

Loads a TF 2.0 model into a PyTorch model.

### Key Implementation Details

* Utilizes `weights` from the input TF 2.0 model
* Calls `load_tf2_weights_in_pytorch_model` with PyTorch model, weights, and optional flags for missing keys and loading info.

### Notable Dependencies/Requirements

* `tf_model.weights`
* Input `pt_model` (PyTorch model)
* Optional input parameters:
	+ `allow_missing_keys=False`: Flag to specify handling of missing keys
	+ `output_loading_info=False`: Flag to specify outputting loading information


key: chroma:document
string_value:
CODE:

def load_tf2_weights_in_pytorch_model(pt_model, tf_weights, allow_missing_keys=False, output_loading_info=False):
    """Load TF2.0 symbolic weights in a PyTorch model"""
    try:
        import tensorflow as tf  # noqa: F401
        import torch  # noqa: F401
    except ImportError:
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see "
            "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
        )
        raise

    tf_state_dict = {tf_weight.name: tf_weight.numpy() for tf_weight in tf_weights}
    return load_tf2_state_dict_in_pytorch_model(
        pt_model, tf_state_dict, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info


SUMMARY:
**load_tf2_weights_in_pytorch_model Function Summary**

**Main Purpose and Functionality**
---------------------------

This function loads TensorFlow 2.0 symbolic weights into a PyTorch model.

* It takes four inputs: `pt_model` (the PyTorch model), `tf_weights` (TensorFlow 2.0 weights), `allow_missing_keys`, and `output_loading_info`.
* The function attempts to fetch the necessary libraries (PyTorch and TensorFlow) and raises an error if neither is installed.

**Key Implementation Details**
-------------------------

* This function uses dictionary comprehension to create a `tf_state_dict` with TF weights' names as keys and their numpy representations as values.
* It then calls another function, `load_tf2_state_dict_in_pytorch_model`, to load the state dict into the PyTorch model.

**Notable Dependencies or Requirements**
--------------------------------------

* TensorFlow (TF) library
* PyTorch library

**Method Signatures and Interfaces**
-----------------------------------

* The function takes four arguments: `pt_model`, `tf_weights`, `allow_missing_keys`, and `output_loading_info`.
* It returns the loaded state dict in the PyTorch model.


key: chroma:document
string_value:
CODE:

def load_tf2_state_dict_in_pytorch_model(pt_model, tf_state_dict, allow_missing_keys=False, output_loading_info=False):
    import torch

    new_pt_params_dict = {}
    current_pt_params_dict = dict(pt_model.named_parameters())

    # Make sure we are able to load PyTorch base models as well as derived models (with heads)
    # TF models always have a prefix, some of PyTorch models (base ones) don't
    start_prefix_to_remove = ""
    if not any(s.startswith(pt_model.base_model_prefix) for s in current_pt_params_dict.keys()):
        start_prefix_to_remove = pt_model.base_model_prefix + "."

    # Build a map from potential PyTorch weight names to TF 2.0 Variables
    tf_weights_map = {}
    for name, tf_weight in tf_state_dict.items():
        pt_name, transpose = convert_tf_weight_name_to_pt_weight_name(
            name, start_prefix_to_remove=start_prefix_to_remove, tf_weight_shape=tf_weight.shape
        )
        tf_weights_map[pt_name] = (tf_weight, transpose)

    all_tf_weights = set(tf_weights_map.keys())
    loaded_pt_weights_data_ptr = {}
    missing_keys_pt = []
    for pt_weight_name, pt_weight in current_pt_params_dict.items():
        # Handle PyTorch shared weight ()not duplicated in TF 2.0
        if pt_weight.data_ptr() in loaded_pt_weights_data_ptr:
            new_pt_params_dict[pt_weight_name] = loaded_pt_weights_data_ptr[pt_weight.data_ptr()]
            continue

        pt_weight_name_to_check = pt_weight_name
        # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
        key_components = pt_weight_name.split(".")
        name = None
        if key_components[-3::2] == ["parametrizations", "original0"]:
            name = key_components[-2] + "_g"
        elif key_components[-3::2] == ["parametrizations", "original1"]:
            name = key_components[-2] + "_v"
        if name is not None:
            key_components = key_components[:-3] + [name]
            pt_weight_name_to_check = ".".join(key_components)

        # Find associated numpy array in pytorch model state dict
        if pt_weight_name_to_check not in tf_weights_map:
            if allow_missing_keys:
                missing_keys_pt.append(pt_weight_name)
                continue

            raise AttributeError(f"{pt_weight_name} not found in TF 2.0 model")

        array, transpose = tf_weights_map[pt_weight_name_to_check]

        array = apply_transpose(transpose, array, pt_weight.shape, pt_to_tf=False)

        if numpy.isscalar(array):
            array = numpy.array(array)
        if not is_torch_tensor(array) and not is_numpy_array(array):
            array = array.numpy()
        if is_numpy_array(array):
            # Convert to torch tensor
            array = torch.from_numpy(array)

        new_pt_params_dict[pt_weight_name] = array
        loaded_pt_weights_data_ptr[pt_weight.data_ptr()] = array
        all_tf_weights.discard(pt_weight_name)

    missing_keys, unexpected_keys = pt_model.load_state_dict(new_pt_params_dict, strict=False)
    missing_keys += missing_keys_pt

    # Some models may have keys that are not in the state by design, removing them before needlessly warning
    # the user.
    if pt_model._keys_to_ignore_on_load_missing is not None:
        for pat in pt_model._keys_to_ignore_on_load_missing:
            missing_keys = [k for k in missing_keys if re.search(pat, k) is None]

    if pt_model._keys_to_ignore_on_load_unexpected is not None:
        for pat in pt_model._keys_to_ignore_on_load_unexpected:
            unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

    if len(unexpected_keys) > 0:
        logger.warning(
            "Some weights of the TF 2.0 model were not used when initializing the PyTorch model"
            f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"
            f" {pt_model.__class__.__name__} from a TF 2.0 model trained on another task or with another architecture"
            " (e.g. initializing a BertForSequenceClassification model from a TFBertForPreTraining model).\n- This IS"
            f" NOT expected if you are initializing {pt_model.__class__.__name__} from a TF 2.0 model that you expect"
            " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"
            " TFBertForSequenceClassification model)."
        )
    else:
        logger.warning(f"All TF 2.0 model weights were used when initializing {pt_model.__class__.__name__}.\n")
    if len(missing_keys) > 0:
        logger.warning(
            f"Some weights of {pt_model.__class__.__name__} were not initialized from the TF 2.0 model and are newly"
            f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"
            " use it for predictions and inference."
        )
    else:
        logger.warning(
            f"All the weights of {pt_model.__class__.__name__} were initialized from the TF 2.0 model.\n"
            "If your task is similar to the task the model of the checkpoint was trained on, "
            f"you can already use {pt_model.__class__.__name__} for predictions without further training."
        )

    logger.info(f"Weights or buffers not loaded from TF 2.0 model: {all_tf_weights}")

    if output_loading_info:
        loading_info = {"missing_keys": missing_keys, "unexpected_keys": unexpected_keys}
        return pt_model, loading_info



SUMMARY:
**Function Summary:**
`load_tf2_state_dict_in_pytorch_model`

This function loads a pre-trained PyTorch model by transferring weights from a TensorFlow 2.0 model.

**Purpose and Functionality:**

* Loads PyTorch model parameters (`pt_model`) based on a provided TensorFlow 2.0 state dictionary (`tf_state_dict`)
* Handles missing keys, duplicate weights, and unexpected key patterns between the two models
* Logs loading information, including missing keys, unexpected keys, and loaded weights
* Returns the updated `pt_model` with loaded parameters

**Key Implementation Details:**

* Uses a mapping function (`convert_tf_weight_name_to_pt_weight_name`) to convert TF2.0 weight names to PyTorch-compatible names
* Applies transpose operations for compatible arrays between TF2.0 and PyTorch
* Handles cases where certain keys are ignored or replaced during loading

**Notable Dependencies:**

* `torch`
* `numpy` (imported implicitly)

**Interfaces and Inherited Interfaces:**

* `pytorch.nn.Module` (`pt_model`)
* `tensorflow.python.framework ops.Eval`
* `numpy.ndarray`


key: chroma:document
string_value:
CODE:

def parse_flag_from_env(key, default=False):
    try:
        value = os.environ[key]
    except KeyError:
        # KEY isn't set, default to `default`.
        _value = default
    else:
        # KEY is set, convert it to True or False.
        try:
            _value = strtobool(value)
        except ValueError:
            # More values are supported, but let's keep the message simple.
            raise ValueError(f"If set, {key} must be yes or no.")


SUMMARY:
**Summary:**

`parse_flag_from_env` function:

* Retrieves environment variable `key` and defaults to `default=False` if not present.
* Converts the value to a boolean (`True`/`False`) or raises a `ValueError` if the value is not in binary format.

**Key Implementation Details:**

* Uses `os.environ[key]` to access environment variable
* Employes `strtobool` function to convert non-boolean values to booleans

**Notable Dependencies and Requirements:**

* `os` module (for accessing environment variables)
* Internal dependency on `_value` variable to store the converted value


key: chroma:document
string_value:
CODE:

def parse_int_from_env(key, default=None):
    try:
        value = os.environ[key]
    except KeyError:
        _value = default
    else:
        try:
            _value = int(value)
        except ValueError:
            raise ValueError(f"If set, {key} must be a int.")


SUMMARY:
**Summary**

`parse_int_from_env`: Extracts an integer value from an environment variable, returning the parsed integer or a default if the variable is not present. If parsing fails due to non-integer value, raises a `ValueError`.

**Key Implementation Details**

* Uses `os.environ[key]` to retrieve environment variable value.
* Attempts to parse the value as an integer using `int(value)`.
* Fails with a `ValueError` if parsing fails.

**Notable Dependencies and Requirements**

* `os` module for accessing environment variables.


key: chroma:document
string_value:
CODE:

def get_device_count():
    import torch

    if is_torch_xpu_available():
        num_devices = torch.xpu.device_count()
    else:
        num_devices = torch.cuda.device_count()



SUMMARY:
**Code Summary**

`get_device_count()` function returns the total number of devices (Accelerators) available in the system. The implementation uses PyTorch's `xpu` and/or `cuda` modules to determine device count.

**Key Details:**

* Primary purpose: Returns the total number of devices.
* Implementation details:
	+ Uses `torch.xpu.device_count()` to query XPU devices (if available).
	+ Fallbacks to `torch.cuda.device_count()` if XPU is not available.
* Notable dependencies:
	+ PyTorch library (`import torch`)
	+ `is_torch_xpu_available()` function (not shown in this snippet)


key: chroma:document
string_value:
CODE:

def is_pt_tf_cross_test(test_case):
    """
    Decorator marking a test as a test that control interactions between PyTorch and TensorFlow.

    PT+TF tests are skipped by default and we can run only them by setting RUN_PT_TF_CROSS_TESTS environment variable
    to a truthy value and selecting the is_pt_tf_cross_test pytest mark.

    """
    if not _run_pt_tf_cross_tests or not is_torch_available() or not is_tf_available():
        return unittest.skip(reason="test is PT+TF test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:


SUMMARY:
**Summary**

`is_pt_tf_cross_test` Decorator:

*   Marks tests as control interactions between PyTorch and TensorFlow (PT+TF) tests.
*   Skips tests by default unless `RUN_PT_TF_CROSS_TESTS` environment variable is set to a truthy value.

**Key Implementation Details:**

*   Returns `unittest.skip` or sets `pytest.mark.is_pt_tf_cross_test` based on availability of PT+TF tests and underlying framework integrations.

**Notable Dependencies/Requirements:**

*   `_run_pt_tf_cross_tests`: a boolean flag indicating whether to run PT+TF cross-tests.
*   `is_torch_available()`: checks if PyTorch is available for use.
*   `is_tf_available()`: checks if TensorFlow is available for use.
*   `pytest` (optional): import statement included in case of pytest integration.


key: chroma:document
string_value:
CODE:

def is_pt_flax_cross_test(test_case):
    """
    Decorator marking a test as a test that control interactions between PyTorch and Flax

    PT+FLAX tests are skipped by default and we can run only them by setting RUN_PT_FLAX_CROSS_TESTS environment
    variable to a truthy value and selecting the is_pt_flax_cross_test pytest mark.

    """
    if not _run_pt_flax_cross_tests or not is_torch_available() or not is_flax_available():
        return unittest.skip(reason="test is PT+FLAX test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:


SUMMARY:
### Code Summary

**Function:** `is_pt_flax_cross_test`

**Purpose:** Decorator marking tests for interactions between PyTorch and Flax.

**Functionality:**

* Skips PT+FLAX tests by default unless environment variable `RUN_PT_FLAX_CROSS_TESTS` is set to truthy value.
* Returns skipped test with reason "test is PT+FLAX test" if dependencies are not met.
* Imports pytest only when running as main library, and marks test with `@pytest.mark.is_pt_flax_cross_test()` otherwise.

### Notable Dependencies:
- `_run_pt_flax_cross_tests`
- `is_torch_available()`
- `is_flax_available()`

### Key Algorithm:
- Conditional checks for dependencies and environment variable.
- Optional pytest import and mark.


key: chroma:document
string_value:
CODE:

def is_staging_test(test_case):
    """
    Decorator marking a test as a staging test.

    Those tests will run using the staging environment of huggingface.co instead of the real model hub.
    """
    if not _run_staging:
        return unittest.skip(reason="test is staging test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:


SUMMARY:
**Function Summary: `is_staging_test` Decorator**

### Purpose and Functionality

The `is_staging_test` function serves as a decorator to mark test cases as staging tests. It substitutes the real model hub with the huggingface.co staging environment.

### Key Implementation Details

*   Uses the `unittest.skip` function with a custom reason if `_run_staging` is `False`.
*   Attempts to import `pytest` when `_run_staging` is `True`, and sets the decorated test case using `pytest.mark.is_staging_test()` if successful.
*   Returns the test case as-is if `pytest` import fails.

### Notable Dependencies and Requirements

*   Relies on `_run_staging` variable
*   Imports the `unittest`, `pytest` libraries (`but only imports pytest in specific condition`)
*   Implies a setup that uses huggingface.co (via model hub) for staging environment


key: chroma:document
string_value:
CODE:

def is_pipeline_test(test_case):
    """
    Decorator marking a test as a pipeline test. If RUN_PIPELINE_TESTS is set to a falsy value, those tests will be
    skipped.
    """
    if not _run_pipeline_tests:
        return unittest.skip(reason="test is pipeline test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:


SUMMARY:
**Summary**

This function is a decorator (`is_pipeline_test`) that marks a test case as belonging to a pipeline. If `RUN_PIPELINE_TESTS` is False, the test is skipped. Otherwise, it tries to import and use Pytest's `@pytest.mark.is_pipeline_test()` marker.

**Key Details:**

* Marks test as pipeline test
* Skips test if `RUN_PIPELINE_TESTS` is False
* Imports Pytest library

**Implementation Notes:**

* Uses `unittest.skip` for skipping if condition is false
* Returns decorated test case with Pytest marker if import successful


key: chroma:document
string_value:
CODE:

def is_agent_test(test_case):
    """
    Decorator marking a test as an agent test. If RUN_TOOL_TESTS is set to a falsy value, those tests will be skipped.
    """
    if not _run_agent_tests:
        return unittest.skip(reason="test is an agent test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:


SUMMARY:
**Summary**

The `is_agent_test` function is a decorator that marks a test as an agent test and skips it if `RUN_TOOL_TESTS` is falsy. It uses the `unittest.skip` and `pytest.mark.is_agent_test` functions to implement this functionality.

**Key Details**

* Marks a test as an agent test using `pytest.mark.is_agent_test()`
* Skips tests if `RUN_TOOL TESTING` is falsy, using `unittest.skip`
* Requires no hard dependency on pytest, only imports it conditionally

**Notable Dependencies/Interfaces**

* `_run_agent_tests`: a flag indicating whether to run agent tests
* `unittest.skip`: a function from the unittest framework for skipping tests
* `pytest.mark.is_agent_test`: a decorator from the pytest framework for marking tests as agent tests


key: chroma:document
string_value:
CODE:

def slow(test_case):
    """
    Decorator marking a test as slow.

    Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.

    """


SUMMARY:
**Summary**

Function: `slow`

Purpose:
Marks a test case as marked for running slowly using the `%s` test marker.

Key Implementation Details:

* Skips test by default unless `_run_slow_tests` environment variable is set to truthy value.
* Uses `unittest.skipUnless` decorater with `RUN_SLOW` marker.

Notable Dependencies:

* `unittest`
* `_run_slow_tests` (environment variable, assuming an external setting)


key: chroma:document
string_value:
CODE:

def tooslow(test_case):
    """
    Decorator marking a test as too slow.

    Slow tests are skipped while they're in the process of being fixed. No test should stay tagged as "tooslow" as
    these will not be tested by the CI.

    """


SUMMARY:
**Summary:**

`tooslow` decorator for marking tests as too slow, skipping them in CI until fixed.

**Key Details:**

* Marks a test case using `unittest.skip`
* Reason for skipping: "test is too slow"
* No tests should be tagged with this reason as they are not tested by CI

**Implementation:**

```
def tooslow(test_case):
    return unittest.skip(reason="test is too slow")(test_case)
```


key: chroma:document
string_value:
CODE:

def skip_if_not_implemented(test_func):
    @functools.wraps(test_func)
    def wrapper(*args, **kwargs):
        try:
            return test_func(*args, **kwargs)
        except NotImplementedError as e:
            raise unittest.SkipTest(f"Test skipped due to NotImplementedError: {e}")



SUMMARY:
**Summary: Skipping Test Functions with Not Implemented Error**

The `skip_if_not_implemented` function is a decorator that wraps a test function, skipping its execution if it raises a `NotImplementedError`. The wrapped test function's original behavior is preserved using `functools.wraps`.

**Key Implementation Details:**

* Uses the `@functools.wraps` decorator to preserve original test function metadata.
* Wraps the original test function with a try-except block, catching `NotImplementedError` exceptions.

**Primary Purpose:** Skipping a test function during execution if it raises a `NotImplementedError`.

**Important Method Signature/Interface:**

* `test_func(*args, **kwargs)`: The original test function to be wrapped and potentially skipped.

**External Dependencies/Requirements:**

* `unittest` library (specifically `SkipTest` exception)
* `functools` library (specifically `wraps` decorator)


key: chroma:document
string_value:
CODE:

def apply_skip_if_not_implemented(cls):
    """
    Class decorator to apply @skip_if_not_implemented to all test methods.
    """
    for attr_name in dir(cls):
        if attr_name.startswith("test_"):
            attr = getattr(cls, attr_name)
            if callable(attr):
                setattr(cls, attr_name, skip_if_not_implemented(attr))


SUMMARY:
**Function Summary**

`apply_skip_if_not_implemented`: Decorator function to apply `skip_if_not_implemented` to all test methods in a class.

**Key Implementation Details:**

* Iterates through attributes of the input class (`cls`)
* Filters and calls `skip_if_not_implemented` on callable attributes with names starting with "test_"
* Modifies the original class by setting decorated attributes

**Notable Dependencies/Requirements:**

* `skip_if_not_implemented` decorator function (not shown)
* No external dependencies mentioned


key: chroma:document
string_value:
CODE:

def custom_tokenizers(test_case):
    """
    Decorator marking a test for a custom tokenizer.

    Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS
    environment variable to a truthy value to run them.
    """


SUMMARY:
**Summary:**

`custom_tokensize_decorator`: A test decorator marking a test for a custom tokenizer. Requires the `RUN_CUSTOM_TOKENIZERS` environment variable to be set to run. Skips tests by default.

**Key Implementation Details:**

* Decorators are used to mark tests for custom tokenizers.
* Tests can only be run if `RUN_CUSTOM_TOKENIZERS` is set to true (`unittest.skipUnless`).

**Notable Dependencies/Requirements:**

* `_run_custom_tokenizers`: an external function or decorator marking the custom tokenizer functionality.


key: chroma:document
string_value:
CODE:

def require_bs4(test_case):
    """
    Decorator marking a test that requires BeautifulSoup4. These tests are skipped when BeautifulSoup4 isn't installed.
    """


SUMMARY:
**Summary**

`require_bs4` is a decorator function marking test cases that require the installation of BeautifulSoup4. It skips tests when BeautifulSoup4 is not available.

**Key Details**

* Marks `test_case` for skipping unless BeautifulSoup4 is installed
* Uses `unittest.skipUnless` with condition checked by `is_bs4_available()`
* Checks if BeautifulSoup4 is installed and requires it for test execution


key: chroma:document
string_value:
CODE:

def require_galore_torch(test_case):
    """
    Decorator marking a test that requires GaLore. These tests are skipped when GaLore isn't installed.
    https://github.com/jiaweizzhao/GaLore
    """


SUMMARY:
**Function Summary: `require_galore_torch`**

This is a decorator function that marks test cases requiring the installation and availability of the GaLore library. It skips tests if GaLore isn't installed, using the `unittest.skipUnless` context manager.

**Key Details:**

- `test_case`: Input test case to be decorated
- `is_galore_torch_available()`: External dependency function to check for GaLore installation availability

This decorator ensures that tests requiring GaLore are skipped if it's not installed, using `unittest.skipUnless`.


key: chroma:document
string_value:
CODE:

def require_lomo(test_case):
    """
    Decorator marking a test that requires LOMO. These tests are skipped when LOMO-optim isn't installed.
    https://github.com/OpenLMLab/LOMO
    """


SUMMARY:
**Function Summary**

`require_lomo` is a decorator that marks test cases as requiring LOMO, skipping them unless LOMO-optim is installed. 

**Key Details:**
- Decorator skips tests when `is_lomo_available()` returns False
- Uses `unittest.skipUnless` to skip testsconditionally


key: chroma:document
string_value:
CODE:

def require_grokadamw(test_case):
    """
    Decorator marking a test that requires GrokAdamW. These tests are skipped when GrokAdamW isn't installed.
    """


SUMMARY:
**Function Summary:**

```python
require_grokadamw(test_case):
  Decors a test case with the requirement to run only when GrokAdamW is installed.
  
  This function utilizes Python's `unittest` module and skips the decorated test case if GrokAdamW is not available.
```

**Key Implementation Details:**

- Utilizes `unittest.skipUnless` decorator to conditionally skip tests based on `is_grokadamw_available()`.
- Requires presence of `GrokAdamW` to be installed for test execution.

**Notable Dependencies/Requirements:**

- Python's `unittest` module
- `is_grokadamw_available()` function (not shown in this snippet)


key: chroma:document
string_value:
CODE:

def require_schedulefree(test_case):
    """
    Decorator marking a test that requires schedulefree. These tests are skipped when schedulefree isn't installed.
    https://github.com/facebookresearch/schedule_free
    """


SUMMARY:
**Summary**

`require_schedulefree` Decorator:

Marks a test case as requiring ScheduleFree (a unit testing library).
This decorator skippers the test when ScheduleFree is not installed.

**Key Details:**

* Returns an `unittest.skipUnless` wrapper around `is_schedulefree_available()`
* Test case is decorated with "test requires schedulefree" error message

**Implementation Details:**

No additional implementation details are provided, as it's a decorator function. 

**Dependencies and Requirements:**

* Unittest library
* schedule_free (available check via `is_schedulefree_available()` method)


key: chroma:document
string_value:
CODE:

def require_cv2(test_case):
    """
    Decorator marking a test that requires OpenCV.

    These tests are skipped when OpenCV isn't installed.

    """


SUMMARY:
**Summary:**

### `require_cv2` Decorator

Marks a test to require OpenCV, skipping if not available. Skips when `is_cv2_available()` returns `False`.

**Key Implementation Details:**

* Uses `unittest.skipUnless` with `is_cv2_available()` check
* Wraps the original test case (`test_case`) in a skip condition

**Context:** OpenCV tests are tagged for skipping on non-OpenCV environments.


key: chroma:document
string_value:
CODE:

def require_levenshtein(test_case):
    """
    Decorator marking a test that requires Levenshtein.

    These tests are skipped when Levenshtein isn't installed.

    """


SUMMARY:
**Summary:**

This function `require_levenshtein` is a decorator that skips test cases marked as requiring Levenshtein unless the Levenshtein algorithm is installed. It uses `unittest.skipUnless` to conditionally skip tests based on the availability of Levenshtein.

**Key Details:**

* **Function Signature:** `def require_levenshtein(test_case):`
* **Purpose:** Marks a test as requiring Levenshtein and skips it if not available.
* **Implementation:** Uses `unittest.skipUnless` to conditionally skip tests.
* **Notable Dependency:** `is_levenshtein_available()` checks for the availability of the Levenshtein algorithm.

**External Dependencies:**

* `unittest`: The unit testing framework used in Python.
* `is_levenshtein_available()`: A function checking for the availability of the Levenshtein algorithm.


key: chroma:document
string_value:
CODE:

def require_nltk(test_case):
    """
    Decorator marking a test that requires NLTK.

    These tests are skipped when NLTK isn't installed.

    """


SUMMARY:
**Function Summary**

`require_nltk`: Decorator marking test cases that require Natural Language Toolkit (NLTK) installation. Skips tests when NLTK is not available.

**Key Details**

* **Purpose**: Ensures integration with NLTK for required tests.
* **Implementation**: Uses `unittest.skipUnless` to skip test case if NLTK is unavailable.
* **Dependency**: `is_nltk_available()` function, which is not shown in this code snippet.


key: chroma:document
string_value:
CODE:

def require_accelerate(test_case, min_version: str = ACCELERATE_MIN_VERSION):
    """
    Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed.
    """
    return unittest.skipUnless(
        is_accelerate_available(min_version), f"test requires accelerate version >= {min_version}"


SUMMARY:
**Summary:**

`require_accelerate` Decorator
================================

This function is a Python decorator that marks tests requiring `accelerate` to automatically skip if the dependency isn't installed. It ensures test cases use a specified minimum version of `accelerate`.

**Key Implementation Details:**
* Checks for availability of `accelerate` using `is_accelerate_available(min_version)`.
* Skips tests unless `accelerate` meets the minimum version requirement (`>= {min_version}`).

**Method Signature:**
- `require_accelerate(test_case, min_version: str = ACCELERATE_MIN_VERSION)`


key: chroma:document
string_value:
CODE:

def require_gguf(test_case, min_version: str = GGUF_MIN_VERSION):
    """
    Decorator marking a test that requires ggguf. These tests are skipped when gguf isn't installed.
    """
    return unittest.skipUnless(is_gguf_available(min_version), f"test requires gguf version >= {min_version}")(
        test_case


SUMMARY:
**Concise Summary:**

`require_gguf` Decorator
==========================

Marks a test as requiring ggguf (>= `GGUF_MIN_VERSION`) and skips it if not installed. If gguf is available, the test is executed.

### Implementation Details:

* Checks if gguf is available using `is_gguf_available(min_version)`
* Skips test using `unittest.skipUnless` if gguf version doesn't meet minimum requirement

### Key Interface/Function:
```python
def is_gguf_available(min_version: str) -> bool
```

### Notable Dependencies:
- `unittest`
- `GGUF_MIN_VERSION` (version of ggguf required for testing)


key: chroma:document
string_value:
CODE:

def require_fsdp(test_case, min_version: str = "1.12.0"):
    """
    Decorator marking a test that requires fsdp. These tests are skipped when fsdp isn't installed.
    """
    return unittest.skipUnless(is_fsdp_available(min_version), f"test requires torch version >= {min_version}")(
        test_case


SUMMARY:
**Function Summary: `require_fsdp`**

**Purpose:** Decorator to mark tests that require fsdp (FSDP) installation and skip if FSDP is not available.

**Primary Purpose:** Ensure fsdp is installed and compatible with the specified version (`min_version`) before running a test case.

**Key Implementation Details:**

* Uses `unittest.skipUnless` to conditionally skip a test case when FSDP is unavailable.
* Invokes `is_fsdp_available(min_version)` function to check for FSDP availability with the specified minimum version.
* Returns a decorated test case with an error message indicating the version requirement.

**Dependencies:**

* `torch`: referenced by `is_fsdp_available(min_version)`

**External Dependencies:** 

* `torch` package
* `unittest` module


key: chroma:document
string_value:
CODE:

def require_g2p_en(test_case):
    """
    Decorator marking a test that requires g2p_en. These tests are skipped when SentencePiece isn't installed.
    """


SUMMARY:
**Function Summary: `require_g2p_en` Decorator**

* **Purpose**: A test case decorator that skips tests requiring `g2p_en` if the dependency is not installed or available.

**Key Details:**
- `is_g2p_en_available()` function is used to check the availability of g2p_en.
- The decorator uses the `unittest.skipUnless` method to skip tests when the dependency is not available.
- The decorated test case is wrapped with a `test_case` object from the unittest framework.

**Notable Dependencies:**
- `g2p_en`: A library that provides a spell checking model.
- `SentencePiece`: A library used to integrate thespell checking functionality.


key: chroma:document
string_value:
CODE:

def require_safetensors(test_case):
    """
    Decorator marking a test that requires safetensors. These tests are skipped when safetensors isn't installed.
    """


SUMMARY:
**Summary:**

This function is a Python decorator named `require_safetensors` that skips a test case (`test_case`) unless the `is_safetensors_available()` function returns True, indicating the presence of the safetensor library.

Key Details:

* Primary purpose: Skip test cases based on availability of safetensor library.
* Implementation: Uses `unittest.skipUnless` with a condition check from `is_safetensors_available()`.
* Notable dependencies: Safetensor library ( referenced by `is_safetensors_available()` ). 

**No external dependencies listed as it's solely focused on its own internal usage.**


key: chroma:document
string_value:
CODE:

def require_rjieba(test_case):
    """
    Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed.
    """


SUMMARY:
**Summary:**

`require_rjieba` Decorator
=============================

* **Functionality:** Skips test cases that require `rjieba` unless it's installed.
* **Design Choice:** Uses `unittest.skipUnless` to conditionally skip tests based on the availability of `rjieba`.
* **Key Implementation:**
	+ Calls `is_rjieba_available()` to determine if `rjieba` is installed.
	+ Applies `unittest.skipUnless` with an error message if `rjieba` isn't available.

**Code Context:**

This decorator enhances unit testing by marking test cases that rely on `rjieba`. When the test case is executed, it checks if `rjieba` is installed and skips the test only if not found.


key: chroma:document
string_value:
CODE:

def require_jieba(test_case):
    """
    Decorator marking a test that requires jieba. These tests are skipped when jieba isn't installed.
    """


SUMMARY:
**Code Summary**

`require_jieba` is a decorator marking test cases that utilize JIEBA (a Chinese word segmentation library) as mandatory. The decorated tests are skipped if JIEBA is not installed.

**Details**

* Marking tests with `@require_jieba(test_case)` indicates they must use JIEBA for execution.
* Tests requiring JIEBA will be skipped when the library is not available (`is_jieba_available()` returns False).
* This decorator ensures that the test suite can run reliably only if JIEBA is installed.


key: chroma:document
string_value:
CODE:

def require_jinja(test_case):
    """
    Decorator marking a test that requires jinja. These tests are skipped when jinja isn't installed.
    """


SUMMARY:
**Summary:**

* **Purpose and Functionality:** Decorator to mark test cases that require Jinja installation, skipping the test if Jinja isn't available.

**Key Implementation Details:**

* Uses `unittest.skipUnless` to skip the test if Jinja isn't available.

**Notable Dependencies/Requirements:**

* [is_jinja_available()](external func) - checks if Jinja is installed.
* `unittest` - testing framework (built-in Python library).
* `unittest.skipUnless` - decorator skipping tests based on a condition.


key: chroma:document
string_value:
CODE:

def require_tf2onnx(test_case):


SUMMARY:
**Summary:**

`require_tf2onnx` test case modifier ensures that tests requiring `tf2onnx` are only executed when `tf2onnx` is available. It uses `unittest.skipUnless` to skip the test if `is_tf2onnx_available()` returns False.

**Key Details:**

* Modifier function: `require_tf2onnx`
* Purpose: Skips tests if `tf2onnx` is not available
* Implementation: Uses `is_tf2onnx_available()` and `unittest.skipUnless`

**Notable Dependencies:**

* `is_tf2onnx_available()`: A test utility function checking for availability of `tf2onnx`
* `unittest`: Unit testing framework

**Interfaces/Classes:**

* `unittest` (unit testing framework)
* `test_case` (unit test case instance)


key: chroma:document
string_value:
CODE:

def require_onnx(test_case):


SUMMARY:
**Summary:**

`require_onnx` is a test case decorator that uses `unittest.skipUnless` to skip a test if ONNX (Open Neural Network Exchange) support is not available, unless testing is explicitly enabled.

**Key Details:**
- `test_case` parameter indicates the decorated test case.
- Uses `is_onnx_available()` function to determine ONNX availability and `"test requires ONNX"` error message for skipping tests.


key: chroma:document
string_value:
CODE:

def require_timm(test_case):
    """
    Decorator marking a test that requires Timm.

    These tests are skipped when Timm isn't installed.

    """


SUMMARY:
**Function Summary: `require_timm`**

* **Purpose**: A decorator that marks test cases that require the Timm library, skipping them unless Timm is installed.

* **Key Implementation Details**:
	+ Uses the `unittest.skipUnless` function to conditionally skip tests based on Timm availability.
	+ Calls `is_timm_available()` to determine whether Timm is installed or not.

* **Dependencies**:
	+ `unittest` module
	+ `is_timm_available()` function (not shown in this code snippet)

* **Notes**: This decorator is designed to protect tests that rely on Timm from running without it, ensuring that the test suite remains stable and predictable.


key: chroma:document
string_value:
CODE:

def require_natten(test_case):
    """
    Decorator marking a test that requires NATTEN.

    These tests are skipped when NATTEN isn't installed.

    """


SUMMARY:
**Summary**

* Function: `require_natten`
* Purpose: A decorator that skips or runs tests depending on the availability of NATTEN.
* Primary purpose: Indicates which tests require NATTEN installation.

**Key implementation details**

* Uses `unittest.skipUnless` with a call to `is_natten_available()`.
* Utilizes the `test_case` object passed as an argument.

**Notable dependencies or requirements**

* None explicitly mentioned, but relies on:
	+ `unittest` module
	+ `is_natten_available()` function (not shown in this snippet)


key: chroma:document
string_value:
CODE:

def require_torch(test_case):
    """
    Decorator marking a test that requires PyTorch.

    These tests are skipped when PyTorch isn't installed.

    """


SUMMARY:
**Summary:**

`require_torch` Decorator:

*   Marks a test as requiring PyTorch.
*   Skips the test if PyTorch is not available; otherwise, runs it.
*   Utilizes `is_torch_available()` to check for PyTorch installation.

Notable Details:

*   Uses `unittest.skipUnless()` with custom error message.
*   Relies on `is_torch_available()` external function or variable (not defined in this snippet).


key: chroma:document
string_value:
CODE:

def require_flash_attn(test_case):
    """
    Decorator marking a test that requires Flash Attention.

    These tests are skipped when Flash Attention isn't installed.

    """


SUMMARY:
**Concise Summary**

`require_flash_attn` is a decorator that marks a test for requirement of Flash Attention. It skips the test if Flash Attention isn't installed, warning about the missing dependency.

**Key Implementation Details:**

* Used `unittest.skipUnless` to skip test if Flash Attention is not available
* Calls `is_flash_attn_2_available()` interface to check Flash Attention availability

**Dependence:**

* Requires Flash Attention feature (`is_flash_attn_2_available()`)


key: chroma:document
string_value:
CODE:

def require_torch_sdpa(test_case):
    """
    Decorator marking a test that requires PyTorch's SDPA.

    These tests are skipped when requirements are not met (torch version).
    """


SUMMARY:
Here is a concise and focused summary of the provided function code chunk:

**Function Summary: `require_torch_sdpa`**

*   Decorator marking a test that requires availability of PyTorch's SDPA (STeady-state Dynamic Programming Algorithm).
*   Skips the test if PyTorch SDPA is not installed.
*   Tests are skipped with an error message indicating requiring PyTorch SDPA.

Key Details:

- **Function Signature:** `require_torch_sdpa(test_case)`
- **Primary Purpose:** Verifies availability of PyTorch's SDPA and enforces its requirement for tests.
- **Implementation:** Utilizes `unittest.skipUnless` to conditionally skip tests based on the availability of PyTorch SDPA, checked by the `is_torch_sdpa_available()` function.


key: chroma:document
string_value:
CODE:

def require_read_token(fn):
    """
    A decorator that loads the HF token for tests that require to load gated models.
    """
    token = os.getenv("HF_HUB_READ_TOKEN")

    @wraps(fn)
    def _inner(*args, **kwargs):
        if token is not None:
            with patch("huggingface_hub.utils._headers.get_token", return_value=token):
                return fn(*args, **kwargs)
        else:  # Allow running locally with the default token env variable
            return fn(*args, **kwargs)



SUMMARY:
**Summary**

The `require_read_token` function is a decorator that loads the HF token from the environment variable "HF_HUB_READ_TOKEN" before running a target function (`fn`). It allows for loading gated models in tests by overriding the default token behavior. 

**Key Implementation Details**

* Uses `patch` to temporarily override the default token retrieval mechanism
* Returns `fn` with modified headers only if a valid token is set in the environment variable

**Dependencies/Requirements**

* `os` (for accessing environment variables)
* `wraps` (from `functools`) or equivalent decorator for preserving original function signature and docstring
* `patch` (from the testing library, e.g., `unittest.mock`)


key: chroma:document
string_value:
CODE:

def require_peft(test_case):
    """
    Decorator marking a test that requires PEFT.

    These tests are skipped when PEFT isn't installed.

    """


SUMMARY:
**Function Summary**

`require_peft`: A decorator function marking a test case that requires the PEFT library. Skips the test when PEFT is not installed.

**Key Details**

* Decorates a `test_case` with a skip condition using `unittest.skipUnless`.
* Conditionally skips tests unless `is_peft_available()` returns `True`, allowing for installation checks.

**Notable Dependencies**

* `unittest`: The standard Python unit testing framework.
* `is_peft_available()`: A function checking if PEFT is installed (implementation not shown).


key: chroma:document
string_value:
CODE:

def require_torchvision(test_case):
    """
    Decorator marking a test that requires Torchvision.

    These tests are skipped when Torchvision isn't installed.

    """


SUMMARY:
**Summary**

This is a Python decorator function `require_torchvision` that marks a test case (`test_case`) as requiring torchvision installation. When torchvision isn't available, the test is skipped using the `unittest.skipUnless` context manager. 

The primary purpose is to ensure tests are only run if torchvision is installed.

**Key Implementation Details:**

* Uses `unittest.skipUnless` with condition `is_torchvision_available()`
* Affects all test cases marked by the decorator

**Notable Dependencies/Requirements:**

* `torchvision`: required for the test case execution
* `is_torchVisionAvailable()` function (exact name not provided)
       


key: chroma:document
string_value:
CODE:

def require_torch_or_tf(test_case):
    """
    Decorator marking a test that requires PyTorch or TensorFlow.

    These tests are skipped when neither PyTorch not TensorFlow is installed.

    """
    return unittest.skipUnless(is_torch_available() or is_tf_available(), "test requires PyTorch or TensorFlow")(
        test_case


SUMMARY:
**Code Summary**

Adds a decorator to mark tests that require either PyTorch or TensorFlow. The decorated tests are skipped when neither library is installed.

**Key Implementation Details**

* Uses `unittest.skipUnless` decorator
* Checks availability of PyTorch and/or TensorFlow using `is_torch_available()` and `is_tf_available()`

**Primary Purpose**
Requires specific libraries for the test execution.

**Notable Dependencies or Requirements**

* `torch` (PyTorch)
* `tf` (TensorFlow)


key: chroma:document
string_value:
CODE:

def require_intel_extension_for_pytorch(test_case):
    """
    Decorator marking a test that requires Intel Extension for PyTorch.

    These tests are skipped when Intel Extension for PyTorch isn't installed or it does not match current PyTorch
    version.

    """
    return unittest.skipUnless(
        is_ipex_available(),
        "test requires Intel Extension for PyTorch to be installed and match current PyTorch version, see"
        " https://github.com/intel/intel-extension-for-pytorch",


SUMMARY:
**Function Summary: `require_intel_extension_for_pytorch`**

This is a decorator function that marks test cases as skipping when Intel Extension for PyTorch is not installed or incompatible with the current PyTorch version.

**Implementation Details:**

* Utilizes `unittest.skipUnless` to skip tests based on the presence and compatibility of Intel Extension.
* Calls `is_ipex_available()` function to verify availability of Intel Extension.
* Passes a custom error message referencing the official GitHub documentation for Intel Extension installation and compatibility.

**Notable Dependencies/Requirements:**

* `unittest` module
* `is_ipex_available()` function (not included in this code snippet)
* Intel Extension for PyTorch (requires installation and matching current PyTorch version)

Note: The exact implementation dependencies are excluded to maintain focus on the primary purpose of the function.


key: chroma:document
string_value:
CODE:

def require_tensorflow_probability(test_case):
    """
    Decorator marking a test that requires TensorFlow probability.

    These tests are skipped when TensorFlow probability isn't installed.

    """
    return unittest.skipUnless(is_tensorflow_probability_available(), "test requires TensorFlow probability")(
        test_case


SUMMARY:
**TF Probability Test Decorator**

* **Purpose and Functionality**: Tests not compatible with TensorFlow Probability will be skipped when `is_tensorflow_probability_available()` returns False.
* **Key Algorithm/Design Choice**: Uses the `unittest.skipUnless` decorator to conditionally skip tests based on availability of TensorFlow Probability.

**Note:** `is.tensorflow_probability_available()` is assumed to be a function that checks if TensorFlow Probability is installed and available.


key: chroma:document
string_value:
CODE:

def require_torchaudio(test_case):
    """
    Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed.
    """


SUMMARY:
**Summary:**
The `require_torchaudio` decorator is a test decorator that skips a test case unless PyTorch Audio (torchaudio) is available.

**Key Implementation Details:**

* Decorates a test case using the `unittest.skipUnless` function.
* Ensures the decorator checks if torchaudio is available by calling the `is_torchaudio_available()` function.
* Skips the test case if torchaudio is not available, and provides a specific error message.

**Notable Dependencies/Requirements:**

* PyTorch Audio (torchaudio) module requires.


key: chroma:document
string_value:
CODE:

def require_tf(test_case):
    """
    Decorator marking a test that requires TensorFlow. These tests are skipped when TensorFlow isn't installed.
    """


SUMMARY:
**Summary**

`require_tf` is a decorator that marks a test case as requiring TensorFlow, skipping the test when TensorFlow is not available. It utilizes `unittest.skipUnless` to check for availability and skips tests if TensorFlow installation is absent.

**Key Details:**

* Decorator marking test cases for TensorFlow dependency
* Skips tests with `"test requires TensorFlow"` error message when TensorFlow is not installed

No notable dependencies or requirements besides TensorFlow, which can be verified using the `is_tf_available` function.


key: chroma:document
string_value:
CODE:

def require_flax(test_case):
    """
    Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed
    """


SUMMARY:
**Code Summary:**

`require_flax` Decorator
==========================

### Purpose and Functionality

Marks test cases requiring JAX and Flax, skipping tests if dependency is not installed.

### Key Implementation Details:

* Uses `unittest.skipUnless` with the result of `is_flax_available()` to skip test case.
* The decorator can be applied using `@require_flax` syntax.

### Notable Dependencies or Requirements

* Requires `jax`, `flax` libraries for functionality.


key: chroma:document
string_value:
CODE:

def require_sentencepiece(test_case):
    """
    Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
    """


SUMMARY:
**Summary:**

`require_sentencepiece` is a unit test decorator marking test cases that require SentencePiece, skipping them if the library isn't installed.

**Key Details:**

* Decorator skips tests if `is_sentencepiece_available()` returns False
* Tests are skipped with explicit error message "test requires SentencePiece"
* Reverts normal test execution using `unittest.skipUnless`


key: chroma:document
string_value:
CODE:

def require_sacremoses(test_case):
    """
    Decorator marking a test that requires Sacremoses. These tests are skipped when Sacremoses isn't installed.
    """


SUMMARY:
**Summary**

`require_sacremoses` is a test case decorator that skips tests requiring Sacremoses when it's not installed. It wraps the `test_case` with `unittest.skipUnless`, which checks if `is_sacremoses_available()` returns True.

**Key Details**

* Uses `unittest.skipUnless`
* Skips tests when `is_sacremoses_available()` returns False
* Requires Sacremosedebugging

**Dependencies**

* `is_sacremoses_available()`: function checking availability of Sacremoses installation.


key: chroma:document
string_value:
CODE:

def require_seqio(test_case):
    """
    Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
    """


SUMMARY:
**Code Summary: `require_seqio` Decorator**

### Main Purpose and Functionality
The `require_seqio` decorator marks test cases that require installation of the SentencePiece library for execution. When Sentinelpiece isn't installed, these tests are skipped.

### Key Implementation Details

* Uses `unittest.skipUnless` to skip tests if SentensePipe is not available.
* Calls the `is_seqio_available()` function to check for Sentinelpipe availability.

### Notable Dependencies or Requirements
Requires installation of `SentencePiece`.

### Interfaces and Inherited Traits
Extends the behavior of `unittest.skipUnless`.


key: chroma:document
string_value:
CODE:

def require_scipy(test_case):
    """
    Decorator marking a test that requires Scipy. These tests are skipped when SentencePiece isn't installed.
    """


SUMMARY:
**Summary:**

This function `require_scipy` is a decorator used to mark test cases that require the SciPy library. When deployed, it skips the test if SciPy is not available.

**Key Implementation Details:**

* Uses `unittest.skipUnless` with a condition based on `is_scipy_available()`
* Throws an error message "test requires Scipy" if SciPy is unavailable

**Notable Dependencies or Requirements:**

* Requires SciPy library (notably, the `is_scipy_available()` function)
* Inherently ties to `unittest` testing framework
* Implements a decorator pattern to dynamically control test execution


key: chroma:document
string_value:
CODE:

def require_tokenizers(test_case):
    """
    Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed.
    """


SUMMARY:
**Summary**

This function, `require_tokenizers`, is a decorator that marks test cases requiring Tokenzie library (not 🤗 Tokenizers as mentioned in the code). It skips tests that are not compatible with Tokenzie when it's not installed.

* **Purpose**: Ensure compatibility of tests with Tokenzie library.
* **Key details**:
	+ Uses `unittest.skipUnless` to skip tests when Tokenzie is unavailable.
	+ Requires Tok McKenzie (`tok McKenzie`) library (noted as Tokenzie in the code comment).
* **Dependencies**:
	+ `Tok McKenzie` library.


key: chroma:document
string_value:
CODE:

def require_tensorflow_text(test_case):
    """
    Decorator marking a test that requires tensorflow_text. These tests are skipped when tensroflow_text isn't
    installed.
    """


SUMMARY:
**Summary:**

`require_tensorflow_text` Decorator:
 
* Marks test cases that require TensorFlow Text installation and automatically skip them if `is_tensorflow_text_available()` returns False.

### Implementation Details:

* Utilizes the `unittest.skipUnless` function to conditionally skip tests based on TensorFlow Text availability.
* Relies on the `is_tensorflow_text_available()` helper function to determine whether TensorFlow Text is installed.

### Notable Dependencies/Requirements:

* `tensorflow_text`: The dependency required for this test case (uninstalling or not installing this package causes the test to be skipped).
* `unittest.skipUnless()`
* `is_tensorflow_text_available()`: An internal function checking the availability of TensorFlow Text.


key: chroma:document
string_value:
CODE:

def require_keras_nlp(test_case):
    """
    Decorator marking a test that requires keras_nlp. These tests are skipped when keras_nlp isn't installed.
    """


SUMMARY:
**Summary**

This is a Python decorator function `require_keras_nlp` that marks test cases requiring the installation of Keras NLP library. It skips tests if `is_keras_nlp_available()` returns False.

**Key Points**

* **Purpose**: Skips non-Keras NLP installed tests.
* **Functionality**:
	+ Uses `unittest.skipUnless` to skip test case if Keras NLP is not available.
	+ Calls `is_keras_nlp_available()` function as a condition for skipping the test.
* **Important details**:
	+ Decorator marks test cases requiring Keras NLP library installation.
	+ Tests are skipped if Keras NLP is not installed.


key: chroma:document
string_value:
CODE:

def require_pandas(test_case):
    """
    Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed.
    """


SUMMARY:
**Function Summary**

`require_pandas` is a decorator function used to mark test cases that require the presence of pandas. It skips tests if pandas is not available.

* **Purpose:** Marks test case as requiring pandas installation.
* **Key Implementation Details:**
	+ Uses `unittest.skipUnless` with condition check from `is_pandas_available()`.
	+ Decorates the `test_case` parameter with a message indicating requirement of pandas.
* **Notable Dependencies/Requirements:** pandas
* **Interfaces Implemented:** `unittest.TestCase`
* **External Dependencies:** None


key: chroma:document
string_value:
CODE:

def require_pytesseract(test_case):
    """
    Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed.
    """


SUMMARY:
**Summary:**

`require_pytesseract` is a test case decorator, marked by PyTesseract's installation status. It uses `unittest.skipUnless` to conditionally skip test runs when PyTesseract is unavailable.

**Detailed Analysis:**

* **Primary Purpose:** This decorator checks if PyTesseract is installed and skips the decorated test case if not available.
* **Key Implementation Details:**
	+ `is_pytesseract_available()` function retrieves PyTesseract's availability status.
	+ `unittest.skipUnless` skips the test case only when `is_pytesseract_available()` returns False.
* **Notable Dependencies or Requirements:** PyTesseract is required for this decorator to work.

**Searchable Context:**

* Classes and functions mentioned in this summary:
	+ `require_pytesseract`
	+ `is_pytesseract_available`
	+ `unittest` (unittest module)
	+ `PyTesseract`
* External Dependencies:
	+ PyTesseract (library)
* Inherited or Implemented Interfaces:
	- None


key: chroma:document
string_value:
CODE:

def require_pytorch_quantization(test_case):
    """
    Decorator marking a test that requires PyTorch Quantization Toolkit. These tests are skipped when PyTorch
    Quantization Toolkit isn't installed.
    """
    return unittest.skipUnless(is_pytorch_quantization_available(), "test requires PyTorch Quantization Toolkit")(
        test_case


SUMMARY:
**Summary:**

`require_pytorch_quantization` decorator: skips test case unless `is_pytorch_quantization_available()` returns True. Test fails with error message if PyTorch Quantization Toolkit is not installed.

**Key Details:**

* Function name: `require_pytorch_quantization`
* Purpose: Conditionally skip a test case based on the availability of PyTorch Quantization Toolkit.
* Implemented by: Utilizes `unittest.skipUnless` decorator and calls virtual function `is_pytorch_quantization_available()`.
* Dependent on: Availability of PyTorch Quantization Toolkit (checked via `is_pytorch_quantization_available()`).
 

**Dependencies:**

* `unittest`: used for decorating `test_case`
* `is_pytorch_quantization_available()`: checks availability of PyTorch Quantization Toolkit


key: chroma:document
string_value:
CODE:

def require_vision(test_case):
    """
    Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't
    installed.
    """


SUMMARY:
**Summary:**

Requires decorator for test cases marked as needing vision dependencies, skipping tests unless torchaudio is available.

* Purpose: Ensures tests are skipped when torchaudio isn't installed and flags vision-related tests.
* Key implementation details:
	+ Uses `unittest.skipUnless` with the result of `is_vision_available()`.
	+ Wraps the test case in a skip instruction if `is_vision_available()` returns False.
* Dependencies:
	+ `unittest`
	+ `torchaudio`
	+ `is_vision_available()` function (not shown)


key: chroma:document
string_value:
CODE:

def require_ftfy(test_case):
    """
    Decorator marking a test that requires ftfy. These tests are skipped when ftfy isn't installed.
    """


SUMMARY:
**Concise Summary:**

`require_ftfy` Decorator:
 
- Purpose: Skips tests unless `ftfy` is installed and available.
  
- Implementation: Uses `unittest.skipUnless` to conditionally skip test cases based on the availability of `is_ftfy_available()`.

 **Key Details:**
 
- Decorator function marking test cases that require `ftfy`.
- Conditionally skips test case using `unittest.skipUnless` with a default message indicating ftfy availability.


key: chroma:document
string_value:
CODE:

def require_spacy(test_case):
    """
    Decorator marking a test that requires SpaCy. These tests are skipped when SpaCy isn't installed.
    """


SUMMARY:
**Summary**

`require_spacy` is a decorator function that marks a test as skipping only if SpaCy is not available. If SpaCy is installed, the test proceeds normally; otherwise, it is automatically skipped.

**Implementation Details**

* Uses `unittest.skipUnless` to conditionally skip tests based on `is_spacy_available()` return value.
* The test case (`test_case`) remains unchanged.
 

**Notable Dependencies/Requirements**

* `SpaCy` library (`is_spacy_available()`)
* `_unittest` module's `skipUnless` function.


key: chroma:document
string_value:
CODE:

def require_torch_multi_gpu(test_case):
    """
    Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without
    multiple GPUs.

    To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu"
    """
    if not is_torch_available():
        return unittest.skip(reason="test requires PyTorch")(test_case)

    device_count = get_device_count()



SUMMARY:
**Summary**

`require_torch_multi_gpu` is a decorator that marks tests requiring PyTorch multi-GPU setup. It skips tests if:

* PyTorch is not available
* There's only one GPU available.

**Key Details:**

- Skips tests without PyTorch or single GPU.
- Checks for multiple GPUs (`get_device_count()`) to conditionally apply the `unittest.skipUnless` decorator.
- Assumes test names contain "multi_gpu" for multi-GPU tests (using pytest).

**Notable Dependencies/Requirements:**

- `is_torch_available()` function from `torch`
- `get_device_count()` function from unknown origin
- `unittest` module

**Code Organization:**

* Decorator pattern used to mark tests.
* Uses Pythonic syntax and naming conventions.


key: chroma:document
string_value:
CODE:

def require_torch_multi_accelerator(test_case):
    """
    Decorator marking a test that requires a multi-accelerator (in PyTorch). These tests are skipped on a machine
    without multiple accelerators. To run *only* the multi_accelerator tests, assuming all test names contain
    multi_accelerator: $ pytest -sv ./tests -k "multi_accelerator"
    """
    if not is_torch_available():
        return unittest.skip(reason="test requires PyTorch")(test_case)

    return unittest.skipUnless(backend_device_count(torch_device) > 1, "test requires multiple accelerators")(
        test_case


SUMMARY:
**Function Summary: `require_torch_multi_accelerator`**

* **Purpose**: A PyTorch multi-accelerator test requirement decorator to mark tests that necessitate accelerated computing.
* 
* **Implementation Details**:
  * Checks for availability of PyTorch and skips the test if not supported (`is_torch_available()`).
  * Verifies that at least two accelerators are present in the system and skips the test otherwise (`backend_device_count`) unless an argument is provided to use an accelerator.
  
* **Notable Dependencies**:
  + `torch`: The PyTorch library required for this function to work.
  + `unittest`: Provides the necessary functionality for unit testing (used by `unittest.skip` and `unittest.skipUnless`).


key: chroma:document
string_value:
CODE:

def require_torch_non_multi_gpu(test_case):
    """
    Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip(reason="test requires PyTorch")(test_case)

    import torch



SUMMARY:
**Summary**

`require_torch_non_multi_gpu` is a unit test decorator that checks if PyTorch (with 0 or 1 GPU availability) is available and ensures the test case is only executed on devices with this configuration.

**Key Details**

* Requires PyTorch installation and functional GPU availability (either 0 or 1 device).
* Skips tests when PyTorch is not available.
* Skips tests that require more than one device using `torch.cuda.device_count() < 2`.

**Notable Dependencies/Requirements**

* PyTorch (`import torch`)
* `unittest` module


key: chroma:document
string_value:
CODE:

def require_torch_non_multi_accelerator(test_case):
    """
    Decorator marking a test that requires 0 or 1 accelerator setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip(reason="test requires PyTorch")(test_case)



SUMMARY:
**Summary:**

This decorator checks if the PyTorch library is available and verifies that the number of accelerators set up is either 0 or 1. 

* Checks for availability of PyTorch using `is_torch_available()`.
* Skips test if PyTorch is not available.
* Verifies accelerator setup condition using `backend_device_count(torch_device)`. 
    * Skips test if the number of accelerators set up is more than 1.

**Implementation Details:**

* Uses `unittest.skip()` and `unittest.skipUnless()` decorators for skipping tests.
* References PyTorch-specific functions/tensors:
	+ `is_torch_available()`
	+ `torch_device` (assumed to be a tensor-like object)
	+ `backend_device_count()`


key: chroma:document
string_value:
CODE:

def require_torch_up_to_2_gpus(test_case):
    """
    Decorator marking a test that requires 0 or 1 or 2 GPU setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip(reason="test requires PyTorch")(test_case)

    import torch



SUMMARY:
**Summary**

`require_torch_up_to_2_gpus` is a decorator for marking test cases that require PyTorch with 0, 1, or 2 GPU setup. It checks the availability of PyTorch and skips the test if it's not installed.

**Key Details:**

* **Purpose**: Marks test cases requiring specific PyTorch GPU configurations.
* **Implementation**: Checks `torch` library availability and skips tests if no GPUs are available (<= 2).
* **Notable Dependencies**: `torch`, `unittest`, `is_torch_available`.


key: chroma:document
string_value:
CODE:

def require_torch_up_to_2_accelerators(test_case):
    """
    Decorator marking a test that requires 0 or 1 or 2 accelerator setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip(reason="test requires PyTorch")(test_case)

    return unittest.skipUnless(backend_device_count(torch_device) < 3, "test requires 0 or 1 or 2 accelerators")(
        test_case


SUMMARY:
**Function Summary**

`require_torch_up_to_2_accelerators` Decorator:
Mark a test as requiring either 0 or 1 or 2 accelerator setups (PyTorch).

**Key Details:**

* Checks availability of PyTorch using `is_torch_available()`.
* Skips tests with insufficient accelerators (<3) using `unittest.skipUnless`.
* Returns the original test case if PyTorch is not available.

**Notable Dependencies/Requirements:**

* `torch_device`: returns the device associated with PyTorch.
* `backend_device_count`: counts devices in the backend (PyTorch or not).
* `is_torch_available()`: checks availability of PyTorch.


key: chroma:document
string_value:
CODE:

def require_torch_xla(test_case):
    """
    Decorator marking a test that requires TorchXLA (in PyTorch).
    """


SUMMARY:
**Summary**
-----------

Function `require_torch_xla` is a decorator marking test cases that require availability of TorchXLA (PyTorch) with skipping the test if unavailable.

**Key Details**

*   Primary purpose: Decorator to mark tests as requiring TorchXLA
*   **Decorating Pattern:** Uses `unittest.skipUnless`
*   **Function Signatures:**
    *   `is_torch_xla_available()`: Internal function checking for TorchXLA availability
    *   `unittest.skipUnless(is_torch_xla_available(),...)`: Skips test if TorchXLA is not available

**Notable Dependencies/Requirements**

*   Relied on internal `is_torch_xla_available()` function for TorchXLA check
*   Utilizes `unittest` and its skipping mechanism (`unittest.skipUnless`)


key: chroma:document
string_value:
CODE:

def require_torch_neuroncore(test_case):
    """
    Decorator marking a test that requires NeuronCore (in PyTorch).
    """
    return unittest.skipUnless(is_torch_neuroncore_available(check_device=False), "test requires PyTorch NeuronCore")(
        test_case


SUMMARY:
**Summary:**

* **Purpose**: Marks a test case as requiring PyTorch NeuronCore and automatically skips the test if it is not available on the specified device.
* **Notable Details**:
	+ Uses `unittest.skipUnless` decorator to skip tests that do not have PyTorch NeuronCore installed.
	+ Calls `is_torch_neuroncore_available` with `check_device=False` to determine availability without checking device supports.
* **Key Implementation**:
	+ The `is_torch_neuroncore_available` function is called to check if the required library is available.

This summary highlights the main functionality of the code while focusing on its purpose, key implementation details, and notable dependencies or requirements.


key: chroma:document
string_value:
CODE:

def require_torch_npu(test_case):
    """
    Decorator marking a test that requires NPU (in PyTorch).
    """


SUMMARY:
**Summary**

A decorator function `require_torch_npu` marks a test case as requiring PyTorch's NPU functionality. It skips the test unless PyTorch NPU is available.

**Key Details:**

* Function signature: `def require_torch_npu(test_case):`
* Primary purpose: Marks a test case as requiring NPU (PyTorch)
* Notable dependencies:
	+ `is_torch_npu_available()`: function to check NPU availability
	+ `unittest.skipUnless()`: testing decorator to skip tests without NPU


key: chroma:document
string_value:
CODE:

def require_torch_multi_npu(test_case):
    """
    Decorator marking a test that requires a multi-NPU setup (in PyTorch). These tests are skipped on a machine without
    multiple NPUs.

    To run *only* the multi_npu tests, assuming all test names contain multi_npu: $ pytest -sv ./tests -k "multi_npu"
    """
    if not is_torch_npu_available():
        return unittest.skip(reason="test requires PyTorch NPU")(test_case)



SUMMARY:
**Summary**

`require_torch_multi_npu` Decorator:

Mark test cases that require PyTorch's multi-NPU setup.

* Checks if NPU support is available (PyTorch) and skips if not.
* Skips tests with only one NPU or no NPUs assigned.
* Can be used to filter and run tests that depend on multiple NPUs.

**Key Implementation Details**

* Uses `unittest.skip` for skipping tests without NPU support.
* Uses `torch.npu.device_count()` to check number of NPUs available.
* Passes test case with reason when checking if NPU setup meets requirements.


key: chroma:document
string_value:
CODE:

def require_torch_xpu(test_case):
    """
    Decorator marking a test that requires XPU (in PyTorch).

    These tests are skipped when XPU backend is not available. XPU backend might be available either via stock
    PyTorch (>=2.4) or via Intel Extension for PyTorch. In the latter case, if IPEX is installed, its version
    must match match current PyTorch version.
    """


SUMMARY:
**Summary**

This function `require_torch_xpu` is a decorator that marks tests as skippable unless the PyTorch XPU (Accelerated Compute Unit) backend is available. It checks for XPU availability using the `is_torch_xpu_available()` function and conditionally applies the default configuration to the test case if XPU is unavailable.

**Key Implement Details**

* The decorator uses `unittest.skipUnless` to mark tests as skippable.
* It relies on the `is_torch_xpu_available()` function to check for XPU availability.

**Notable Dependencies/Requirements**

* PyTorch.
* `unittest`.
* `is_torch_xpu_available()`.


key: chroma:document
string_value:
CODE:

def require_non_xpu(test_case):
    """
    Decorator marking a test that should be skipped for XPU.
    """


SUMMARY:
**Summary**

`require_non_xpu` is a decorator function that skips a `unittest.testCase` instance if running on an XPU device. It checks the value of `torch_device` and applies the `unittest.skipUnless` marker with a custom failure message.

**Key Details**

* Decorator function applies `unittest.skipUnless` to a test case
* Checks `torchDevice != "xpu"` condition for skipping
* Custom failure message: `"test requires a non-XPU"`

**No notable dependencies or requirements beyond the given code.**


key: chroma:document
string_value:
CODE:

def require_torch_multi_xpu(test_case):
    """
    Decorator marking a test that requires a multi-XPU setup (in PyTorch). These tests are skipped on a machine without
    multiple XPUs.

    To run *only* the multi_xpu tests, assuming all test names contain multi_xpu: $ pytest -sv ./tests -k "multi_xpu"
    """
    if not is_torch_xpu_available():
        return unittest.skip(reason="test requires PyTorch XPU")(test_case)



SUMMARY:
**Summary:**

`require_torch_multi_xpu` Decorator:

* Marks a test that relies on PyTorch's multi-XPU feature
* Skips tests without multiple available XPUs or with invalid device count
* Run all multi-xpu tests using `$ pytest -sv ./tests -k "multi_xpu"`

**Key Implementation Details:**

* Checks availability of PyTorch XPU and ensures proper device count (> 1)
* Uses `unittest.skip` and `skipsUnless` to skip or run tests based on conditions

Notable Dependencies/Requirements:

* `torch.xpu`
* `unittest`
* `pytest`


key: chroma:document
string_value:
CODE:

def require_torchdynamo(test_case):
    """Decorator marking a test that requires TorchDynamo"""


SUMMARY:
**Summary:**

**Function Purpose and Functionality:** `require_torchdynamo` is a decorator for marking test cases that require TorchDynamo.

**Key Implementation Details:**

* Emulates `unittest.skipUnless` to skip tests unless TorchDynamo is available
* Uses `is_torchdynamo_available()` to determine availability

**Notable Dependencies or Requirements:**

* `is_torchdynamo_available()`: a function checking whether TorchDynamo is installed and accessible


key: chroma:document
string_value:
CODE:

def require_torchao(test_case):
    """Decorator marking a test that requires torchao"""


SUMMARY:
**Code Summary**

`require_torchao` is a decorator that marks a test case as requiring `torchao`. It wraps the `unittest.skipUnless` context manager with the condition to check if `torchao` is available. If unavailable, the test will be skipped.

**Key Details:**

* Decorator marking test as requiring `torchao`
* Uses `unittest.skipUnless` to skip test if `torchao` is not available
* Condition checked by the `is_torchao_available()` function


key: chroma:document
string_value:
CODE:

def require_torch_tensorrt_fx(test_case):
    """Decorator marking a test that requires Torch-TensorRT FX"""


SUMMARY:
**Summary:**

`require_torch_tensorrt_fx` Decorator
=====================================


Marks a test as requiring Torch-TensorRT FX functionality. If Torch-TensorRT FX is not available, the test is skipped.

**Key Implementation Details:**

*   Uses `unittest.skipUnless` to ensure the test only runs if Torch-TensorRT FX is enabled.
*   Calls `is_torch_tensorrt_fx_available()` to check for availability of Torch-TensorRT FX functionality.

**Notable Dependencies/Requirements:**

*   `torch`
*   `torch.tortensortfx` ( Torch-TensorRT FX module)
*   `unittest` framework

Note: The actual import is not shown in this snippet, but it's assumed to be part of the larger test framework.


key: chroma:document
string_value:
CODE:

def require_torch_gpu(test_case):
    """Decorator marking a test that requires CUDA and PyTorch."""


SUMMARY:
**Summary:**

`require_torch_gpu` Decorator:

* A decorator function for marking tests that require CUDA and PyTorch support.

**Key Details:**

* **Function Signature:** `require_torch_gpu(test_case)`
* **Purpose:** Checks if the test case is running on a device other than CUDA. If not, skips the test.
* **Imported Dependencies:**
	+ `unittest`
	+ `torch_device` (assumed to be a variable)
	+ `torch` (imported implicitly)
* **Notable Interface:** Uses the `unittest.skipUnless` function with a custom error message.

**Class/Variable Name:** `test_case`

Note that this summary focuses on the key aspects of the code, omitting unnecessary details such as the exact implementation or surrounding context.


key: chroma:document
string_value:
CODE:

def require_torch_gpu_if_bnb_not_multi_backend_enabled(test_case):
    """
    Decorator marking a test that requires a GPU if bitsandbytes multi-backend feature is not enabled.
    """
    if is_bitsandbytes_available() and is_bitsandbytes_multi_backend_available():
        return test_case


SUMMARY:
**Summary:**

`require_torch_gpu_if_bnb_not_multi_backend_enabled` Decorator:

* Requires GPU if `bitsandbytes` multi-backend feature is not enabled.

**Key Implementation Details:**

* Checks availability of `bitsandbytes` and its multi-backend functionality
* Returns the test case unchanged if both are available
* Otherwise, calls `require_torch_gpu` with the test case

**Notable Dependencies/Requirements:**

* `is_bitsandbytes_available()`
* `is_bitsandbytes_multi_backend_available()`
* `require_torch_gpu(test_case)`


key: chroma:document
string_value:
CODE:

def require_torch_accelerator(test_case):
    """Decorator marking a test that requires an accessible accelerator and PyTorch."""
    return unittest.skipUnless(torch_device is not None and torch_device != "cpu", "test requires accelerator")(
        test_case


SUMMARY:
**Summary:**

```
Decorates a unit test to run only if PyTorch is executed on an accelerator device (not CPU) using the `torch_device` variable.
```   

**Key Implementation Details:**

* Utilizes the `unittest.skipUnless` decorator to skip tests without a valid `torch_device`.
* Requires access to both `torch.device` and its value, which must be non-`cpu`.

**Notable Dependencies/Requirements:**

* `torch`: PyTorch library.
* `unittest`: Unit testing framework.
```


key: chroma:document
string_value:
CODE:

def require_torch_fp16(test_case):
    """Decorator marking a test that requires a device that supports fp16"""
    return unittest.skipUnless(
        is_torch_fp16_available_on_device(torch_device), "test requires device with fp16 support"


SUMMARY:
**Summary:**

`require_torch_fp16` Decorator
=====================================

This decorator marks a unit test for execution only if the specified device supports Floating-Point 16 (fp16). It checks if `torch_device` is compatible with fp16 using the `is_torch_fp16_available_on_device` function.

**Implementation Details:**

* Requires torch 1.4+ due to its implementation of fp16 arithmetic.
* Utilizes the `unittest.skipUnless` context manager to execute the test only when the device supports fp16.

**Dependencies and Requirements:**

* torch (version 1.4+) or later, due to support for fp16 precision.


key: chroma:document
string_value:
CODE:

def require_fp8(test_case):
    """Decorator marking a test that requires supports for fp8"""
    return unittest.skipUnless(is_accelerate_available() and is_fp8_available(), "test requires fp8 support")(
        test_case


SUMMARY:
**Summary**

`require_fp8` Decorator:
------------------------

*   Purpose: Marks tests that rely on FP8 (16-bit floating-point) support as skipable if required components are not available.

**Implementation Details**
---------------------------

*   Uses `unittest.skipUnless` to conditionally skip the test case unless two prerequisites are met:
    *   `is_accelerate_available()`
    *   `is_fp8_available()`

These dependencies check for the availability of corresponding features in the system.


key: chroma:document
string_value:
CODE:

def require_torch_bf16(test_case):
    """Decorator marking a test that requires a device that supports bf16"""
    return unittest.skipUnless(
        is_torch_bf16_available_on_device(torch_device), "test requires device with bf16 support"


SUMMARY:
**Summary**

`require_torch_bf16` Decorator: Marks a test as requiring a device that supports Torch BF16. The decorator skips the test unless the required device is available and supports TF32/F16.

**Key Details**

* Test case wrapper using `unittest.skipUnless`
* Checks availability of the required device with `is_torch_bf16_available_on_device(torch_device)`
* Requires Torch BF16 support on device

**External Dependencies**

* `torch_device`: Represents the specific tensor device (e.g., GPU, CPU) being used.
* `is_torch_bf16_available_on_device`: Function checks if the specified device has BF16 support.


key: chroma:document
string_value:
CODE:

def require_torch_bf16_gpu(test_case):
    """Decorator marking a test that requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0"""
    return unittest.skipUnless(
        is_torch_bf16_gpu_available(),
        "test requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0",


SUMMARY:
**Function Summary:**

`require_torch_bf16_gpu` is a decorator that skips the test case if PyTorch (torch) version 1.10 or higher is not installed, and it is running on an Ampere GPU or a newer architecture with CUDA version 11.0 or higher.

**Key Details:**

* Decorator function
* Requires PyTorch >= 1.10
* Checks for Ampere GPU or newer architecture with CUDA >= 11.0

**Notable Dependencies and Requirements:**

* torch (PyTorch)
	+ Version >= 1.10
* cuda (CUDA)
	+ Version >= 11.0


key: chroma:document
string_value:
CODE:

def require_torch_bf16_cpu(test_case):
    """Decorator marking a test that requires torch>=1.10, using CPU."""
    return unittest.skipUnless(
        is_torch_bf16_cpu_available(),
        "test requires torch>=1.10, using CPU",


SUMMARY:
**Function Summary:**

`require_torch_bf16_cpu` Decorator:
---------------------------

This function is used as a decorator to mark tests that require PyTorch version 1.10 or higher (`torch>=1.10`) and run on CPU.

**Key Details:**

* Checks availability of `bf16` support in PyTorch using `is_torch_bf16_cpu_available()`.
* Skips test if `bf16` is not available, with a warning message indicating the required PyTorch version.

**Notable Dependencies:**

* `unittest`
* `is_torch_bf16_cpu_available()` (referenced class/function not shown in this snippet)

Note: The implementation of `is_torch_bf16_cpu_available()` is not provided in this code chunk.


key: chroma:document
string_value:
CODE:

def require_deterministic_for_xpu(test_case):
    if is_torch_xpu_available():
        return unittest.skipUnless(is_torch_deterministic(), "test requires torch to use deterministic algorithms")(
            test_case
        )
    else:


SUMMARY:
**Summary:**

`require_deterministic_for_xpu` function:
 
- Skips or runs a test case based on availability of PyTorch XPU engine.
- Returns a `unittest.skipUnless` wrapper for tests that require PyTorch deterministic algorithms.

**Key Details:**

- Conditionally checks for PyTorch XPU engine availability using `is_torch_xpu_available()`.
- Skips the test (`return unittest.skipUnless(...)`) if both PyTorch XPU and deterministic algorithms are available.
- Otherwise, runs the test unchanged (`return test_case`).

**Notable Dependencies:**

- `unittest`
- `torch.xpu`
    - `is_torch_xpu_available()`
    - `is_torch_deterministic()`


key: chroma:document
string_value:
CODE:

def require_torch_tf32(test_case):
    """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7."""
    return unittest.skipUnless(
        is_torch_tf32_available(), "test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7"


SUMMARY:
**Summary**

`require_torch_tf32` Decorator:

* A test case decorator marking the test as skipping if ` Torch-TF32` framework is not available.
* Requires:
	+ GPUs with Ampere or newer arch (torch>=1.7).
	+ CUDA 11 or higher (`cuda>=11`).
* Conditionally skips test execution using `unittest.skipUnless`.

**Implementation Details**

* Invokes the `is_torch_tf32_available()` function to check availability.
* Uses `unittest.skipunless` with a custom condition to determine whether to skip the test.

Note: Exactly preserved class and method names (`test_case`, `unittest`) and references to dependent methods (`is_torch_tf32_available()`).


key: chroma:document
string_value:
CODE:

def require_detectron2(test_case):
    """Decorator marking a test that requires detectron2."""


SUMMARY:
**Summary:**
 
This function, `require_detectron2`, is a decorator used to mark a test as requiring Detectron 2. It skips the test unless `is_detectron2_available()` returns True.

**Key Details:**

- Decorator marking tests that require detectron2
- Skips test if detectron2 is not available

**Notable Dependencies/Requirements:**

- `is_detectron2_available()`
- `unittest`


key: chroma:document
string_value:
CODE:

def require_faiss(test_case):
    """Decorator marking a test that requires faiss."""


SUMMARY:
**Function Summary: `require_faiss`**

This function acts as a decorator to mark a test case that requires the Faiss library to be available.

**Implementation Details:**

- Utilizes `unittest.skipUnless` to skip testing if Faiss is not available and displays an error message.
- Calls the `is_faiss_available()` function before testing, indicating its reliance on the result of this external check.

**Key Dependencies/Requirements:**

- Relies on `faiss` being installed and available for testing.


key: chroma:document
string_value:
CODE:

def require_optuna(test_case):
    """
    Decorator marking a test that requires optuna.

    These tests are skipped when optuna isn't installed.

    """


SUMMARY:
**Summary:**

`require_optuna` Decorator
=============================

**Purpose and Functionality**

Marks a test as requiring Optuna, skipping it unless Optuna is installed.

**Key Implementation Details**

* Skips test using `unittest.skipUnless` if Optuna is not available.
* Uses the `is_optuna_available()` function to check for Optuna installation.

**Notable Dependencies or Requirements**

* `optuna`: dependency required for Optunas availability check
* `unittest`: testing framework used by the decorator.

**Method Signatures and Interfaces**

None explicitly defined, relying on custom `is_optuna_available()` function.


key: chroma:document
string_value:
CODE:

def require_ray(test_case):
    """
    Decorator marking a test that requires Ray/tune.

    These tests are skipped when Ray/tune isn't installed.

    """


SUMMARY:
**Summary**

 decorate function `require_ray`:

*   Marks test cases that require `Ray/tune`.
*   Skips tests when `is_ray_available()` returns False.

**Key Implementation Details**

*   Uses `unittest.skipUnless` decorator to conditionally skip tests based on availability of `Ray/tune`.
*   Calls `is_ray_available()` function to determine the availability of `Ray/tune`.

**Notable Dependencies or Requirements**

*   Utilizes `unittest` module for testing framework support.
*   Relies on `is_ray_available()` function to check if `Ray/tune` is installed.


key: chroma:document
string_value:
CODE:

def require_sigopt(test_case):
    """
    Decorator marking a test that requires SigOpt.

    These tests are skipped when SigOpt isn't installed.

    """


SUMMARY:
**Summary:**

`require_sigopt` is a decorator that marks a test case as skipping until SigOpt availability. It checks if SigOpt is installed and uses the `unittest.skipUnless` function to skip tests unless SigOpt is available.

**Key Details:**

* Primary purpose: Skip test cases when SigOpt isn't installed.
* Implementation:
	+ Checks if SigOpt is available using the `is_sigopt_available()` function.
	+ Uses `unittest.skipUnless` with a custom error message ("test requires SigOpt").
	+ Wraps the decorated test case in this context manager.

**Notable Dependencies:**

* `unittest`: used for skipping tests
* `is_sigopt_available()`: function checking SigOpt availability (not shown in this code snippet)

Note: This summary preserves the function's name, function signature, and main purpose, while extracting key implementation details and dependencies.


key: chroma:document
string_value:
CODE:

def require_wandb(test_case):
    """
    Decorator marking a test that requires wandb.

    These tests are skipped when wandb isn't installed.

    """


SUMMARY:
**Function Summary**

`require_wandb` Decorator
==========================

Mark test case for execution with WandB logging enabled. Skips tests without WandB installation.

### Key Implementation Details

* Uses `unittest.skipUnless` to skip tests without WandB availability
* Invokes `is_wandb_available()` function to check WandB installation status

### Notable Dependencies and Requirements

* `unittest`: Unit testing framework
* `is_wandb_available()`: Function checking WandB installation status


key: chroma:document
string_value:
CODE:

def require_clearml(test_case):
    """
    Decorator marking a test requires clearml.

    These tests are skipped when clearml isn't installed.

    """


SUMMARY:
**Summary**

`require_clearml` Decorator:

* Marks a test case for skipping if ClearML is not installed.
* Uses `unittest.skipUnless` with the result of `is_clearml_available()` to determine whether to skip the test.

**Key Details**

* Requires ClearML installation (`is_clearml_available()`)
* Skips test when ClearML not available
* Utilizes `unittest.skipUnless` decorator


key: chroma:document
string_value:
CODE:

def require_soundfile(test_case):
    """
    Decorator marking a test that requires soundfile

    These tests are skipped when soundfile isn't installed.

    """


SUMMARY:
**Summary: `require_soundfile` Decorator**

### Purpose and Functionality
The `require_soundfile` decorator marks a test case as requiring the `soundfile` library. If `soundfile` is not installed or available, the decorated test case is skipped.

### Key Implementation Details

* Uses `unittest.skipUnless` to skip tests when `soundfile` is unavailable.
* Calls `is_soundfile_available()` to check availability of `soundfile`.
* Returns a decorated test case that skips if `soundfile` is not available.

### Notable Dependencies or Requirements
* `soundfile` library (imported by calling `is_soundfile_available()`)


key: chroma:document
string_value:
CODE:

def require_deepspeed(test_case):
    """
    Decorator marking a test that requires deepspeed
    """


SUMMARY:
**Summary:**

`require_deepspeed` is a test case decorator that ensures a test only runs if DeepSpeed is available. It uses the `unittest.skipUnless` context manager to skip the test if DeepSpeed is not present.

**Key Details:**

* Ensures test suite checks for DeepSpeed availability before execution.
* Uses `is_deepspeed_available()` function to determine availability.
* Deciphers as a wrapper around `skipUnless` to conditionally enable or disable individual test cases.


key: chroma:document
string_value:
CODE:

def require_apex(test_case):
    """
    Decorator marking a test that requires apex
    """


SUMMARY:
**Summary:**

`require_apex`: A decorator function that marks a test case as requiring Apex, skipping the test unless Apex is available.

**Key Details:**

* Utilizes `unittest.skipUnless` to skip the test if Apex is not available.
* Relies on the `is_apex_available()` function to determine Apex availability.

**Notable Dependencies:**

* `unittest`
* `is_apex_available()`: A separate function indicating whether Apex is available.


key: chroma:document
string_value:
CODE:

def require_aqlm(test_case):
    """
    Decorator marking a test that requires aqlm
    """


SUMMARY:
**Summary**

`require_aqlm`: A decorator function marking a test case as requiring `aqlm` (Advanced Query Language Macro) functionality. It skips the test if `aqlm` is not available.

**Key Details:**

* **Function Signature:** `def require_aqlm(test_case):`
* **Primary Purpose:** To mark necessary dependencies
* **Implementation Notes:** Uses the `unittest.skipUnless` decorator to conditionally skip tests based on the `is_aqlm_available()` function.
* **Notable Dependencies:** None specified, but relies on `aqlm`.
* **Related Interfaces:** None.


key: chroma:document
string_value:
CODE:

def require_vptq(test_case):
    """
    Decorator marking a test that requires vptq
    """


SUMMARY:
**Function Summary**

`require_vptq` is a decorator that marks a test case as requiring VPTQ (Version-Patient-Target-Quota) functionality. If VPTQ is not available, the decorated test case will be skipped using `unittest.skipUnless`.

**Key Implementation Details:**

* `is_vptq_available()` function is called to determine whether VPTQ is available
* The `unittest.skipUnless` context manager is used to skip the test case if VPTQ is not available
* Decorated test case is returned

**Notable Dependencies:**

* `unittest` module (built-in Python library)
* `is_vptq_available()` function (unknown implementation, possibly defined elsewhere in the codebase)


key: chroma:document
string_value:
CODE:

def require_eetq(test_case):
    """
    Decorator marking a test that requires eetq
    """
    eetq_available = is_eetq_available()
    if eetq_available:
        try:
            import eetq  # noqa: F401
        except ImportError as exc:
            if "shard_checkpoint" in str(exc):
                # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed
                # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34.
                # TODO: Remove once eetq releases a fix and this release is used in CI
                eetq_available = False


SUMMARY:
**Summary**

The `require_eetq` function is a decorator that marks test cases requiring EETQ (Elastic Enriched Test Query). It checks availability of `is_eetq_available()` and conditional imports or overrides based on exceptions thrown. The decorated test case is skipped unless EETQ is available, with a corresponding error message.

**Key Implementation Details**

* **Decoration**: Marks a test case to be skipped if EETQ is not available.
* **Dependency Check**: `is_eetq_available()` function checks EETQ availability.
* **Conditional Import**: Attempts to import `eetq` module; overrides availability based on exceptions thrown.

**Notable Dependencies or Requirements**

* `is_eetq_available()`: Checks if EETQ is available.
* `eetq`: Import attempt for EETQ module.


key: chroma:document
string_value:
CODE:

def require_av(test_case):
    """
    Decorator marking a test that requires av
    """


SUMMARY:
**Summary:**

* Purpose: Decorates aunittest test case to skip unless AV is available.
* Functionality: Uses `unittest.skipUnless` and calls `is_av_available()` to determine the required availability of AV.
* Key details:
	+ Calls `is_av_available()`
	+ Skips test if AV is not available
	+ Raises an error message "test requires av" if AV is not available

**Notable dependencies:**
* `unittest`
* `is_av_available()`


key: chroma:document
string_value:
CODE:

def require_bitsandbytes(test_case):
    """
    Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed.
    """
    if is_bitsandbytes_available() and is_torch_available():
        try:
            import pytest

            return pytest.mark.bitsandbytes(test_case)
        except ImportError:
            return test_case
    else:


SUMMARY:
**Function Summary:**

`require_bitsandbytes`
========================

**Purpose:** Marks a test case as requiring the `bitsandbytes` library. Skips tests when the library is not installed.

**Implementation Details:**

*   Uses a decorator to conditionally apply the `pytest.mark.bitsandbytes` marker.
*   Verifies availability of both `bitsandbytes` and its hard dependency, `torch`, before attempting import.
*   Returns the original test case in two scenarios:
    *   When the required libraries are available.
    *   Otherwise, skips the test with a provided reason.

**Key Components:**

*   `is_bitsandbytes_available()` and `is_torch_available()` checks for library availability.
*   `pytest.mark.bitsandbytes` decorator marker application.
*   `unittest.skip` method usage for skipping tests.

**Example Usage:**
```python
@require_bitsandbytes('test_case')
def my_test():
    # test code here
    pass
```
In this example, the `my_test` function will be skipped if the `bitsandbytes` library is not installed or its hard dependency (`torch`) is missing.


key: chroma:document
string_value:
CODE:

def require_optimum(test_case):
    """
    Decorator for optimum dependency
    """


SUMMARY:
**Summary:**

`require_optimum` is a decorator that skips test cases unless `isOptimumAvailable()` returns True. It uses the `unittest.skipUnless` context manager to conditionally skip tests.

**Key Details:**

* Purpose: Conditionally skip non-optimal test cases
* Implementation:
	+ Calls `isOptimumAvailable()`
	+ Uses `unittest.skipUnless` to skip non-optimal tests if `is_optimum_available()` returns False
	+ Returns decorated `test_case`

**Notable Dependencies/Interfaces:**

* `unittest`
* `is_optimum_available()`: a function within the same module being analyzed


key: chroma:document
string_value:
CODE:

def require_tensorboard(test_case):
    """
    Decorator for `tensorboard` dependency
    """


SUMMARY:
**Summary**

Decorates a test case to skip unless the `tensorboard` dependency is available.

**Key Implementation Details**

* Applies `unittest.skipUnless` with condition checking if `is_tensorboard_available()` returns True
* Throws an error message "test requires tensorboard" otherwise

**Notable Dependencies or Requirements**

* `is_tensorboard_available()` function (interface)
* External dependency on the `tensorboard` library not explicitly managed within this snippet, but suggested by context.


key: chroma:document
string_value:
CODE:

def require_auto_gptq(test_case):
    """
    Decorator for auto_gptq dependency
    """


SUMMARY:
**Summary**

* **Purpose**: Applies the `unittest.skipUnless` decorator to skip test cases unless the `is_auto_gptq_available()` function returns True, indicating the presence of an auto-GPTQ dependency.
* **Key Implementation Details**:
	+ Uses `unittest.skipUnless` with a custom failure message referencing `is_auto_gptq_available()`.
* **Dependencies/Requirements**: `unittest`, `is_auto_gptq_available()` (assuming a separate function definition)


key: chroma:document
string_value:
CODE:

def require_auto_awq(test_case):
    """
    Decorator for auto_awq dependency
    """


SUMMARY:
**Summary**

`require_auto_awq` is a decorator factory that conditionally skips tests using `unittest.skipUnless` if the `auto_awq` dependency is not available.

**Key Implementation Details**

* Conditionally skips tests based on availability of `auto_awq`
* Uses `is_auto_awq_available()` as the criterion
* Invokes `unittest.skipUnless` with the result

**Notable Dependencies or Requirements**

* `unit_test`: implicitly referenced through `unittest`

No external dependencies are listed.


key: chroma:document
string_value:
CODE:

def require_optimum_quanto(test_case):
    """
    Decorator for quanto dependency
    """


SUMMARY:
**Code Summary:**

**Decorator for Checking Quanto Dependency in Unit Test Cases**

* Purpose: Verifies if the "quanto" dependency is available before running unit test cases using the `unittest` framework.
* Primary Functionality: Skips unit tests unless the "is_optimum_quanto_available()" function returns True.

**Implementation Details:**

* Uses a decorator to wrap the original unit test case in a conditional skip.
* Calls the `is_optimum_quanto_available()` function to check for availability of the "quanto" dependency.


key: chroma:document
string_value:
CODE:

def require_compressed_tensors(test_case):
    """
    Decorator for compressed_tensors dependency
    """


SUMMARY:
**Function Summary: `require_compressed_tensors`**

* **Purpose:** A decorator function that skips unit test execution if `compressed_tensors` is not available.

**Key Implementation Details:**
* Uses a decorator pattern with `unittest.skipUnless`.
* Checks availability of `is_compressed_tensors_available()` function as a requirement.
* Passed `test_case` object to be skipped or executed based on the result.

**Dependencies and Requirements:**
* Requires `compressed_tensors` dependency (`is_compressed_tensors_available()` function).
* Inherently depends on `unittest` framework.


key: chroma:document
string_value:
CODE:

def require_fbgemm_gpu(test_case):
    """
    Decorator for fbgemm_gpu dependency
    """


SUMMARY:
**Summary**

`require_fbgemm_gpu` is a decorator function that skips test execution unless the `fbgemm_gpu` dependency is available. It calls `is_fbgemm_gpu_available()` to check for availability and uses the result to conditionally skip the test.

**Key Details:**

* Decorates test cases with skipping behavior
* Checks availability of `fbgemm_gpu`
* Conditionally skips tests based on availability

**Notable Dependencies/Requirements:**

* `unittest` (testing framework)
* `is_fbgemm_gpu_available()` function
* `fbgemm-gpu` dependency


key: chroma:document
string_value:
CODE:

def require_flute_hadamard(test_case):
    """
    Decorator marking a test that requires higgs and hadamard
    """
    return unittest.skipUnless(
        is_flute_available() and is_hadamard_available(), "test requires flute and fast_hadamard_transform"


SUMMARY:
**Summary:**

`require_flute_hadamard` Decorator:

* Marks a test case requiring Higgs and Hadamard availability (using `is_flute_available()` and `is_hadamard_available()` checks).
* Skips the test unless both conditions are met, with an error message indicating required dependencies (`flute` and `fast_hadamard_transform`).

**Key Details:**

* Decorator function for marking test cases.
* Uses `unittest.skipUnless` to conditionally skip tests based on availability checks.
* Calls `is_flute_available()` and `is_hadamard_available()` functions (details not provided in this snippet).


key: chroma:document
string_value:
CODE:

def require_phonemizer(test_case):
    """
    Decorator marking a test that requires phonemizer
    """


SUMMARY:
**Summary:**

`require_phonemizer` is a decorator function that skips a test unless `is_phonemizer_available()` returns `True`, ensuring the phonemizer is available for tests.

**Key Implementation Details:**

* Decorator function `require_phonemizer` uses `unittest.skipUnless()` to skip or run tests based on the availability of `phonemizer`.
* Checks if `is_phonemizer_available()` method returns `True` before executing the test.

**Notable Dependencies or Requirements:**

* `unittest`
* `is_phonemizer_available()`, a related function (not shown in this code snippet)


key: chroma:document
string_value:
CODE:

def require_pyctcdecode(test_case):
    """
    Decorator marking a test that requires pyctcdecode
    """


SUMMARY:
**Function Summary**

`require_pyctcdecode` is a decorator function used to mark tests that require the `pyctcdecode` library. It skips the test if `pyctcdecode` is not available.

**Key Implementation Details**

- **Decorator Function**: Marks a test case (`test_case`) as required for `pyctcdecode` using `unittest.skipUnless`.
- **Availability Check**: Validates `pyctcdecode` availability with `is_pyctcdecode_available()`.

**Notable Dependencies/Requirements**

- `pyctcdecode`: The Python implementation of CTC (Concurrent-Time-Correcting) decoding
- `unittest`: A built-in Python testing framework used for running the test case

**Interfaces Implementations**

- None explicitly mentioned, but implies compatibility with `unittest` framework.


key: chroma:document
string_value:
CODE:

def require_librosa(test_case):
    """
    Decorator marking a test that requires librosa
    """


SUMMARY:
**Summary**

```
[librosa Availability Check Decorator]
 decorate a test with `unittest.skipUnless`,
 ensuring it's skipped if librosa is not available.

Requires `is_librosa_available()` function.
```

**Breakdown**

* Main purpose: To mark tests that require the presence of librosa
* Key implementation detail: Uses `unittest.skipUnless` to conditionally skip a test based on librosa availability
* Notable dependency: `librosa` library, checked using the `is_librosa_available()` function


key: chroma:document
string_value:
CODE:

def require_liger_kernel(test_case):
    """
    Decorator marking a test that requires liger_kernel
    """


SUMMARY:
**Summary:**

Decorates a test case to skip execution unless the `liger_kernel` is available. Utilizes the `unittest.skipUnless` decorator with a custom test condition evaluated by `is_liger_kernel_available()`.

**Key Details:**

* Primary purpose: Decorate test cases that require `liger_kernel`.
* Implementation details:
	+ Uses `unittest.skipUnless` decorator.
	+ Custom test condition is checked using `is_liger_kernel_available()`.


key: chroma:document
string_value:
CODE:

def require_essentia(test_case):
    """
    Decorator marking a test that requires essentia
    """


SUMMARY:
**Summary:**

`require_essentia` is a decorator function marking tests as requiring the Essentia library. It skips the test unless `is_essentia_available()` indicates that Essentia is available, and wraps the test case with `unittest.skipUnless` to achieve this.

**Key Details:**

* Marks tests as requiring Essentia using decorator syntax
* Skips the test if Essentia is not available (via `is_essentia_available()`)
* Uses `unittest skippedUnless` to conditionally skip the test

No notable dependencies or requirements are listed in this code snippet.


key: chroma:document
string_value:
CODE:

def require_pretty_midi(test_case):
    """
    Decorator marking a test that requires pretty_midi
    """


SUMMARY:
**Summary:**

This function is a decorator for unit tests (`unittest`) that checks if the `pretty_midi` library is available. If it's not, the test is skipped using `unittest.skipUnless`.

* **Main Purpose:** Enforces the usage of `pretty_midi` in individual test cases.
* **Key Implementation Details:**
	+ Uses `@unittest.skipUnless` decorator to skip tests without `pretty_midi`.
	+ Calls helper function `is_pretty_midi_available()` to check library availability.
* **Notable Dependencies/Requirements:** `pretty_midi`, `unittest`.


key: chroma:document
string_value:
CODE:

def cmd_exists(cmd):


SUMMARY:
**Summary:**

### Function Purpose

The `cmd_exist` (corrected to match the provided function name) function checks if a specific command (`cmd`) exists in the system.

### Primary Purpose

Uses `shutil.which` to verify the presence of an executable file or command in the current environment, returning `True` if found and `False` otherwise.

### Key Implementation Details:

- Utilizes `shutil.which(cmd)`: searches for the specified command.
- Returns a boolean value (`is not None`) indicating existence.

### Notable Dependencies/Requirements:

- Requires `shutil`, a built-in Python module for interacting with the file system.

### Interface and Interfaces:

None specified, but inherits from standard Python functions and types.


key: chroma:document
string_value:
CODE:

def require_usr_bin_time(test_case):
    """
    Decorator marking a test that requires `/usr/bin/time`
    """


SUMMARY:
**Function Summary**

`require_usr_bin_time`: A decorator function used to mark tests that require running with `/usr/bin/time`. It checks if the command exists and skips the test otherwise.

**Key Implementation Details**

* Uses `unittest.skipUnless` to conditionally skip tests based on command existence.
* References external dependency: `cmd_exists("/usr/bin/time")`.

**Notable Dependencies**

* No explicit dependencies, but assumes existence of `cmd_exists()` function.


key: chroma:document
string_value:
CODE:

def require_sudachi(test_case):
    """
    Decorator marking a test that requires sudachi
    """


SUMMARY:
**Summary:**

Function `require_sudachi` is a decorator that skips a test unless the `is_sudachi_available()` function returns True.

**Key Details:**

* Checks availability of `sudachi` functionality using `is_sudachi_available()`
* Skips test if sudachi is not available, otherwise runs with the provided test case (`test_case`)
* Returns a decorated test case that can be used in unit testing (using `unittest.skipUnless` method)


key: chroma:document
string_value:
CODE:

def require_sudachi_projection(test_case):
    """
    Decorator marking a test that requires sudachi_projection
    """
    return unittest.skipUnless(is_sudachi_projection_available(), "test requires sudachi which supports projection")(
        test_case


SUMMARY:
**Code Summary:**

`require_sudachi_projection` Decorator

**Purpose:** Marks test cases that require sudachi_projection library, skipping tests if not available.

**Functionality:**

* Takes a unit test case as input `test_case`
* Returns a decorated test case using `unittest.skipUnless`, conditional on availability of `is_sudachi_projection_available()`
* Raises an exception if `sudachi_projection` is not supported

**Key Implementation Details:**

* Uses `unittest.skipUnless` to conditionally skip tests
* Calls `is_sudachi_projection_available()` method, which is assumed to be a separate function defining the library availability check.

**Notable Dependencies and Requirements:**

* `unittest`
* `sudachi_projection` library (checked for availability using an external method `is_sudachi_projection_available()`)


key: chroma:document
string_value:
CODE:

def require_jumanpp(test_case):
    """
    Decorator marking a test that requires jumanpp
    """


SUMMARY:
**Function Summary**

`require_jumanpp` is a test decorator that skips a test unless `jumanpp` is available. It uses the `unittest` framework and relies on the existence of a function `is_jumanpp_available()` to determine availability.

**Key Implementation Details**

* Skips a test if `jumanpp` is not available
* Uses the `unittest.skipUnless` context manager to achieve this

**Notable Dependencies/Requirements**

* `unittest`
* `is_jumanpp_available() function` (exact implementation not shown)


key: chroma:document
string_value:
CODE:

def require_cython(test_case):
    """
    Decorator marking a test that requires jumanpp
    """


SUMMARY:
**Summary:**

This function, `require_cython`, is a decorator used to mark tests that require the Cython compiler. It skips the test if Cython is not available and displays an error message.

**Key Implementation Details:**

- Uses the `unittest.skipUnless` context manager to skip the test if Cython is not available.
- Calls the `is_cython_available()` function, which is not shown in this snippet.

**Notable Dependencies/Requirements:**

- Utilizes the `unittest` and `is_cython_available()` modules (not shown).
- Expects the availability of Cython compiler.


key: chroma:document
string_value:
CODE:

def require_tiktoken(test_case):
    """
    Decorator marking a test that requires TikToken. These tests are skipped when TikToken isn't installed.
    """


SUMMARY:
**Function Summary**

`require_tiktoken` Decorator
==========================

Marks a test case as requiring TikToken. Tests with this decorator are skipped if TikToken is not installed.

Functions:
* `is_tiktoken_available()`: checks if TikToken is installed (not shown in this code snippet)
* `unittest.skipUnless()`: skips the test if `is_tiktoken_available()` returns False

Key Implementation Details:

* Uses a decorator to mark tests that require TikToken
* Skips tests if TikToken is not available due to `unittest.skipUnless`


key: chroma:document
string_value:
CODE:

def get_gpu_count():
    """
    Return the number of available gpus (regardless of whether torch, tf or jax is used)
    """
    if is_torch_available():
        import torch

        return torch.cuda.device_count()
    elif is_tf_available():
        import tensorflow as tf

        return len(tf.config.list_physical_devices("GPU"))
    elif is_flax_available():
        import jax

        return jax.device_count()
    else:


SUMMARY:
**Function Summary: `get_gpu_count`**

Retrieves the number of available GPUs, leveraging various frameworks (torch, tf, jax) or defaulting to 0 if none is supported.

### Key Functionality:

- Returns GPU count using:
  - `torch.cuda.device_count()` when torch is available
  - `len(tf.config.list_physical_devices("GPU"))` when tf is available
  - `jax.device_count()` when flax/jax is available

### Notable Dependencies/Requirements:

- `torch`
- `tensorflow` (tf)
- `jax` (or `flax`)
- `is_torch_available()`, `is_tf_available()`, and `is_flax_available()` internal utility functions (not shown in this code snippet)


key: chroma:document
string_value:
CODE:

def get_tests_dir(append_path=None):
    """
    Args:
        append_path: optional path to append to the tests dir path

    Return:
        The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is
        joined after the `tests` dir the former is provided.

    """
    # this function caller's __file__
    caller__file__ = inspect.stack()[1][1]
    tests_dir = os.path.abspath(os.path.dirname(caller__file__))

    while not tests_dir.endswith("tests"):
        tests_dir = os.path.dirname(tests_dir)

    if append_path:
        return os.path.join(tests_dir, append_path)
    else:


SUMMARY:
**Function Summary**

`get_tests_dir(append_path=None)`:

Returns the full path to the `tests` directory. optionally appends `append_path` to it.

**Key Implementation Details:**

* Uses caller's file path (`caller__file__`) to determine the tests dir
* Ensures the resulting path ends with "tests" using a loop

**Notable Dependencies/Requirements:**

* `inspect`
* `os`

**Primary Purpose:** Returns an absolute path to the `tests` directory.


key: docstring
string_value:
#
# Helper functions for dealing with testing text outputs
# The original code came from:
# https://github.com/fastai/fastai/blob/master/tests/utils/text.py
# When any function contains print() calls that get overwritten, like progress bars,
# a special care needs to be applied, since under pytest -s captured output (capsys
# or contextlib.redirect_stdout) contains any temporary printed strings, followed by
# \r's. This helper function ensures that the buffer will contain the same output
# with and without -s in pytest, by turning:
# foo bar\r tar mar\r final message
# into:
# final message
# it can handle a single string or a multiline buffer


key: chroma:document
string_value:
CODE:
# it can handle a single string or a multiline buffer
def apply_print_resets(buf):


SUMMARY:
**Function Summary: `apply_print_resets(buf)`**

**Purpose and Functionality:**
Removes trailing `\r` characters from the end of a given input buffer (`buf`) to ensure consistent behavior with and without `pytest -s` captor output.

**Implementation Details:**

* Uses `re.sub()` to replace trailing `\r` characters with an empty string.
* Sets `re.M` flag for multiline matching.
* Returns the modified buffer (`buf`).

**Notable Dependencies/Requirements:**
- Utilizes `re` (regular expression) module.
- Relies on `pytest` framework with `-s` option.

**Interfaces/Methods:**
- `re.sub()`: replaces substrings in a given string.

Note: The code uses a helper function from an external repository, but the implementation details are self-contained within this summary.


key: chroma:document
string_value:
CODE:

def assert_screenout(out, what):
    out_pr = apply_print_resets(out).lower()
    match_str = out_pr.find(what.lower())


SUMMARY:
**Summary**

Function `assert_screenout(out, what)` asserts that a specified output string (`out`) contains the expected strings (`what`). It normalizes the output and searches for matching values.

**Key Details**

- Normalizes output using `apply_print_resets(out)` before searching.
- Searches for lower-case version of input (`what`).
- Asserts presence of match string.


key: chroma:document
string_value:
CODE:

def set_model_tester_for_less_flaky_test(test_case):
    target_num_hidden_layers = 1
    # TODO (if possible): Avoid exceptional cases
    exceptional_classes = [
        "ZambaModelTester",
        "RwkvModelTester",
        "AriaVisionText2TextModelTester",
        "GPTNeoModelTester",
        "DPTModelTester",
    ]
    if test_case.model_tester.__class__.__name__ in exceptional_classes:
        target_num_hidden_layers = None
    if hasattr(test_case.model_tester, "out_features") or hasattr(test_case.model_tester, "out_indices"):
        target_num_hidden_layers = None

    if hasattr(test_case.model_tester, "num_hidden_layers") and target_num_hidden_layers is not None:
        test_case.model_tester.num_hidden_layers = target_num_hidden_layers
    if (
        hasattr(test_case.model_tester, "vision_config")
        and "num_hidden_layers" in test_case.model_tester.vision_config
        and target_num_hidden_layers is not None
    ):
        test_case.model_tester.vision_config = copy.deepcopy(test_case.model_tester.vision_config)
        test_case.model_tester.vision_config["num_hidden_layers"] = target_num_hidden_layers
    if (
        hasattr(test_case.model_tester, "text_config")
        and "num_hidden_layers" in test_case.model_tester.text_config
        and target_num_hidden_layers is not None
    ):
        test_case.model_tester.text_config = copy.deepcopy(test_case.model_tester.text_config)
        test_case.model_tester.text_config["num_hidden_layers"] = target_num_hidden_layers

    # A few model class specific handling

    # For Albert
    if hasattr(test_case.model_tester, "num_hidden_groups"):


SUMMARY:
**Summary: `set_model_tester_for_less_flaky_test` Function**

* **Purpose:** Adjusts model tester's configuration to reduce flakiness in tests for specific models (e.g., ZambaModelTester, AriaVisionText2TextModelTester).
* **Functionality:**
	+ Toggles or sets `num_hidden_layers` based on presence of certain attributes.
	+ Copies and overwrites `vision_config` and `text_config` dictionaries as needed.

**Key Implementation Details:**

* Iterates through exceptional classes (e.g., `ZambaModelTester`) and sets `target_num_hidden_layers` to `None` if the condition is met.
* Modifies `num_hidden_layers`, `num_hidden_groups`, and attributes of related configurations (`vision_config` and `text_config`) based on presence of certain attributes.

**Notable Dependencies/Requirements:**

* `copy`
* Model-specific configurations (e.g., `ZambaModelTester`) with attributes:
	+ `num_hidden_layers`
	+ `num_hidden_groups`
	+ `vision_config`
	+ `text_config`


key: chroma:document
string_value:
CODE:

def set_config_for_less_flaky_test(config):
    target_attrs = [
        "rms_norm_eps",
        "layer_norm_eps",
        "norm_eps",
        "norm_epsilon",
        "layer_norm_epsilon",
        "batch_norm_eps",
    ]
    for target_attr in target_attrs:
        setattr(config, target_attr, 1.0)

    # norm layers (layer/group norm, etc.) could cause flaky tests when the tensors have very small variance.
    # (We don't need the original epsilon values to check eager/sdpa matches)
    attrs = ["text_config", "vision_config", "text_encoder", "audio_encoder", "decoder"]
    for attr in attrs:
        if hasattr(config, attr):
            for target_attr in target_attrs:


SUMMARY:
**Function Summary: `set_config_for_less_flaky_test`**

**Purpose:** This function sets default values for specific configuration attributes (`rms_norm_eps`, `layer_norm_eps`, etc.) in the provided `config` object.

**Key Implementation Details:**

*   Sets default norm epsilon values for all applicable configuration attributes.
*   Iterates through attribute lists and recursively updates nested configurations.

**External Dependencies/Notable Requirements:**

*   `config` object must be a dictionary-like structure.
*   The `setattr` function is used to dynamically update the configuration attributes.

This function appears to be part of a test configuration adjustment mechanism, aiming to reduce flaky test issues caused by negligible tensor variances in norm layers.


key: chroma:document
string_value:
CODE:

def set_model_for_less_flaky_test(model):
    # Another way to make sure norm layers have desired epsilon. (Some models don't set it from its config.)
    target_names = ("LayerNorm", "GroupNorm", "BatchNorm", "RMSNorm", "BatchNorm2d", "BatchNorm1d")
    target_attrs = ["eps", "epsilon", "variance_epsilon"]
    if is_torch_available() and isinstance(model, torch.nn.Module):
        for module in model.modules():
            if type(module).__name__.endswith(target_names):
                for attr in target_attrs:
                    if hasattr(module, attr):


SUMMARY:
**Summary:**

`set_model_for_less_flaky_test(model)` 

* Sets the desired norms with `epsilon` values to various models, ensuring norm layers have predictable behavior.

**Key Implementation Details:**

- Loops through model's modules and sets `eps`, `epsilon`, or `variance_epsilon` attributes of norm layers to 1.0.
  
**Notable Dependencies or Requirements:**

- `torch.nn.Module` and `torch` library for availability check (`is_torch_available()`)


key: chroma:document
string_value:
CODE:

class CaptureStd:
    """
    Context manager to capture:

        - stdout: replay it, clean it up and make it available via `obj.out`
        - stderr: replay it and make it available via `obj.err`

    Args:
        out (`bool`, *optional*, defaults to `True`): Whether to capture stdout or not.
        err (`bool`, *optional*, defaults to `True`): Whether to capture stderr or not.
        replay (`bool`, *optional*, defaults to `True`): Whether to replay or not.
            By default each captured stream gets replayed back on context's exit, so that one can see what the test was
            doing. If this is a not wanted behavior and the captured data shouldn't be replayed, pass `replay=False` to
            disable this feature.

    Examples:

    ```python
    # to capture stdout only with auto-replay
    with CaptureStdout() as cs:
        print("Secret message")
    assert "message" in cs.out

    # to capture stderr only with auto-replay
    import sys

    with CaptureStderr() as cs:
        print("Warning: ", file=sys.stderr)
    assert "Warning" in cs.err

    # to capture both streams with auto-replay
    with CaptureStd() as cs:
        print("Secret message")
        print("Warning: ", file=sys.stderr)
    assert "message" in cs.out
    assert "Warning" in cs.err

    # to capture just one of the streams, and not the other, with auto-replay
    with CaptureStd(err=False) as cs:
        print("Secret message")
    assert "message" in cs.out
    # but best use the stream-specific subclasses

    # to capture without auto-replay
    with CaptureStd(replay=False) as cs:
        print("Secret message")
    assert "message" in cs.out
    ```"""

    def __init__(self, out=True, err=True, replay=True):
        self.replay = replay

        if out:
            self.out_buf = StringIO()
            self.out = "error: CaptureStd context is unfinished yet, called too early"
        else:
            self.out_buf = None
            self.out = "not capturing stdout"

        if err:
            self.err_buf = StringIO()
            self.err = "error: CaptureStd context is unfinished yet, called too early"
        else:
            self.err_buf = None
            self.err = "not capturing stderr"

    def __enter__(self):
        if self.out_buf:
            self.out_old = sys.stdout
            sys.stdout = self.out_buf

        if self.err_buf:
            self.err_old = sys.stderr
            sys.stderr = self.err_buf

        return self

    def __exit__(self, *exc):
        if self.out_buf:
            sys.stdout = self.out_old
            captured = self.out_buf.getvalue()
            if self.replay:
                sys.stdout.write(captured)
            self.out = apply_print_resets(captured)

        if self.err_buf:
            sys.stderr = self.err_old
            captured = self.err_buf.getvalue()
            if self.replay:
                sys.stderr.write(captured)
            self.err = captured

    def __repr__(self):
        msg = ""
        if self.out_buf:
            msg += f"stdout: {self.out}\n"
        if self.err_buf:
            msg += f"stderr: {self.err}\n"


SUMMARY:
**Class CaptureStd Context Manager**
=====================================

### Purpose and Functionality

The `CaptureStd` class is a context manager that captures and replays output to `stdout` and `stderr`. It provides flexible management of captured streams, including replaying them back on exit.

### Key Implementation Details

*   Initializes contexts for capturing `stdout` and/or `stderr`, using `StringIO` buffers.
*   Sets up `sys.stdout` and `sys.stderr` to capture output in `__enter__`.
*   Restores original `sys.stdout` and `sys.stderr` values on exit (`__exit__`).
*   Replays captured streams back to the same stream or writes them directly to it, depending on the `replay` flag.
*   Manages stored output via method `apply_print_resets`.

### Notable Dependencies and Requirements

*   Uses `StringIO`, a built-in Python module for easily creating file-like objects as strings.

This code provides a concise way to manage capturing and replaying output streams, simplifying the process of testing and debugging print statements.


key: docstring
string_value:
# in tests it's the best to capture only the stream that's wanted, otherwise
# it's easy to miss things, so unless you need to capture both streams, use the
# subclasses below (less typing). Or alternatively, configure `CaptureStd` to
# disable the stream you don't need to test.


key: chroma:document
string_value:
CODE:

class CaptureStdout(CaptureStd):
    """Same as CaptureStd but captures only stdout"""

    def __init__(self, replay=True):


SUMMARY:
## Class Summary: CaptureStdOut

### Purpose and Functionality

The `CaptureStdOut` class is a subclass of `CaptureStd` that captures only the standard output (`stdout`) stream, inheriting its functionality while filtering out error output. It allows for capturing only the desired stream in tests.

#### Primary Purpose
Captures only `stdout` stream with optional replay feature.

### Key Implementation Details

- Utilizes `super()` to initialize `CaptureStd` with `err=False` and `replay=replay`.
- Includes a docstring advice on not capturing both streams or configuring `CaptureStd` to disable unnecessary streams.

#### Important Method Signatures and Interfaces
- `__init__(self, replay: bool)`: Initializes the object with optional replay feature.
  - `replay: bool`
  
### Dependencies and Requirements

- Subclasses `CaptureStd`
- Requires `super()` for inheritance and method call chaining.


key: chroma:document
string_value:
CODE:

class CaptureStderr(CaptureStd):
    """Same as CaptureStd but captures only stderr"""

    def __init__(self, replay=True):


SUMMARY:
**Summary:**

The `CaptureStderr` class extends the base `CaptureStd` class to specifically capture standard error output. It inherits the functionality of its parent class while modifying the initialization with only stderr enabled.

**Key Details:**

* Extends `CaptureStd` and overrides it with `out=False`
* Initializes with replay enabled (`replay=True`) by default

**Dependencies:**

* Inheriting from `CaptureStd`
* External dependency on `super()` for inheritance


key: chroma:document
string_value:
CODE:

class CaptureLogger:
    """
    Context manager to capture `logging` streams

    Args:
        logger: 'logging` logger object

    Returns:
        The captured output is available via `self.out`

    Example:

    ```python
    >>> from transformers import logging
    >>> from transformers.testing_utils import CaptureLogger

    >>> msg = "Testing 1, 2, 3"
    >>> logging.set_verbosity_info()
    >>> logger = logging.get_logger("transformers.models.bart.tokenization_bart")
    >>> with CaptureLogger(logger) as cl:
    ...     logger.info(msg)
    >>> assert cl.out, msg + "\n"
    ```
    """

    def __init__(self, logger):
        self.logger = logger
        self.io = StringIO()
        self.sh = logging.StreamHandler(self.io)
        self.out = ""

    def __enter__(self):
        self.logger.addHandler(self.sh)
        return self

    def __exit__(self, *exc):
        self.logger.removeHandler(self.sh)
        self.out = self.io.getvalue()

    def __repr__(self):


SUMMARY:
**Summary:`

The `CaptureLogger` class is a context manager that captures the output of a given `logging` logger. It ensures that log messages are redirected to a string stream (`StringIO`) instead of being displayed in the console.

**Key Implementation Details:**

* Redirects log messages to a string stream (`self.io`) using `logging.StreamHandler`.
* Allows capture of log output via attribute `self.out`.
* Implements context manager protocol for use with `with` statement.

**Notable Dependencies/Requirements:**

* `logging` module
* `StringIO` class from `io` module

**Primary Purpose:** Capture and redirect logging messages to a string stream.


key: chroma:document
string_value:
CODE:

class TemporaryHubRepo:
    """Create a temporary Hub repository and return its `RepoUrl` object. This is similar to
    `tempfile.TemporaryDirectory` and can be used as a context manager. For example:

        with TemporaryHubRepo(token=self._token) as temp_repo:
            ...

    Upon exiting the context, the repository and everything contained in it are removed.

    Example:

    ```python
    with TemporaryHubRepo(token=self._token) as temp_repo:
        model.push_to_hub(tmp_repo.repo_id, token=self._token)
    ```
    """

    def __init__(self, namespace: Optional[str] = None, token: Optional[str] = None) -> None:
        self.token = token
        with tempfile.TemporaryDirectory() as tmp_dir:
            repo_id = Path(tmp_dir).name
            if namespace is not None:
                repo_id = f"{namespace}/{repo_id}"
            self.repo_url = huggingface_hub.create_repo(repo_id, token=self.token)

    def __enter__(self):
        return self.repo_url

    def __exit__(self, exc, value, tb):


SUMMARY:
**Summary: Temporary Hub Repository Context Manager**

```python
TemporaryHubRepo:
  - **Purpose:** Create a temporary Hub repository with optional namespace and return its RepoUrl object as a context manager.
 
  Attributes:
    - token (str, Optional): API token for accessing the Hub repository.

  Methods:
    - `__enter__`: Returns the RepoUrl object when entering the context.
    - `__exit__`: Deletes the created Hub repository when exiting the context.

  Key Operations:
   - Creates a temporary directory using `tempfile.TemporaryDirectory`.
   - Generates a unique repository ID from the directory name.
   - Creates and initializes a new Hub repository with the given namespace (if provided) and token.
   - Uses `huggingface_hub.create_repo` to create the repository.
   - Deletes the repository using `delete_repo` when exiting the context.

  Examples:
    ```python
    # Create and exit the context
    with TemporaryHubRepo(token='my_token') as repo_url:
      pass
    
    # With namespace
    with TemporaryHubRepo(namespace='my_namespace', token='my_token') as repo_url:
      pass
    ```

  External Dependencies:
   - `tempfile`
   - `huggingface_hub` (via `create_repo` and `delete_repo`)
```

This summary aims to provide a concise overview of the key aspects of the `TemporaryHubRepo` class, including its purpose, attributes, methods, and external dependencies.


key: chroma:document
string_value:
CODE:

class TestCasePlus(unittest.TestCase):
    """
    This class extends *unittest.TestCase* with additional features.

    Feature 1: A set of fully resolved important file and dir path accessors.

    In tests often we need to know where things are relative to the current test file, and it's not trivial since the
    test could be invoked from more than one directory or could reside in sub-directories with different depths. This
    class solves this problem by sorting out all the basic paths and provides easy accessors to them:

    - `pathlib` objects (all fully resolved):

       - `test_file_path` - the current test file path (=`__file__`)
       - `test_file_dir` - the directory containing the current test file
       - `tests_dir` - the directory of the `tests` test suite
       - `examples_dir` - the directory of the `examples` test suite
       - `repo_root_dir` - the directory of the repository
       - `src_dir` - the directory of `src` (i.e. where the `transformers` sub-dir resides)

    - stringified paths---same as above but these return paths as strings, rather than `pathlib` objects:

       - `test_file_path_str`
       - `test_file_dir_str`
       - `tests_dir_str`
       - `examples_dir_str`
       - `repo_root_dir_str`
       - `src_dir_str`

    Feature 2: Flexible auto-removable temporary dirs which are guaranteed to get removed at the end of test.

    1. Create a unique temporary dir:

    ```python
    def test_whatever(self):
        tmp_dir = self.get_auto_remove_tmp_dir()
    ```

    `tmp_dir` will contain the path to the created temporary dir. It will be automatically removed at the end of the
    test.


    2. Create a temporary dir of my choice, ensure it's empty before the test starts and don't
    empty it after the test.

    ```python
    def test_whatever(self):
        tmp_dir = self.get_auto_remove_tmp_dir("./xxx")
    ```

    This is useful for debug when you want to monitor a specific directory and want to make sure the previous tests
    didn't leave any data in there.

    3. You can override the first two options by directly overriding the `before` and `after` args, leading to the
        following behavior:

    `before=True`: the temporary dir will always be cleared at the beginning of the test.

    `before=False`: if the temporary dir already existed, any existing files will remain there.

    `after=True`: the temporary dir will always be deleted at the end of the test.

    `after=False`: the temporary dir will always be left intact at the end of the test.

    Note 1: In order to run the equivalent of `rm -r` safely, only subdirs of the project repository checkout are
    allowed if an explicit `tmp_dir` is used, so that by mistake no `/tmp` or similar important part of the filesystem
    will get nuked. i.e. please always pass paths that start with `./`

    Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested
    otherwise.

    Feature 3: Get a copy of the `os.environ` object that sets up `PYTHONPATH` specific to the current test suite. This
    is useful for invoking external programs from the test suite - e.g. distributed training.


    ```python
    def test_whatever(self):
        env = self.get_env()
    ```"""

    def setUp(self):
        # get_auto_remove_tmp_dir feature:
        self.teardown_tmp_dirs = []

        # figure out the resolved paths for repo_root, tests, examples, etc.
        self._test_file_path = inspect.getfile(self.__class__)
        path = Path(self._test_file_path).resolve()
        self._test_file_dir = path.parents[0]
        for up in [1, 2, 3]:
            tmp_dir = path.parents[up]
            if (tmp_dir / "src").is_dir() and (tmp_dir / "tests").is_dir():
                break
        if tmp_dir:
            self._repo_root_dir = tmp_dir
        else:
            raise ValueError(f"can't figure out the root of the repo from {self._test_file_path}")
        self._tests_dir = self._repo_root_dir / "tests"
        self._examples_dir = self._repo_root_dir / "examples"
        self._src_dir = self._repo_root_dir / "src"

    @property
    def test_file_path(self):
        return self._test_file_path

    @property
    def test_file_path_str(self):
        return str(self._test_file_path)

    @property
    def test_file_dir(self):
        return self._test_file_dir

    @property
    def test_file_dir_str(self):
        return str(self._test_file_dir)

    @property
    def tests_dir(self):
        return self._tests_dir

    @property
    def tests_dir_str(self):
        return str(self._tests_dir)

    @property
    def examples_dir(self):
        return self._examples_dir

    @property
    def examples_dir_str(self):
        return str(self._examples_dir)

    @property
    def repo_root_dir(self):
        return self._repo_root_dir

    @property
    def repo_root_dir_str(self):
        return str(self._repo_root_dir)

    @property
    def src_dir(self):
        return self._src_dir

    @property
    def src_dir_str(self):
        return str(self._src_dir)

    def get_env(self):
        """
        Return a copy of the `os.environ` object that sets up `PYTHONPATH` correctly, depending on the test suite it's
        invoked from. This is useful for invoking external programs from the test suite - e.g. distributed training.

        It always inserts `./src` first, then `./tests` or `./examples` depending on the test suite type and finally
        the preset `PYTHONPATH` if any (all full resolved paths).

        """
        env = os.environ.copy()
        paths = [self.src_dir_str]
        if "/examples" in self.test_file_dir_str:
            paths.append(self.examples_dir_str)
        else:
            paths.append(self.tests_dir_str)
        paths.append(env.get("PYTHONPATH", ""))

        env["PYTHONPATH"] = ":".join(paths)
        return env

    def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None):
        """
        Args:
            tmp_dir (`string`, *optional*):
                if `None`:

                   - a unique temporary path will be created
                   - sets `before=True` if `before` is `None`
                   - sets `after=True` if `after` is `None`
                else:

                   - `tmp_dir` will be created
                   - sets `before=True` if `before` is `None`
                   - sets `after=False` if `after` is `None`
            before (`bool`, *optional*):
                If `True` and the `tmp_dir` already exists, make sure to empty it right away if `False` and the
                `tmp_dir` already exists, any existing files will remain there.
            after (`bool`, *optional*):
                If `True`, delete the `tmp_dir` at the end of the test if `False`, leave the `tmp_dir` and its contents
                intact at the end of the test.

        Returns:
            tmp_dir(`string`): either the same value as passed via *tmp_dir* or the path to the auto-selected tmp dir
        """
        if tmp_dir is not None:
            # defining the most likely desired behavior for when a custom path is provided.
            # this most likely indicates the debug mode where we want an easily locatable dir that:
            # 1. gets cleared out before the test (if it already exists)
            # 2. is left intact after the test
            if before is None:
                before = True
            if after is None:
                after = False

            # using provided path
            path = Path(tmp_dir).resolve()

            # to avoid nuking parts of the filesystem, only relative paths are allowed
            if not tmp_dir.startswith("./"):
                raise ValueError(
                    f"`tmp_dir` can only be a relative path, i.e. `./some/path`, but received `{tmp_dir}`"
                )

            # ensure the dir is empty to start with
            if before is True and path.exists():
                shutil.rmtree(tmp_dir, ignore_errors=True)

            path.mkdir(parents=True, exist_ok=True)

        else:
            # defining the most likely desired behavior for when a unique tmp path is auto generated
            # (not a debug mode), here we require a unique tmp dir that:
            # 1. is empty before the test (it will be empty in this situation anyway)
            # 2. gets fully removed after the test
            if before is None:
                before = True
            if after is None:
                after = True

            # using unique tmp dir (always empty, regardless of `before`)
            tmp_dir = tempfile.mkdtemp()

        if after is True:
            # register for deletion
            self.teardown_tmp_dirs.append(tmp_dir)

        return tmp_dir

    def python_one_liner_max_rss(self, one_liner_str):
        """
        Runs the passed python one liner (just the code) and returns how much max cpu memory was used to run the
        program.

        Args:
            one_liner_str (`string`):
                a python one liner code that gets passed to `python -c`

        Returns:
            max cpu memory bytes used to run the program. This value is likely to vary slightly from run to run.

        Requirements:
            this helper needs `/usr/bin/time` to be installed (`apt install time`)

        Example:

        ```
        one_liner_str = 'from transformers import AutoModel; AutoModel.from_pretrained("google-t5/t5-large")'
        max_rss = self.python_one_liner_max_rss(one_liner_str)
        ```
        """

        if not cmd_exists("/usr/bin/time"):
            raise ValueError("/usr/bin/time is required, install with `apt install time`")

        cmd = shlex.split(f"/usr/bin/time -f %M python -c '{one_liner_str}'")
        with CaptureStd() as cs:
            execute_subprocess_async(cmd, env=self.get_env())
        # returned data is in KB so convert to bytes
        max_rss = int(cs.err.split("\n")[-2].replace("stderr: ", "")) * 1024
        return max_rss

    def tearDown(self):
        # get_auto_remove_tmp_dir feature: remove registered temp dirs
        for path in self.teardown_tmp_dirs:
            shutil.rmtree(path, ignore_errors=True)
        self.teardown_tmp_dirs = []
        if is_accelerate_available():
            AcceleratorState._reset_state()
            PartialState._reset_state()

            # delete all the env variables having `ACCELERATE` in them
            for k in list(os.environ.keys()):
                if "ACCELERATE" in k:


SUMMARY:
This code is a test class written in Python, which provides several utility methods for running tests and measuring CPU memory usage. Here's a concise summary:

**Main Features:**

1. **Temporary Directory Management**: The `get_auto_remove_tmp_dir` method creates a temporary directory with optional custom behavior (e.g., empty before test or kept after test).
2. **CPU Memory Measurement**: The `python_one_liner_max_rss` method runs a Python one-liner code and returns the maximum CPU memory usage in bytes.
3. **Environment Management**: The `get_env` method returns an environment dictionary with relevant settings for testing, including `/usr/bin/time` location.

**Dependent Services:**

1. `/usr/bin/time` is required for measuring CPU memory usage (install it if not available).
2. `Accelerate` framework is used for distributed computing (reset state after test).

Overall, this class provides a set of convenient methods for managing temporary directories and measuring CPU memory usage, making it easier to write tests with these features.


key: chroma:document
string_value:
CODE:

def mockenv(**kwargs):
    """
    this is a convenience wrapper, that allows this ::

    @mockenv(RUN_SLOW=True, USE_TF=False) def test_something():
        run_slow = os.getenv("RUN_SLOW", False) use_tf = os.getenv("USE_TF", False)

    """


SUMMARY:
**Summary**

The `mockenv` function is a convenience wrapper that enables environment variable mocking for testing purposes. It replaces the target environment variables with values provided through keyword arguments (`kwargs`). The function leverages the `unittest.mock.patch.dict` function to achieve this.

**Key Implementation Details:**

* Replaces `os.environ` with mocked value using `patch.dict`
* Allows mocking of key environment variables using `__kwargs__`

**Notable Dependencies/Requirements:**

* `os` module (specifically `os.getenv` and `os.environ`)
* `unittest.mock` module (specifically `patch.dict`)


key: chroma:document
string_value:
CODE:

def pytest_addoption_shared(parser):
    """
    This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there.

    It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest`
    option.

    """
    option = "--make-reports"
    if option not in pytest_opt_registered:
        parser.addoption(
            option,
            action="store",
            default=False,
            help="generate report files. The value of this option is used as a prefix to report names",
        )


SUMMARY:
**Summary: `pytest_addoption_shared` Function**

This function, called from `conftest.py`, adds a custom option `--make-reports` to the pytest parser. It registers the option with a default value of `False`.

**Key details:**

* **Option signature:** `--make-reports`
* **Action:** Store the option value as an integer (`store`)
* **Default value:** `False`
* **Help message:** Generates report files and uses the option value as a prefix for report names.

**Dependencies/erequisites:**

None external dependencies mentioned. The code assumes the existence of:
- `pytest_opt_registered` (a dictionary to track registered options)
- A `conftest.py` file referencing this function via `pytest_addoption`


key: chroma:document
string_value:
CODE:

def pytest_terminal_summary_main(tr, id):
    """
    Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current
    directory. The report files are prefixed with the test suite name.

    This function emulates --duration and -rA pytest arguments.

    This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined
    there.

    Args:
    - tr: `terminalreporter` passed from `conftest.py`
    - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is
      needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other.

    NB: this functions taps into a private _pytest API and while unlikely, it could break should pytest do internal
    changes - also it calls default internal methods of terminalreporter which can be hijacked by various `pytest-`
    plugins and interfere.

    """
    from _pytest.config import create_terminal_writer

    if not len(id):
        id = "tests"

    config = tr.config
    orig_writer = config.get_terminal_writer()
    orig_tbstyle = config.option.tbstyle
    orig_reportchars = tr.reportchars

    dir = f"reports/{id}"
    Path(dir).mkdir(parents=True, exist_ok=True)
    report_files = {
        k: f"{dir}/{k}.txt"
        for k in [
            "durations",
            "errors",
            "failures_long",
            "failures_short",
            "failures_line",
            "passes",
            "stats",
            "summary_short",
            "warnings",
        ]
    }

    # custom durations report
    # note: there is no need to call pytest --durations=XX to get this separate report
    # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66
    dlist = []
    for replist in tr.stats.values():
        for rep in replist:
            if hasattr(rep, "duration"):
                dlist.append(rep)
    if dlist:
        dlist.sort(key=lambda x: x.duration, reverse=True)
        with open(report_files["durations"], "w") as f:
            durations_min = 0.05  # sec
            f.write("slowest durations\n")
            for i, rep in enumerate(dlist):
                if rep.duration < durations_min:
                    f.write(f"{len(dlist)-i} durations < {durations_min} secs were omitted")
                    break
                f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n")

    def summary_failures_short(tr):
        # expecting that the reports were --tb=long (default) so we chop them off here to the last frame
        reports = tr.getreports("failed")
        if not reports:
            return
        tr.write_sep("=", "FAILURES SHORT STACK")
        for rep in reports:
            msg = tr._getfailureheadline(rep)
            tr.write_sep("_", msg, red=True, bold=True)
            # chop off the optional leading extra frames, leaving only the last one
            longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S)
            tr._tw.line(longrepr)
            # note: not printing out any rep.sections to keep the report short

    # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each
    # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814
    # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g.
    # pytest-instafail does that)

    # report failures with line/short/long styles
    config.option.tbstyle = "auto"  # full tb
    with open(report_files["failures_long"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.summary_failures()

    # config.option.tbstyle = "short" # short tb
    with open(report_files["failures_short"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        summary_failures_short(tr)

    config.option.tbstyle = "line"  # one line per error
    with open(report_files["failures_line"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.summary_failures()

    with open(report_files["errors"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.summary_errors()

    with open(report_files["warnings"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.summary_warnings()  # normal warnings
        tr.summary_warnings()  # final warnings

    tr.reportchars = "wPpsxXEf"  # emulate -rA (used in summary_passes() and short_test_summary())

    # Skip the `passes` report, as it starts to take more than 5 minutes, and sometimes it timeouts on CircleCI if it
    # takes > 10 minutes (as this part doesn't generate any output on the terminal).
    # (also, it seems there is no useful information in this report, and we rarely need to read it)
    # with open(report_files["passes"], "w") as f:
    #     tr._tw = create_terminal_writer(config, f)
    #     tr.summary_passes()

    with open(report_files["summary_short"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.short_test_summary()

    with open(report_files["stats"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.summary_stats()

    # restore:
    tr._tw = orig_writer
    tr.reportchars = orig_reportchars


SUMMARY:
**Function Summary: `pytest_terminal_summary_main`**

Generate multiple reports at the end of test suite run.

**Purpose:** Emulate `-rA pytest` arguments by generating report files with specific formats and content.

**Key Implementation Details:**

* Generates reports for:
	+ Durations
	+ Failures (long, short, line)
	+ Errors
	+ Warnings
* Uses `terminalreporter` to write reports to files in a designated directory.
* Hijacks internal methods of `terminalreporter` to generate custom report layouts.

**Notable Dependencies and Requirements:**

* `_pytest.config`: used to access configuration settings and terminal writer functionality.
* `terminalreporter`: required for writing reports to files.
* `config.option.tbstyle`: controls the format of test results output (auto, short, line).


key: chroma:document
string_value:
CODE:

class _RunOutput:
    def __init__(self, returncode, stdout, stderr):
        self.returncode = returncode
        self.stdout = stdout


SUMMARY:
**Summary**

Class `_RunOutput` encapsulates the output of a runtime process. It stores three attributes:

* `returncode`: The exit status of the process (int)
* `stdout`: The standard output of the process (str or bytes)
* `stderr`: The standard error of the process (str or bytes)

The class facilitates the structured storage and retrieval of this information.

**Implementation Details**

Key implementation details include:

* An initialization method `__init__` that accepts a return code, stdout, and stderr as parameters.
* A docstring comment (`# noqa`) indicating adherence to Python's PEP-257 style guide for docstrings.


key: chroma:document
string_value:
CODE:

async def _read_stream(stream, callback):
    while True:
        line = await stream.readline()
        if line:
            callback(line)
        else:


SUMMARY:
**Summary:**

`_read_stream` Function:

- Purpose: Reads and asynchronously processes a stream of lines using the provided `callback` function.

**Key Implementation Details:**

* Iterates indefinitely until an empty line (`''`) is encountered.
* Calls `callback` with each non-empty read line.

**Dependencies:**

* None explicitly listed; however, assumes access to `stream.readline()` method.

Note: This summary captures the essential functionality and implementation details of the `_read_stream` function.


key: chroma:document
string_value:
CODE:

async def _stream_subprocess(cmd, env=None, stdin=None, timeout=None, quiet=False, echo=False) -> _RunOutput:
    if echo:
        print("\nRunning: ", " ".join(cmd))

    p = await asyncio.create_subprocess_exec(
        cmd[0],
        *cmd[1:],
        stdin=stdin,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env,
    )

    # note: there is a warning for a possible deadlock when using `wait` with huge amounts of data in the pipe
    # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.wait
    #
    # If it starts hanging, will need to switch to the following code. The problem is that no data
    # will be seen until it's done and if it hangs for example there will be no debug info.
    # out, err = await p.communicate()
    # return _RunOutput(p.returncode, out, err)

    out = []
    err = []

    def tee(line, sink, pipe, label=""):
        line = line.decode("utf-8").rstrip()
        sink.append(line)
        if not quiet:
            print(label, line, file=pipe)

    # XXX: the timeout doesn't seem to make any difference here
    await asyncio.wait(
        [
            _read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:")),
            _read_stream(p.stderr, lambda l: tee(l, err, sys.stderr, label="stderr:")),
        ],
        timeout=timeout,
    )


SUMMARY:
**Function Summary: `_stream_subprocess`**

**Purpose and Functionality**
-------------------------

The `_stream_subprocess` function asynchronously executes a subprocess using `asyncio.create_subprocess_exec`. It takes optional arguments for input/output streams and timeout.

**Key Implementation Details**
-----------------------------

*   The function uses `tee` to forward output lines to both the main stream and a separate file (stdout/stderr).
*   It waits for the subprocess to finish with a timeout, preventing potential deadlocks.
*   If the subprocess hangs, it falls back to waiting for completion without debug info.

**Notable Dependencies or Requirements**
-----------------------------------------

*   `asyncio`
*   `sys` for printing output
*   `_ReadStream` abstraction (not shown in this code snippet)

The function returns an instance of `_RunOutput`, containing the subprocess' returncode, stdout, and stderr.


key: chroma:document
string_value:
CODE:

def execute_subprocess_async(cmd, env=None, stdin=None, timeout=180, quiet=False, echo=True) -> _RunOutput:
    loop = asyncio.get_event_loop()
    result = loop.run_until_complete(
        _stream_subprocess(cmd, env=env, stdin=stdin, timeout=timeout, quiet=quiet, echo=echo)
    )

    cmd_str = " ".join(cmd)
    if result.returncode > 0:
        stderr = "\n".join(result.stderr)
        raise RuntimeError(
            f"'{cmd_str}' failed with returncode {result.returncode}\n\n"
            f"The combined stderr from workers follows:\n{stderr}"
        )

    # check that the subprocess actually did run and produced some output, should the test rely on
    # the remote side to do the testing
    if not result.stdout and not result.stderr:
        raise RuntimeError(f"'{cmd_str}' produced no output.")



SUMMARY:
**Summary of `execute_subprocess_async` Function:**

**Purpose and Functionality:** 
The `_execute_subprocess_async` function executes a subprocess asynchronously, handling the outcome in an event loop. It captures the standard input, output, and error streams from the subprocess.

**Key Implementation Details:**

*   Uses asynchronous execution with `asyncio.get_event_loop()` and `loop.run_until_complete()`.
*   Calls `_stream_subprocess` internally to manage the subprocess.
*   Verifies the return code and combined error messages from the subprocess.

**Notable Dependencies or Requirements:**

*   *_RunOutput:* A custom class (not shown in this snippet) that encapsulates the output of the subprocess.
*   *_stream_subprocess:* Another function (not shown in this snippet) handling the subprocess execution, which is imported.

This summary aims to preserve key elements of the original function without including external functions and unnecessary details.


key: chroma:document
string_value:
CODE:

def pytest_xdist_worker_id():
    """
    Returns an int value of worker's numerical id under `pytest-xdist`'s concurrent workers `pytest -n N` regime, or 0
    if `-n 1` or `pytest-xdist` isn't being used.
    """
    worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
    worker = re.sub(r"^gw", "", worker, 0, re.M)


SUMMARY:
**Summary:**

The `pytest_xdist_worker_id` function returns the numerical ID of a worker under pytest-xdist's concurrent testing framework. It retrieves the worker ID from the environment variable "PYTEST_XDIST_WORKER", cleans it by removing "gw" prefix, and returns it as an integer. If no valid worker ID is found or -n 1 flag is used, it returns 0.

**Key Implementation Details:**

* Uses `os.environ.get()` to retrieve the "PYTEST_XDIST_WORKER" environment variable
* Applies regular expression substitution (`re.sub`) to remove "gw" prefix from the worker ID

**Notable Dependencies and Requirements:**

* `os` and `re` modules are used for environment variable retrieval and regular expression operations, respectively.


key: chroma:document
string_value:
CODE:

def get_torch_dist_unique_port():
    """
    Returns a port number that can be fed to `torch.distributed.launch`'s `--master_port` argument.

    Under `pytest-xdist` it adds a delta number based on a worker id so that concurrent tests don't try to use the same
    port at once.
    """
    port = 29500
    uniq_delta = pytest_xdist_worker_id()


SUMMARY:
**Summary**

`get_torch_dist_unique_port()` returns a unique port number for PyTorch distributed training, combining a base port (29500) with a delta value generated by `pytest_xdist_worker_id()` to prevent concurrent tests from using the same port.

**Key details**

* Port base: 29500
* Uses `pytest_xdist_worker_id()` to generate a unique delta value

**Implementation notes**

* Adds delta value to base port to create unique port number
* Returns port number suitable for `torch.distributed.launch`'s `--master_port` argument


key: chroma:document
string_value:
CODE:

def nested_simplify(obj, decimals=3):
    """
    Simplifies an object by rounding float numbers, and downcasting tensors/numpy arrays to get simple equality test
    within tests.
    """
    import numpy as np

    if isinstance(obj, list):
        return [nested_simplify(item, decimals) for item in obj]
    if isinstance(obj, tuple):
        return tuple([nested_simplify(item, decimals) for item in obj])
    elif isinstance(obj, np.ndarray):
        return nested_simplify(obj.tolist())
    elif isinstance(obj, Mapping):
        return {nested_simplify(k, decimals): nested_simplify(v, decimals) for k, v in obj.items()}
    elif isinstance(obj, (str, int, np.int64)):
        return obj
    elif obj is None:
        return obj
    elif is_torch_available() and isinstance(obj, torch.Tensor):
        return nested_simplify(obj.tolist(), decimals)
    elif is_tf_available() and tf.is_tensor(obj):
        return nested_simplify(obj.numpy().tolist())
    elif isinstance(obj, float):
        return round(obj, decimals)
    elif isinstance(obj, (np.int32, np.float32, np.float16)):
        return nested_simplify(obj.item(), decimals)
    else:


SUMMARY:
**Nested Simplify Function Summary**

### Purpose and Functionality

The `nested_simplify` function simplifies an object by rounding float numbers, downcasting tensors/numpy arrays for equality tests in tests. It recursively processes various data types (e.g., lists, tuples, dictionaries, torch/tf tensors, numpy arrays, float, int).

### Key Implementation Details

* Rounds float numbers using `round()`.
* Downcasts tensors to numpy arrays for equality tests.
* Uses function nesting and recursion.

### Notable Dependencies and Requirements

* `numpy` for numpy array support
* Torch (`torch`) for torch tensor support (if available)
* TensorFlow (`tf`) for tf tensor support (if available)

### Method Signatures and Interfaces

* Applies to various data types, including `list`, `tuple`, `np.ndarray`, `Mapping`, `str`, `int`, `float`
* Utilizes recursive function calls
* Returns simplified object


key: chroma:document
string_value:
CODE:

def check_json_file_has_correct_format(file_path):
    with open(file_path, "r") as f:
        lines = f.readlines()
        if len(lines) == 1:
            # length can only be 1 if dict is empty
            assert lines[0] == "{}"
        else:
            # otherwise make sure json has correct format (at least 3 lines)
            assert len(lines) >= 3
            # each key one line, ident should be 2, min length is 3
            assert lines[0].strip() == "{"
            for line in lines[1:-1]:
                left_indent = len(lines[1]) - len(lines[1].lstrip())
                assert left_indent == 2


SUMMARY:
**Summary:**

### Code Purpose and Functionality:

The `check_json_file_has_correct_format` function validates if a provided JSON file follows the correct format. It checks for specific requirements, including:
- Exact structure consisting of one dictionary with no nested elements.
- A minimum of three lines.
- Indentation pattern where each key is on its own line and has a left indentation of 2 spaces.

### Key Implementation Details:

* Reads the JSON file by line and compares it to expected patterns using assertions.
* Uses `assert` statements to validate specific conditions such as the length of the JSON data, line endings at specific points (e.g., key starting lines), and correct closing brackets.

### Notable Dependencies or Requirements:

No external dependencies are referenced. The function only relies on Python's built-in libraries for file operations.

**RAG System Summary:**
`check_json_file_has_correct_format(file_path)`: Validates a JSON file by checking its structure, line counts, and indentation patterns using assertions. **Assumes correct file existence and readable formats.**


key: chroma:document
string_value:
CODE:

def to_2tuple(x):
    if isinstance(x, collections.abc.Iterable):
        return x


SUMMARY:
**Function Summary: `to_2tuple`**

* **Purpose:** Converts an iterable object to a tuple containing two identical elements.
	+ Primary purpose: Ensures input consistency by forcing iterables into a predictable format.
* **Implementation Details:**
	+ Utilizes Python's built-in `isinstance` check for iterability.
	+ Returns the input object unchanged if it is already an iterable.
	+ Consistently returns a tuple with two identical elements `(x, x)` for non-iterable inputs.
* **Dependencies:** None.
 
Note: The `collections.abc` module is required as an external dependency to use the `Iterable` abstract base class.


key: chroma:document
string_value:
CODE:
# These utils relate to ensuring the right error message is received when running scripts
class SubprocessCallException(Exception):


SUMMARY:
**Summary**

`SubprocessCallException` is a custom exception class used for handling subprocess-related errors, providing a centralized method for propagating and handling exceptions when running external scripts.

**Key Details**

* Exception type: `SubprocessCallException`
* Primary purpose: Handling subprocess-related errors
* Interfaces:
	+ Inherited from `Exception`

**Note**: The code does not specify any specific error messages or behavior; it only provides the exception class.


key: chroma:document
string_value:
CODE:

def run_command(command: List[str], return_stdout=False):
    """
    Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture
    if an error occured while running `command`
    """
    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT)
        if return_stdout:
            if hasattr(output, "decode"):
                output = output.decode("utf-8")
            return output
    except subprocess.CalledProcessError as e:
        raise SubprocessCallException(
            f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}"


SUMMARY:
**Function Summary:**

`run_command` executes a command using `subprocess.check_output`, capturing both stdout and stderr. It returns the output as a string if `return_stdout` is True, or raises an exception with the error message otherwise.

**Key Implementation Details:**

* Uses `subprocess.check_output` to run the command
* Capures stderr in addition to stdout using `stderr=subprocess.STDOUT`
* Decodes output to text if it's a bytes object

**Notable Dependencies and Requirements:**

* `subprocess`: used for running external commands
* `subprocess.check_output`: used to execute the command and capture its output
* `SubprocessCallException`: custom exception class raised on error


key: chroma:document
string_value:
CODE:

class RequestCounter:
    """
    Helper class that will count all requests made online.

    Might not be robust if urllib3 changes its logging format but should be good enough for us.

    Usage:
    ```py
    with RequestCounter() as counter:
        _ = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
    assert counter["GET"] == 0
    assert counter["HEAD"] == 1
    assert counter.total_calls == 1
    ```
    """

    def __enter__(self):
        self._counter = defaultdict(int)
        self._thread_id = threading.get_ident()
        self._extra_info = []

        def patched_with_thread_info(func):
            def wrap(*args, **kwargs):
                self._extra_info.append(threading.get_ident())
                return func(*args, **kwargs)

            return wrap

        self.patcher = patch.object(
            urllib3.connectionpool.log, "debug", side_effect=patched_with_thread_info(urllib3.connectionpool.log.debug)
        )
        self.mock = self.patcher.start()
        return self

    def __exit__(self, *args, **kwargs) -> None:
        assert len(self.mock.call_args_list) == len(self._extra_info)

        for thread_id, call in zip(self._extra_info, self.mock.call_args_list):
            if thread_id != self._thread_id:
                continue
            log = call.args[0] % call.args[1:]
            for method in ("HEAD", "GET", "POST", "PUT", "DELETE", "CONNECT", "OPTIONS", "TRACE", "PATCH"):
                if method in log:
                    self._counter[method] += 1
                    break
        self.patcher.stop()

    def __getitem__(self, key: str) -> int:
        return self._counter[key]

    @property
    def total_calls(self) -> int:


SUMMARY:
**RequestCounter Class Summary**

*   **Primary Purpose:** This class helps track all requests made online by counting the number of instances of HTTP methods (GET, HEAD, etc.) that are logged using the `urllib3` library.

**Key Implementation Details:**

*   Uses `defaultdict` to maintain a counter for each HTTP method.
*   Patchs the `debug` method of `urllib3.connectionpool.log` to log method names with thread ID information.
*   Verifies that the logged debug messages match expected behavior, incrementing the corresponding counter value.

**Notable Dependencies and Requirements:**

*   `defaultdict`
*   `threading`
*   `patch`

**Method Signatures and Interfaces:**

*   `__enter__`: Initializes the RequestCounter instance.
*   `__exit__`: Finalizes the request counting process with thread-specific logs.
*   `__getitem__`: Returns the count for a specific HTTP method.
*   `total_calls`: Provides the total number of requests made.


key: chroma:document
string_value:
CODE:

def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None):
    """
    To decorate flaky tests. They will be retried on failures.

    Args:
        max_attempts (`int`, *optional*, defaults to 5):
            The maximum number of attempts to retry the flaky test.
        wait_before_retry (`float`, *optional*):
            If provided, will wait that number of seconds before retrying the test.
        description (`str`, *optional*):
            A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors,
            etc.)
    """

    def decorator(test_func_ref):
        @functools.wraps(test_func_ref)
        def wrapper(*args, **kwargs):
            retry_count = 1

            while retry_count < max_attempts:
                try:
                    return test_func_ref(*args, **kwargs)

                except Exception as err:
                    print(f"Test failed with {err} at try {retry_count}/{max_attempts}.", file=sys.stderr)
                    if wait_before_retry is not None:
                        time.sleep(wait_before_retry)
                    retry_count += 1

            return test_func_ref(*args, **kwargs)

        return wrapper



SUMMARY:
**Function Summary: is_flaky Decorator**

The `is_flaky` function is a decorator that wraps a test function to implement flaky test retry mechanisms. It takes three optional parameters:
- `max_attempts`: The maximum number of attempts to retry the test (default 5).
- `wait_before_retry`: Time in seconds to wait before retrying the test (optional).
- `description`: A string describing the flaky situation (optional).

The decorator wraps a test function and retries it on failure, waiting for the specified duration if an exception occurs. If the maximum number of attempts is reached without success, the original test result is returned.

**Key Implementation Details:**

- Uses `functools.wraps` to preserve the test function's metadata.
- Inherent recursive retry behavior with infinite loop limitation.
- Waits for a specified amount before automatically retrying if wait_before_retry is provided.
 
**Notable Dependencies or Requirements:**

* The decorator uses the `time` module for waiting periods.
* It also assumes the presence of a `test_func_ref` to decorate, which could be another function definition not shown here.
 

Please note your RAG reference, in this case it is ' Decorator Function '.


key: chroma:document
string_value:
CODE:

def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None):
    """
    To run a test in a subprocess. In particular, this can avoid (GPU) memory issue.

    Args:
        test_case (`unittest.TestCase`):
            The test that will run `target_func`.
        target_func (`Callable`):
            The function implementing the actual testing logic.
        inputs (`dict`, *optional*, defaults to `None`):
            The inputs that will be passed to `target_func` through an (input) queue.
        timeout (`int`, *optional*, defaults to `None`):
            The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env.
            variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`.
    """
    if timeout is None:
        timeout = int(os.environ.get("PYTEST_TIMEOUT", 600))

    start_methohd = "spawn"
    ctx = multiprocessing.get_context(start_methohd)

    input_queue = ctx.Queue(1)
    output_queue = ctx.JoinableQueue(1)

    # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle.
    input_queue.put(inputs, timeout=timeout)

    process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout))
    process.start()
    # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents
    # the test to exit properly.
    try:
        results = output_queue.get(timeout=timeout)
        output_queue.task_done()
    except Exception as e:
        process.terminate()
        test_case.fail(e)
    process.join(timeout=timeout)

    if results["error"] is not None:


SUMMARY:
**Summary:**

`run_test_in_subprocess` function runs a test case in a subprocess to avoid memory issues, particularly with GPU. It takes a `unittest.TestCase`, a callable testing function (`target_func`), and optional input dictionary (`inputs`) and timeout value (`timeout`).

**Main Purpose and Functionality:**

* Runs `test_case` using `target_func` in a separate process
* Passes `inputs` to `target_func` through a queue with specified timeout
* Retrieves results from output queue, handling exceptions and timing out if necessary

**Key Implementation Details:**

* Uses `multiprocessing` library for inter-process communication (IPC)
* Creates input and output queues for IPC
* Starts a process using `ctx.Process`, passing `target_func` and queues as arguments

**Notable Dependencies or Requirements:**

* `unittest` for testing framework
* `multiprocessing` for process management
* `os` for environment variable access (`PYTEST_TIMEOUT`)


key: chroma:document
string_value:
CODE:

def run_test_using_subprocess(func):
    """
    To decorate a test to run in a subprocess using the `subprocess` module. This could avoid potential GPU memory
    issues (GPU OOM or a test that causes many subsequential failing with `CUDA error: device-side assert triggered`).
    """
    import pytest

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if os.getenv("_INSIDE_SUB_PROCESS", None) == "1":
            func(*args, **kwargs)
        else:
            test = " ".join(os.environ.get("PYTEST_CURRENT_TEST").split(" ")[:-1])
            try:
                import copy

                env = copy.deepcopy(os.environ)
                env["_INSIDE_SUB_PROCESS"] = "1"
                # This prevents the entries in `short test summary info` given by the subprocess being truncated. so the
                # full information can be passed to the parent pytest process.
                # See: https://docs.pytest.org/en/stable/explanation/ci.html
                env["CI"] = "true"

                # If not subclass of `unitTest.TestCase` and `pytestconfig` is used: try to grab and use the arguments
                if "pytestconfig" in kwargs:
                    command = list(kwargs["pytestconfig"].invocation_params.args)
                    for idx, x in enumerate(command):
                        if x in kwargs["pytestconfig"].args:
                            test = test.split("::")[1:]
                            command[idx] = "::".join([f"{func.__globals__['__file__']}"] + test)
                    command = [f"{sys.executable}", "-m", "pytest"] + command
                    command = [x for x in command if x not in ["--no-summary"]]
                # Otherwise, simply run the test with no option at all
                else:
                    command = [f"{sys.executable}", "-m", "pytest", f"{test}"]

                subprocess.run(command, env=env, check=True, capture_output=True)
            except subprocess.CalledProcessError as e:
                exception_message = e.stdout.decode()
                lines = exception_message.split("\n")
                # Add a first line with more informative information instead of just `= test session starts =`.
                # This makes the `short test summary info` section more useful.
                if "= test session starts =" in lines[0]:
                    text = ""
                    for line in lines[1:]:
                        if line.startswith("FAILED "):
                            text = line[len("FAILED ") :]
                            text = "".join(text.split(" - ")[1:])
                        elif line.startswith("=") and line.endswith("=") and " failed in " in line:
                            break
                        elif len(text) > 0:
                            text += f"\n{line}"
                    text = "(subprocess) " + text
                    lines = [text] + lines
                exception_message = "\n".join(lines)
                raise pytest.fail(exception_message, pytrace=False)



SUMMARY:
**Function Summary:**

`run_test_using_subprocess(func)`

**Purpose and Functionality:**

This function decorates a test to run in a subprocess using the `subprocess` module, utilizing pytest. Its primary purpose is to avoid potential GPU memory issues by running tests in separate processes.

**Key Implementation Details:**

* Creates a subprocess with the `subprocess.run()` function.
* Overwrites environment variables for the subprocess to include `_INSIDE_SUB_PROCESS` and `CI`.
* Uses `pytestconfig` to modify test arguments, if applicable.

**Notable Dependencies or Requirements:**

* pytest
* subprocess
* functools
* os
* sys

**Design Choices:**

* Wraps around the original function with a wrapper using `functools.wraps()`.
* Uses a environment dictionary (`env`) to modify and pass environment variables to the subprocess.

This decorated function can help avoid GPU memory issues by running tests in separate processes.


key: docstring
string_value:
"""
The following contains utils to run the documentation tests without having to overwrite any files.

The `preprocess_string` function adds `# doctest: +IGNORE_RESULT` markers on the fly anywhere a `load_dataset` call is
made as a print would otherwise fail the corresonding line.

To skip cuda tests, make sure to call `SKIP_CUDA_DOCTEST=1 pytest --doctest-modules <path_to_files_to_test>
"""


key: chroma:document
string_value:
CODE:

def preprocess_string(string, skip_cuda_tests):
    """Prepare a docstring or a `.md` file to be run by doctest.

    The argument `string` would be the whole file content if it is a `.md` file. For a python file, it would be one of
    its docstring. In each case, it may contain multiple python code examples. If `skip_cuda_tests` is `True` and a
    cuda stuff is detective (with a heuristic), this method will return an empty string so no doctest will be run for
    `string`.
    """
    codeblock_pattern = r"(```(?:python|py)\s*\n\s*>>> )((?:.*?\n)*?.*?```)"
    codeblocks = re.split(re.compile(codeblock_pattern, flags=re.MULTILINE | re.DOTALL), string)
    is_cuda_found = False
    for i, codeblock in enumerate(codeblocks):
        if "load_dataset(" in codeblock and "# doctest: +IGNORE_RESULT" not in codeblock:
            codeblocks[i] = re.sub(r"(>>> .*load_dataset\(.*)", r"\1 # doctest: +IGNORE_RESULT", codeblock)
        if (
            (">>>" in codeblock or "..." in codeblock)
            and re.search(r"cuda|to\(0\)|device=0", codeblock)
            and skip_cuda_tests
        ):
            is_cuda_found = True
            break

    modified_string = ""
    if not is_cuda_found:
        modified_string = "".join(codeblocks)



SUMMARY:
**Function Summary: `preprocess_string`**

* **Purpose**: Prepare a docstring or markdown file for running doctest tests, skipping CUDA-related code if enabled.

**Key Implementation Details:**

* Uses regular expressions to identify and split codeblocks from input string (`codeblock_pattern`, `re.split`)
* Applies modifications to codeblocks (adding `# doctest: +IGNORE_RESULT` markers) based on presence of CUDA-related keywords in skip_cuda_tests mode
* Checks for existence of CUDA-related code by scanning through each codeblock with a specific heuristic

**Notable Dependencies and Requirements:**

* Python re module (`re.split`, `re.compile`)
* No external dependencies listed, but assumed to be available due to usage in pytest-based testing context


key: chroma:document
string_value:
CODE:

class HfDocTestParser(doctest.DocTestParser):
    """
    Overwrites the DocTestParser from doctest to properly parse the codeblocks that are formatted with black. This
    means that there are no extra lines at the end of our snippets. The `# doctest: +IGNORE_RESULT` marker is also
    added anywhere a `load_dataset` call is made as a print would otherwise fail the corresponding line.

    Tests involving cuda are skipped base on a naive pattern that should be updated if it is not enough.
    """

    # This regular expression is used to find doctest examples in a
    # string.  It defines three groups: `source` is the source code
    # (including leading indentation and prompts); `indent` is the
    # indentation of the first (PS1) line of the source code; and
    # `want` is the expected output (including leading indentation).
    # fmt: off
    _EXAMPLE_RE = re.compile(r'''
        # Source consists of a PS1 line followed by zero or more PS2 lines.
        (?P<source>
            (?:^(?P<indent> [ ]*) >>>    .*)    # PS1 line
            (?:\n           [ ]*  \.\.\. .*)*)  # PS2 lines
        \n?
        # Want consists of any non-blank lines that do not start with PS1.
        (?P<want> (?:(?![ ]*$)    # Not a blank line
             (?![ ]*>>>)          # Not a line starting with PS1
             # !!!!!!!!!!! HF Specific !!!!!!!!!!!
             (?:(?!```).)*        # Match any character except '`' until a '```' is found (this is specific to HF because black removes the last line)
             # !!!!!!!!!!! HF Specific !!!!!!!!!!!
             (?:\n|$)  # Match a new line or end of string
          )*)
        ''', re.MULTILINE | re.VERBOSE
    )
    # fmt: on

    # !!!!!!!!!!! HF Specific !!!!!!!!!!!
    skip_cuda_tests: bool = bool(os.environ.get("SKIP_CUDA_DOCTEST", False))
    # !!!!!!!!!!! HF Specific !!!!!!!!!!!

    def parse(self, string, name="<string>"):
        """
        Overwrites the `parse` method to incorporate a skip for CUDA tests, and remove logs and dataset prints before
        calling `super().parse`
        """
        string = preprocess_string(string, self.skip_cuda_tests)


SUMMARY:
**Class HfDocTestParser Summary**

### Purpose and Functionality

The `HfDocTestParser` class overrides the default `DocTestParser` from `doctest` to parse codeblocks formatted with black. It ensures no extra lines at the end of snippets and ignores CUDA tests based on an environment variable.

### Key Implementation Details

* Uses a regular expression `_EXAMPLE_RE` to extract source, indent, and want parts of doctest examples.
* Introduces skip for CUDA tests via `skip_cuda_tests` flag.
* Parses input string using `preprocess_string` method and calls `super().parse` with modified string.

### Notable Dependencies and Requirements

* Re dependency on `black`
* External dependencies:
	+ `os` module (for environment variable access)

### Class Structure and Interfaces

* Inherited from `docetest.DocTestParser`
* Implements custom parsing logic for codeblocks formatted with black.


key: chroma:document
string_value:
CODE:

class HfDoctestModule(Module):
    """
    Overwrites the `DoctestModule` of the pytest package to make sure the HFDocTestParser is used when discovering
    tests.
    """

    def collect(self) -> Iterable[DoctestItem]:
        class MockAwareDocTestFinder(doctest.DocTestFinder):
            """A hackish doctest finder that overrides stdlib internals to fix a stdlib bug.

            https://github.com/pytest-dev/pytest/issues/3456 https://bugs.python.org/issue25532
            """

            def _find_lineno(self, obj, source_lines):
                """Doctest code does not take into account `@property`, this
                is a hackish way to fix it. https://bugs.python.org/issue17446

                Wrapped Doctests will need to be unwrapped so the correct line number is returned. This will be
                reported upstream. #8796
                """
                if isinstance(obj, property):
                    obj = getattr(obj, "fget", obj)

                if hasattr(obj, "__wrapped__"):
                    # Get the main obj in case of it being wrapped
                    obj = inspect.unwrap(obj)

                # Type ignored because this is a private function.
                return super()._find_lineno(  # type:ignore[misc]
                    obj,
                    source_lines,
                )

            def _find(self, tests, obj, name, module, source_lines, globs, seen) -> None:
                if _is_mocked(obj):
                    return
                with _patch_unwrap_mock_aware():
                    # Type ignored because this is a private function.
                    super()._find(  # type:ignore[misc]
                        tests, obj, name, module, source_lines, globs, seen
                    )

        if self.path.name == "conftest.py":
            module = self.config.pluginmanager._importconftest(
                self.path,
                self.config.getoption("importmode"),
                rootpath=self.config.rootpath,
            )
        else:
            try:
                module = import_path(
                    self.path,
                    root=self.config.rootpath,
                    mode=self.config.getoption("importmode"),
                )
            except ImportError:
                if self.config.getvalue("doctest_ignore_import_errors"):
                    skip("unable to import module %r" % self.path)
                else:
                    raise

        # !!!!!!!!!!! HF Specific !!!!!!!!!!!
        finder = MockAwareDocTestFinder(parser=HfDocTestParser())
        # !!!!!!!!!!! HF Specific !!!!!!!!!!!
        optionflags = get_optionflags(self)
        runner = _get_runner(
            verbose=False,
            optionflags=optionflags,
            checker=_get_checker(),
            continue_on_failure=_get_continue_on_failure(self.config),
        )
        for test in finder.find(module, module.__name__):
            if test.examples:  # skip empty doctests and cuda


SUMMARY:
**Class Analysis**

`HfDoctestModule(Module)`

*   **Primary Purpose**: Overwrites the `DoctestModule` of the pytest package to use `HFDocTestParser` when discovering tests. 
*   
*   **Implementation Details**
    *   Uses a custom doctest finder (`MockAwareDocTestFinder`) to bypass stdlib bug and correctly resolve property usage in doctests.
    *   Applies HF-specific patches, including wrapping Doctests to fix line number issues with Python's `property` keyword. 

**Notable Dependencies and Requirements**

*   `HfDocTestParser`: Used by the custom doctest finder (`MockAwareDocTestFinder`) to correctly parse and process doctests.
    *   References: `HFDocTestParser`

*   External Dependencies
    + `doctest`
        - Python module for structured testing. 
    + `inspect` 
        - Module providing functions to inspect live objects, such as getting the source code or line number of a function.

**Key Method Signatures and Interfaces**

*   `_find_lineno`: Method from `MockAwareDocTestFinder` class.
*   `_is_mocked`
*   `_patch_unwrap_mock_aware`

        - Used for wrapping objects to unwrap property types 
    + `_find`
        - A private method of `MockAwareDocTestFinder` class.

        **Note:** Some details have been omitted as per the original request focusing on key implementation details in an easy-understand manner.


key: chroma:document
string_value:
CODE:

def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs):
    if device not in dispatch_table:
        return dispatch_table["default"](*args, **kwargs)

    fn = dispatch_table[device]

    # Some device agnostic functions return values. Need to guard against `None`
    # instead at user level.
    if fn is None:
        return None


SUMMARY:
**Summary:**

 `_device_agnostic_dispatch` function implements a dynamic dispatch mechanism for device-agnostic functionality. It takes a `device`, `dispatch_table`, and variable arguments/keyword arguments, then calls the corresponding function from the table while ensuring default behavior for non-existent devices.

**Key Implementation Details:**

* The function utilizes a dictionary (`dispatch_table`) to store callable functions for each device.
* A check is performed to ensure `default` function returns if a specified device is not present in the table.
* The function calls the selected function with variable arguments/keyword arguments while preventing potential `None` values from being returned.

**Notable Dependencies:**

* `dispatch_table`: a dictionary of callable functions for each device
* `Callable`: refers to an interface that defines call operation

These details are presented in this concise summary, avoiding unnecessary information and focusing on the essential functionality.


key: chroma:document
string_value:
CODE:

def backend_manual_seed(device: str, seed: int):


SUMMARY:
**Summary**

This is the `backend_manual_seed` function, which dispatches to an underlying engine for manual seeding of a data storage device. It takes two inputs: `device` and `seed`, and returns the result of the dispatch using `_device_agnostic_dispatch`.

**Key Details**

- Purpose: Manual device seeding
- Primary Signature: `backend_manual_seed(device: str, seed: int) -> ...`
- Notable Dispatch to: `_device_agnostic_dispatch`

**Dependencies/Requirements**

* `_device_agnostic_dispatch` (function)
        The function is likely part of a backend engine, possibly supporting multiple devices.
* `BACKEND_MANUAL_SEED`: Enum or constant representing manual seeding mode.


key: chroma:document
string_value:
CODE:

def backend_empty_cache(device: str):


SUMMARY:
**Function Summary**

`backend_empty_cache` : Clears the cache for a specific device by dispatching to an underlying device-agnostic function.

**Implementation Details**

* Dispatches to `_device_agnostic_dispatch` with `BACKEND_EMPTY_CACHE` as an argument
* No additional implementation details provided

**Notable Dependencies/Requirements**

* `_device_agnostic_dispatch`
	+ References related class or function: `_device_agnostic_dispatch`
	+ External Dependency: None mentioned


key: chroma:document
string_value:
CODE:

def backend_device_count(device: str):


SUMMARY:
**Summary:**

The `backend_device_count` function maps the device type as input to its corresponding backend count. It utilizes the `_device_agnostic_dispatch` dispatching mechanism to determine the appropriate backend implementation.

Key details:

* Input: Device type (string)
* Returns: Backend count
* Implementation details hidden behind `_device_agnostic_dispatch` proxy

Notable dependencies/requirements:
- `BACKEND_DEVICE_COUNT`: External constant or configuration, likely related to available backend devices.


key: chroma:document
string_value:
CODE:

def compare_pipeline_output_to_hub_spec(output, hub_spec):
    missing_keys = []
    unexpected_keys = []
    all_field_names = {field.name for field in fields(hub_spec)}
    matching_keys = sorted([key for key in output.keys() if key in all_field_names])

    # Fields with a MISSING default are required and must be in the output
    for field in fields(hub_spec):
        if field.default is MISSING and field.name not in output:
            missing_keys.append(field.name)

    # All output keys must match either a required or optional field in the Hub spec
    for output_key in output:
        if output_key not in all_field_names:
            unexpected_keys.append(output_key)

    if missing_keys or unexpected_keys:
        error = ["Pipeline output does not match Hub spec!"]
        if matching_keys:
            error.append(f"Matching keys: {matching_keys}")
        if missing_keys:
            error.append(f"Missing required keys in pipeline output: {missing_keys}")
        if unexpected_keys:
            error.append(f"Keys in pipeline output that are not in Hub spec: {unexpected_keys}")


SUMMARY:
**Summarized Code Analysis**

### Function Purpose and Functionality
`compare_pipeline_output_to_hub_spec(output, hub_spec)`: Compares the pipeline output with the Hub specification to identify missing or unexpected keys and raises a `KeyError` if discrepancies are found.

### Key Implementation Details

*   Filters matching fields between `output` and `hub_spec.keys()`.
*   Checks for required fields in `missing_keys` when their default is `MISSING` and they don't exist in `output`.
*   Verifies all output keys against the corresponding fields in `hub_spec`.

### Notable Dependencies or Requirements

None.

### Interface and Method Signatures

*   Takes two parameters: `output` and `hub_spec`.
*   Returns no value (`None`) if the comparison is successful; raises a `KeyError` otherwise.
*   Utilizes the `fields()` function from the `dataclasses` module (not shown in this summary).
*   Implements the `__iter__()` interface from the `dict` type (not shown in this summary).


key: chroma:document
string_value:
CODE:

class HyperParamSearchBackendBase:
    name: str
    pip_package: str = None

    @staticmethod
    def is_available():
        raise NotImplementedError

    def run(self, trainer, n_trials: int, direction: str, **kwargs):
        raise NotImplementedError

    def default_hp_space(self, trial):
        raise NotImplementedError

    def ensure_available(self):
        if not self.is_available():
            raise RuntimeError(
                f"You picked the {self.name} backend, but it is not installed. Run {self.pip_install()}."
            )

    @classmethod
    def pip_install(cls):


SUMMARY:
**Summary:**

The `HyperParamSearchBackendBase` class is the base class for hyperparameter search backends. It provides basic functionality and interfaces for registering, installing, and running hyperparameter tuning processes.

**Key Features:**

* **Registration**: The class maintains attributes for storing hyperparameter search backend details (name and pip package).
* **Installation Check**: The `ensure_available` method checks if the backend is available to run by calling `is_available()`. If not installed, it installs using `pip_install()` in a runtime context (`self.pip_install()`).
* **Hyperparameter Tuning Interface**:
	+ `is_available()`: Raised as an abstract method; must be implemented by child classes.
	+ `run(self, trainer, n_trials: int, direction: str)`: Raised as an abstract method; requires implementation to define the hyperparameter tuning process.
	+ `default_hp_space(trial)`: Raised as an abstract method; should return default hyperparameter space for a given trial.

**Notable Dependencies and Requirements:**

* The class depends on `pip` package installed with the name specified as `pip_package` or `name`.
* External dependency on `numpy`, not explicitly listed, assumed for potential usage in hyperparameter search process. 

The provided class serves as an abstraction layer to handle different backend implementations and facilitates easy switching between available backends by providing default hp space through abstract method calls.


key: chroma:document
string_value:
CODE:

class OptunaBackend(HyperParamSearchBackendBase):
    name = "optuna"

    @staticmethod
    def is_available():
        return is_optuna_available()

    def run(self, trainer, n_trials: int, direction: str, **kwargs):
        return run_hp_search_optuna(trainer, n_trials, direction, **kwargs)

    def default_hp_space(self, trial):


SUMMARY:
**OptunaBackend Class Summary**

**Main Purpose and Functionality:**
The `OptunaBackend` class is a HyperParamSearchBackend implementation that utilizes Optuna for Bayesian optimization. It provides methods to determine availability, run searches, and define the default hyperparameter space.

**Key Implementation Details:**

*   **is_available() method:** Checks if Optuna is available using the `is_optuna_available()` function.
*   **run() method:** Runs an Optuna-based Bayesian optimization search with specified parameters (`n_trials` and `direction`) using `run_hp_search_optuna(trainer, n_trials, direction, ...)`.

**Notable Dependencies or Requirements:**

*   References:
    *   `is_optuna_available`: a function that checks if Optuna is installed
    *   `run_hp_search_optuna`, `default_hp_space_optuna`
*   External Dependencies:
    + [Optuna](https://optuna.com/) (Python library for Bayesian optimization)
*   Implemented Interfaces:
    None


key: chroma:document
string_value:
CODE:

class RayTuneBackend(HyperParamSearchBackendBase):
    name = "ray"
    pip_package = "'ray[tune]'"

    @staticmethod
    def is_available():
        return is_ray_tune_available()

    def run(self, trainer, n_trials: int, direction: str, **kwargs):
        return run_hp_search_ray(trainer, n_trials, direction, **kwargs)

    def default_hp_space(self, trial):


SUMMARY:
**Summary**

The `RayTuneBackend` class is a HyperParamSearch backend that utilizes the Ray Tune library. Its primary purpose is to perform high-performance hypersparameter tuning using distributed computing.

**Main Purpose and Functionality**

This class provides a way to search for optimal hyperparameters by leveraging Ray Tune's capabilities, specifically targeting the "ray" package. The main functionality includes:

* Checking availability of Ray Tune.
* Running a hypersparameter search with specified trials and direction.
* Handling default search space generation using Ray-based implementation.

**Implementation Details**

Key details include:

* `run_hp_search_ray`, `default_hp_space_ray` methods, which perform the actual search and generate default spaces.
* Dependency on "ray[tune]" package for distributed computing.

**Notable Dependencies and Requirements**

* `ray_tune_available()` function (not shown in this code snippet).
* External dependency: `"ray[tune]"` package.


key: chroma:document
string_value:
CODE:

class SigOptBackend(HyperParamSearchBackendBase):
    name = "sigopt"

    @staticmethod
    def is_available():
        return is_sigopt_available()

    def run(self, trainer, n_trials: int, direction: str, **kwargs):
        return run_hp_search_sigopt(trainer, n_trials, direction, **kwargs)

    def default_hp_space(self, trial):


SUMMARY:
**Summary: SigOptBackend Class**

The `SigOptBackend` class is a custom search backend for hyperparameter tuning using Bayesian optimization. It inherits from `HyperParamSearchBackendBase`.

**Main Purpose and Functionality**

* Provides a standardized interface for Hyperopt-style Bayesian optimization with SigOpt (a cloud-based optimization platform).

**Key Implementation Details**

* Reuses the following functions:
	+ `run_hp_search_sigopt`
	+ `default_hp_space_sigopt`

**Notable Dependencies or Requirements**

* Uses SigOpt's `is_sigopt_available` function to determine availability, which is a separate module.
* Requires the `SigOpt` platform for cloud-based optimization.

This class simplifies integration with SigOpt by providing a dedicated backend for hyperparameter tuning. Developers can leverage the existing `HyperParamSearchBackendBase` framework while leveraging SigOpt's capabilities for Bayesian optimization.


key: chroma:document
string_value:
CODE:

class WandbBackend(HyperParamSearchBackendBase):
    name = "wandb"

    @staticmethod
    def is_available():
        return is_wandb_available()

    def run(self, trainer, n_trials: int, direction: str, **kwargs):
        return run_hp_search_wandb(trainer, n_trials, direction, **kwargs)

    def default_hp_space(self, trial):


SUMMARY:
**Summary**

The `WandbBackend` class is a HyperParameterSearchBackendBase implementation that utilizes the WANDB (Weasel) library for hyperparameter search.

**Main Purpose and Functionality**
---------------------------

* Utilizes WANDB for hyperparameter search
* Facilitates collaboration with a shared dashboard

**Key Implementation Details**
-----------------------------

* Leverages `is_wandb_available()` static method to check availability of WANDB
* Extends HyperParamSearchBackendBase, implying the use of `run_hp_search_wandb` and other related methods
* Implements a default `default_hp_space_wandb` to generate the initial parameter space

**Notable Dependencies or Requirements**
-----------------------------------------

* Relies on WANDB library (`is_wandb_available()`, `run_hp_search_wandb`, `default_hp_space_wandb`)
* Uses HyperParamSearchBackendBase as its base class
* Expects `trainers` to be passed with built-in instance of the `WandbBackend`.


key: chroma:document
string_value:
CODE:

def default_hp_search_backend() -> str:
    available_backends = [backend for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() if backend.is_available()]
    if len(available_backends) > 0:
        name = available_backends[0].name
        if len(available_backends) > 1:
            logger.info(
                f"{len(available_backends)} hyperparameter search backends available. Using {name} as the default."
            )
        return name
    raise RuntimeError(
        "No hyperparameter search backend available.\n"
        + "\n".join(
            f" - To install {backend.name} run {backend.pip_install()}"
            for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values()
        )


SUMMARY:
**Summary: `default_hp_search_backend` Function**

```markdown
### Purpose
Automatically determines the default hyperparameter search backend to use.

### Implementation Details

*   Retrieves a list of available backends from `ALL_HYPERPARAMETER_SEARCH_BACKENDS`.
*   If multiple backends are available, logs a message indicating this and uses the first one.
*   If only one backend is available or no backends are found, raises a `RuntimeError` with installation instructions for each affected backend.

### Notable Dependencies

*   `ALL_HYPERPARAMETER_SEARCH_BACKENDS`: A dictionary mapping backend names to their respective instances.
    *   The dictionary contains the following keys/values:
        - `backend.name`
        - `backend.is_available()`
        - `backend.pip_install()`

### External Requirements

No external dependencies are listed, but these may be used internally by the backends (e.g., `pip` package manager).

### Primary Method Signatures and Interfaces

*   `backend.is_available() -> bool`: Checks if a backend is available.
*   `backend.name -> str`: Returns the name of a backend instance.
*   `backend.pip_install() -> str`: Generates the pip install command for a given backend.
```

This code provides an additional layer of flexibility in selecting a hyperparameter search backend when none are available by default. If multiple backends can be utilized, it selects the first one but still maintains log information to inform about the selection process. The method raises a notable exception with explicit commands to install all relevant backends if they are discovered to be unavailable, making recovery and upgrade easier for developers or users who may have forgotten installing those packages during configuration or deployment stages in your application.


key: chroma:document
string_value:
CODE:

def _is_peft_model(model):
    if is_peft_available():
        classes_to_check = (PeftModel,) if is_peft_available() else ()
        # Here we also check if the model is an instance of `PeftMixedModel` introduced in peft>=0.7.0: https://github.com/huggingface/transformers/pull/28321
        if version.parse(importlib.metadata.version("peft")) >= version.parse("0.7.0"):
            from peft import PeftMixedModel

            classes_to_check = (*classes_to_check, PeftMixedModel)
        return isinstance(model, classes_to_check)


SUMMARY:
Here is a concise and focused summary of the function code chunk:

**Description:** Checks if a given model instance conforms to the `PeftModel` class or its subclass `PeftMixedModel` based on Peft version availability.

```markdown
### Function
`_is_peft_model(model)`: Returns True if `model` is an instance of `PeftModel` or `PeftMixedModel`, False otherwise.

### Conditions

- If Peft is available.
- Based on the Peft version (>= 0.7.0).

### Return Value
Boolean indicating model compatibility with Peft classes.
```

**Key Implementation Details:**

* Uses `is_peft_available()` to enable or disable `PeftModel` class checking.
* Accounts for `PeftMixedModel` introduction in Peft >= 0.7.0.

**Notable Dependencies/Requirements:**

* `peft` library (version >= 0.7.0)
* `transformers` library (referenced but not used directly)


key: chroma:document
string_value:
CODE:

def _get_fsdp_ckpt_kwargs():
    # TODO: @AjayP13, @younesbelkada replace this check with version check at the next `accelerate` release
    if is_accelerate_available() and "adapter_only" in list(inspect.signature(save_fsdp_model).parameters):
        return {"adapter_only": True}
    else:


SUMMARY:
**Summary**

The `_get_fsdp_ckpt_kwargs` function generates keyword arguments (kwargs) for saving a model. It checks if the `accelerate` library is available and if the `save_fsdp_model` function supports the `"adapter_only"` parameter. If `accelerate` is available, it returns `{"adapter_only": True}`, otherwise an empty dictionary.

**Key aspects:**

* Purpose: Generate kwargs for saving a model
* Key implementation details:
	+ Uses `inspect.signature` to check if `save_fsdp_model` supports `"adapter_only"`
	+ Availability of `accelerate` library is checked
* Notable dependencies/requirements:
	+ `inspect`
	+ `accelerate` (library)
	+ `save_fsdp_model` function


key: chroma:document
string_value:
CODE:

def safe_globals():
    # Starting from version 2.4 PyTorch introduces a check for the objects loaded
    # with torch.load(weights_only=True). Starting from 2.6 weights_only=True becomes
    # a default and requires allowlisting of objects being loaded.
    # See: https://github.com/pytorch/pytorch/pull/137602
    # See: https://pytorch.org/docs/stable/notes/serialization.html#torch.serialization.add_safe_globals
    # See: https://github.com/huggingface/accelerate/pull/3036
    if version.parse(torch.__version__).release < version.parse("2.6").release:
        return contextlib.nullcontext()

    np_core = np._core if version.parse(np.__version__) >= version.parse("2.0.0") else np.core
    allowlist = [np_core.multiarray._reconstruct, np.ndarray, np.dtype]
    # numpy >1.25 defines numpy.dtypes.UInt32DType, but below works for
    # all versions of numpy
    allowlist += [type(np.dtype(np.uint32))]



SUMMARY:
**Function Summary: `safe_globals()`**

`safeoglobales()` is a function that returns a context manager for safely loading PyTorch weights. Its primary purpose is to prevent potential security issues with dynamic imports.

**Key Implementation Details:**

* Uses Python version checks to determine if the `torch` and NumPy versions meet specific requirements.
* Allowing list generation, including:
	+ NumPy multiarray reconstruction functions
	+ NumPy arrays and data types (`np.ndarray`, `np.dtype`)
	+ Certain NumPy data type wrappers

**Notable Dependencies and Requirements:**

* PyTorch >= 2.6
* NumPy >= 2.0.0 (for multi-array reconstruction)
* External Dependency:
	+ Python's `torch.serialization` module


key: chroma:document
string_value:
CODE:

class Trainer:
    """
    Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers.

    Args:
        model ([`PreTrainedModel`] or `torch.nn.Module`, *optional*):
            The model to train, evaluate or use for predictions. If not provided, a `model_init` must be passed.

            <Tip>

            [`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use
            your own models defined as `torch.nn.Module` as long as they work the same way as the 🤗 Transformers
            models.

            </Tip>

        args ([`TrainingArguments`], *optional*):
            The arguments to tweak for training. Will default to a basic instance of [`TrainingArguments`] with the
            `output_dir` set to a directory named *tmp_trainer* in the current directory if not provided.
        data_collator (`DataCollator`, *optional*):
            The function to use to form a batch from a list of elements of `train_dataset` or `eval_dataset`. Will
            default to [`default_data_collator`] if no `processing_class` is provided, an instance of
            [`DataCollatorWithPadding`] otherwise if the processing_class is a feature extractor or tokenizer.
        train_dataset (Union[`torch.utils.data.Dataset`, `torch.utils.data.IterableDataset`, `datasets.Dataset`], *optional*):
            The dataset to use for training. If it is a [`~datasets.Dataset`], columns not accepted by the
            `model.forward()` method are automatically removed.

            Note that if it's a `torch.utils.data.IterableDataset` with some randomization and you are training in a
            distributed fashion, your iterable dataset should either use a internal attribute `generator` that is a
            `torch.Generator` for the randomization that must be identical on all processes (and the Trainer will
            manually set the seed of this `generator` at each epoch) or have a `set_epoch()` method that internally
            sets the seed of the RNGs used.
        eval_dataset (Union[`torch.utils.data.Dataset`, Dict[str, `torch.utils.data.Dataset`, `datasets.Dataset`]), *optional*):
             The dataset to use for evaluation. If it is a [`~datasets.Dataset`], columns not accepted by the
             `model.forward()` method are automatically removed. If it is a dictionary, it will evaluate on each
             dataset prepending the dictionary key to the metric name.
        processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*):
            Processing class used to process the data. If provided, will be used to automatically process the inputs
            for the model, and it will be saved along the model to make it easier to rerun an interrupted training or
            reuse the fine-tuned model.
            This supercedes the `tokenizer` argument, which is now deprecated.
        model_init (`Callable[[], PreTrainedModel]`, *optional*):
            A function that instantiates the model to be used. If provided, each call to [`~Trainer.train`] will start
            from a new instance of the model as given by this function.

            The function may have zero argument, or a single one containing the optuna/Ray Tune/SigOpt trial object, to
            be able to choose different architectures according to hyper parameters (such as layer count, sizes of
            inner layers, dropout probabilities etc).
        compute_loss_func (`Callable`, *optional*):
            A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated
            batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default [loss function](https://github.com/huggingface/transformers/blob/052e652d6d53c2b26ffde87e039b723949a53493/src/transformers/trainer.py#L3618) used by [`Trainer`].
        compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*):
            The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return
            a dictionary string to metric values. *Note* When passing TrainingArgs with `batch_eval_metrics` set to
            `True`, your compute_metrics function must take a boolean `compute_result` argument. This will be triggered
            after the last eval batch to signal that the function needs to calculate and return the global summary
            statistics rather than accumulating the batch-level statistics
        callbacks (List of [`TrainerCallback`], *optional*):
            A list of callbacks to customize the training loop. Will add those to the list of default callbacks
            detailed in [here](callback).

            If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method.
        optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`):
            A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your
            model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`.
        optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*):
            A tuple containing the optimizer class and keyword arguments to use.
            Overrides `optim` and `optim_args` in `args`. Incompatible with the `optimizers` argument.

            Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before initializing the Trainer.
        preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*):
            A function that preprocess the logits right before caching them at each evaluation step. Must take two
            tensors, the logits and the labels, and return the logits once processed as desired. The modifications made
            by this function will be reflected in the predictions received by `compute_metrics`.

            Note that the labels (second parameter) will be `None` if the dataset does not have them.

    Important attributes:

        - **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`]
          subclass.
        - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the
          original model. This is the model that should be used for the forward pass. For example, under `DeepSpeed`,
          the inner model is wrapped in `DeepSpeed` and then again in `torch.nn.DistributedDataParallel`. If the inner
          model hasn't been wrapped, then `self.model_wrapped` is the same as `self.model`.
        - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from
          data parallelism, this means some of the model layers are split on different GPUs).
        - **place_model_on_device** -- Whether or not to automatically place the model on the device - it will be set
          to `False` if model parallel or deepspeed is used, or if the default
          `TrainingArguments.place_model_on_device` is overridden to return `False` .
        - **is_in_train** -- Whether or not a model is currently running `train` (e.g. when `evaluate` is called while
          in `train`)

    """

    # Those are used as methods of the Trainer in examples.
    from .trainer_pt_utils import _get_learning_rate, log_metrics, metrics_format, save_metrics, save_state

    @deprecate_kwarg("tokenizer", new_name="processing_class", version="5.0.0", raise_if_both_names=True)
    def __init__(
        self,
        model: Union[PreTrainedModel, nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset], "datasets.Dataset"]] = None,
        processing_class: Optional[
            Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin]
        ] = None,
        model_init: Optional[Callable[[], PreTrainedModel]] = None,
        compute_loss_func: Optional[Callable] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None),
        optimizer_cls_and_kwargs: Optional[Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]] = None,
        preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
    ):
        if args is None:
            output_dir = "tmp_trainer"
            logger.info(f"No `TrainingArguments` passed, using `output_dir={output_dir}`.")
            args = TrainingArguments(output_dir=output_dir)
        if args.batch_eval_metrics and compute_metrics is not None:
            if "compute_result" not in inspect.signature(compute_metrics).parameters.keys():
                raise ValueError(
                    "When using `batch_eval_metrics`, your `compute_metrics` function must take a `compute_result`"
                    " boolean argument which will be triggered after the last batch of the eval set to signal that the"
                    " summary statistics should be returned by the function."
                )
        if args.eval_strategy is not None and args.eval_strategy != "no" and eval_dataset is None:
            raise ValueError(
                f"You have set `args.eval_strategy` to {args.eval_strategy} but you didn't pass an `eval_dataset` to `Trainer`. Either set `args.eval_strategy` to `no` or pass an `eval_dataset`. "
            )
        if args.save_strategy == SaveStrategy.BEST or args.load_best_model_at_end:
            if args.metric_for_best_model is None:
                raise ValueError(
                    "`args.metric_for_best_model` must be provided when using 'best' save_strategy or if `args.load_best_model_at_end` is set to `True`."
                )

        self.args = args
        self.compute_loss_func = compute_loss_func
        # Seed must be set before instantiating the model when using model
        enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed)

        self.hp_name = None
        self.deepspeed = None
        self.is_in_train = False

        self.create_accelerator_and_postprocess()

        # memory metrics - must set up as early as possible
        self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics)
        self._memory_tracker.start()

        # set the correct log level depending on the node
        log_level = args.get_process_log_level()
        logging.set_verbosity(log_level)

        # force device and distributed setup init explicitly
        args._setup_devices

        if model is None:
            if model_init is not None:
                self.model_init = model_init
                model = self.call_model_init()
            else:
                raise RuntimeError("`Trainer` requires either a `model` or `model_init` argument")
        else:
            if model_init is not None:
                warnings.warn(
                    "`Trainer` requires either a `model` or `model_init` argument, but not both. `model_init` will"
                    " overwrite your model when calling the `train` method. This will become a fatal error in the next"
                    " release.",
                    FutureWarning,
                )
            self.model_init = model_init

        if model.__class__.__name__ in MODEL_MAPPING_NAMES:
            raise ValueError(
                f"The model you have picked ({model.__class__.__name__}) cannot be used as is for training: it only "
                "computes hidden states and does not accept any labels. You should choose a model with a head "
                "suitable for your task like any of the `AutoModelForXxx` listed at "
                "https://huggingface.co/docs/transformers/model_doc/auto"
            )

        if getattr(model, "is_parallelizable", False) and getattr(model, "model_parallel", False):
            self.is_model_parallel = True
        else:
            self.is_model_parallel = False

        if getattr(model, "hf_device_map", None) is not None:
            devices = [device for device in set(model.hf_device_map.values()) if device not in ["cpu", "disk"]]
            if len(devices) > 1:
                self.is_model_parallel = True
            elif len(devices) == 1:
                self.is_model_parallel = self.args.device != torch.device(devices[0])
            else:
                self.is_model_parallel = False

            # warn users
            if self.is_model_parallel:
                logger.info(
                    "You have loaded a model on multiple GPUs. `is_model_parallel` attribute will be force-set"
                    " to `True` to avoid any unexpected behavior such as device placement mismatching."
                )

        if self.args.use_liger_kernel:
            if is_liger_kernel_available():
                from liger_kernel.transformers import _apply_liger_kernel_to_instance

                if isinstance(model, PreTrainedModel):
                    # Patch the model with liger kernels. Use the default kernel configurations.
                    _apply_liger_kernel_to_instance(model=model)
                else:
                    logger.warning(
                        "The model is not an instance of PreTrainedModel. No liger kernels will be applied."
                    )
            else:
                raise ImportError(
                    "You have set `use_liger_kernel` to `True` but liger-kernel >= 0.3.0 is not available. "
                    "Please install it with `pip install liger-kernel`"
                )

        _is_quantized_and_base_model = getattr(model, "is_quantized", False) and not getattr(
            model, "_hf_peft_config_loaded", False
        )
        _quantization_method_supports_training = (
            getattr(model, "hf_quantizer", None) is not None and model.hf_quantizer.is_trainable
        )

        _is_model_quantized_and_qat_trainable = getattr(model, "hf_quantizer", None) is not None and getattr(
            model.hf_quantizer, "is_qat_trainable", False
        )

        # Filter out quantized + compiled models
        if _is_quantized_and_base_model and hasattr(model, "_orig_mod"):
            raise ValueError(
                "You cannot fine-tune quantized model with `torch.compile()` make sure to pass a non-compiled model when fine-tuning a quantized model with PEFT"
            )

        # At this stage the model is already loaded
        if _is_quantized_and_base_model and not _is_peft_model(model) and not _is_model_quantized_and_qat_trainable:
            raise ValueError(
                "You cannot perform fine-tuning on purely quantized models. Please attach trainable adapters on top of"
                " the quantized model to correctly perform fine-tuning. Please see: https://huggingface.co/docs/transformers/peft"
                " for more details"
            )
        elif _is_quantized_and_base_model and not _quantization_method_supports_training:
            raise ValueError(
                f"The model you are trying to fine-tune is quantized with {model.hf_quantizer.quantization_config.quant_method}"
                " but that quantization method do not support training. Please open an issue on GitHub: https://github.com/huggingface/transformers"
                f" to request the support for training support for {model.hf_quantizer.quantization_config.quant_method}"
            )

        self.is_fsdp_xla_enabled = args.fsdp_config["xla"]
        if len(args.fsdp) > 0:
            if self.is_deepspeed_enabled:
                raise ValueError(
                    "Using --fsdp xxx together with --deepspeed is not possible, deactivate one of those flags."
                )
            if not args.fsdp_config["xla"] and args.parallel_mode != ParallelMode.DISTRIBUTED:
                raise ValueError("Using fsdp only works in distributed training.")

        # one place to sort out whether to place the model on device or not
        # postpone switching model to cuda when:
        # 1. MP - since we are trying to fit a much bigger than 1 gpu model
        # 2. fp16-enabled DeepSpeed loads the model in half the size and it doesn't need .to() anyway,
        #    and we only use deepspeed for training at the moment
        # 3. full bf16 or fp16 eval - since the model needs to be cast to the right dtype first
        # 4. FSDP - same as MP
        self.place_model_on_device = args.place_model_on_device
        if (
            self.is_model_parallel
            or self.is_deepspeed_enabled
            or ((args.fp16_full_eval or args.bf16_full_eval) and not args.do_train)
            or self.is_fsdp_xla_enabled
            or self.is_fsdp_enabled
        ):
            self.place_model_on_device = False

        default_collator = (
            DataCollatorWithPadding(processing_class)
            if processing_class is not None
            and isinstance(processing_class, (PreTrainedTokenizerBase, SequenceFeatureExtractor))
            else default_data_collator
        )
        self.data_collator = data_collator if data_collator is not None else default_collator
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.processing_class = processing_class

        # Bnb Quantized models doesn't support `.to` operation.
        if (
            self.place_model_on_device
            and not getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES
        ):
            self._move_model_to_device(model, args.device)

        # Force n_gpu to 1 to avoid DataParallel as MP will manage the GPUs
        if self.is_model_parallel:
            self.args._n_gpu = 1

        # later use `self.model is self.model_wrapped` to check if it's wrapped or not
        self.model_wrapped = model
        self.model = model

        # Just in case the model was wrapped outside of the `Trainer`
        unwrapped_model = self.accelerator.unwrap_model(model)
        model_forward = (
            unwrapped_model.forward
            if not _is_peft_model(unwrapped_model)
            else unwrapped_model.get_base_model().forward
        )
        forward_params = inspect.signature(model_forward).parameters

        # Check if the model has explicit setup for loss kwargs,
        # if not, check if `**kwargs` are in model.forward
        if hasattr(model, "accepts_loss_kwargs"):
            self.model_accepts_loss_kwargs = model.accepts_loss_kwargs
        else:
            self.model_accepts_loss_kwargs = any(
                k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values()
            )

        self.neftune_noise_alpha = args.neftune_noise_alpha

        self.compute_metrics = compute_metrics
        self.preprocess_logits_for_metrics = preprocess_logits_for_metrics
        self.optimizer, self.lr_scheduler = optimizers
        self.optimizer_cls_and_kwargs = optimizer_cls_and_kwargs
        if self.optimizer_cls_and_kwargs is not None and self.optimizer is not None:
            raise RuntimeError("Passing both `optimizers` and `optimizer_cls_and_kwargs` arguments is incompatible.")
        if model_init is not None and (self.optimizer is not None or self.lr_scheduler is not None):
            raise RuntimeError(
                "Passing a `model_init` is incompatible with providing the `optimizers` argument. "
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )
        if is_torch_xla_available() and self.optimizer is not None:
            for param in self.model.parameters():
                model_device = param.device
                break
            for param_group in self.optimizer.param_groups:
                if len(param_group["params"]) > 0:
                    optimizer_device = param_group["params"][0].device
                    break
            if model_device != optimizer_device:
                raise ValueError(
                    "The model and the optimizer parameters are not on the same device, which probably means you"
                    " created an optimizer around your model **before** putting on the device and passing it to the"
                    " `Trainer`. Make sure the lines `import torch_xla.core.xla_model as xm` and"
                    " `model.to(xm.xla_device())` is performed before the optimizer creation in your script."
                )
        if (self.is_fsdp_xla_enabled or self.is_fsdp_enabled) and (
            self.optimizer is not None or self.lr_scheduler is not None
        ):
            raise RuntimeError(
                "Passing `optimizers` is not allowed if PyTorch FSDP is enabled. "
                "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method."
            )
        default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to)
        callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks
        self.callback_handler = CallbackHandler(
            callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler
        )
        self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK)

        # Will be set to True by `self._setup_loggers()` on first call to `self.log()`.
        self._loggers_initialized = False

        # Create distant repo and output directory if needed
        self.hub_model_id = None
        if self.args.push_to_hub:
            self.init_hf_repo()
        if self.args.should_save:
            os.makedirs(self.args.output_dir, exist_ok=True)

        if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)):
            raise ValueError("The `data_collator` should be a simple callable (function, class with `__call__`).")

        if args.max_steps > 0 and args.num_train_epochs > 0:
            logger.info("max_steps is given, it will override any value given in num_train_epochs")

        if train_dataset is not None and not has_length(train_dataset) and args.max_steps <= 0:
            raise ValueError(
                "The train_dataset does not implement __len__, max_steps has to be specified. "
                "The number of steps needs to be known in advance for the learning rate scheduler."
            )

        if (
            train_dataset is not None
            and isinstance(train_dataset, torch.utils.data.IterableDataset)
            and args.group_by_length
        ):
            raise ValueError("the `--group_by_length` option is only available for `Dataset`, not `IterableDataset")

        self._signature_columns = None

        # Mixed precision setup
        self.use_apex = False
        self.use_cpu_amp = False

        # Mixed precision setup for SageMaker Model Parallel
        if is_sagemaker_mp_enabled():
            # BF16 + model parallelism in SageMaker: currently not supported, raise an error
            if args.bf16:
                raise ValueError("SageMaker Model Parallelism does not support BF16 yet. Please use FP16 instead ")

            if IS_SAGEMAKER_MP_POST_1_10:
                # When there's mismatch between SMP config and trainer argument, use SMP config as truth
                if args.fp16 != smp.state.cfg.fp16:
                    logger.warning(
                        f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                        f"but FP16 provided in trainer argument is {args.fp16}, "
                        f"setting to {smp.state.cfg.fp16}"
                    )
                    args.fp16 = smp.state.cfg.fp16
            else:
                # smp < 1.10 does not support fp16 in trainer.
                if hasattr(smp.state.cfg, "fp16"):
                    logger.warning(
                        f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, "
                        "but SageMaker Model Parallelism < 1.10 does not support FP16 in trainer."
                    )
        if (args.fp16 or args.bf16) and args.half_precision_backend == "auto":
            if args.device == torch.device("cpu"):
                if args.fp16:
                    if not is_torch_greater_or_equal_than_2_3:
                        raise ValueError("Tried to use `fp16` but it is not supported on cpu")
                else:
                    args.half_precision_backend = "cpu_amp"
            logger.info(f"Using {args.half_precision_backend} half precision backend")

        if (args.fp16 or args.bf16) and not (self.is_deepspeed_enabled or is_sagemaker_mp_enabled()):
            # deepspeed and SageMaker Model Parallel manage their own half precision
            if args.half_precision_backend == "cpu_amp":
                self.use_cpu_amp = True
                self.amp_dtype = torch.bfloat16
            elif args.half_precision_backend == "apex":
                if not is_apex_available():
                    raise ImportError(
                        "Using FP16 with APEX but APEX is not installed, please refer to"
                        " https://www.github.com/nvidia/apex."
                    )
                self.use_apex = True

        # Label smoothing
        if self.args.label_smoothing_factor != 0:
            self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor)
        else:
            self.label_smoother = None

        self.control = TrainerControl()

        self.state = TrainerState(
            is_local_process_zero=self.is_local_process_zero(),
            is_world_process_zero=self.is_world_process_zero(),
            stateful_callbacks=[
                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
            ],
        )
        # Internal variable to count flos in each process, will be accumulated in `self.state.total_flos` then
        # returned to 0 every time flos need to be logged
        self.current_flos = 0
        self.hp_search_backend = None
        default_label_names = find_labels(self.model.__class__)
        self.label_names = default_label_names if self.args.label_names is None else self.args.label_names
        self.can_return_loss = can_return_loss(self.model.__class__)
        self.control = self.callback_handler.on_init_end(self.args, self.state, self.control)

        # Internal variables to help with automatic batch size reduction
        self._train_batch_size = args.train_batch_size
        self._created_lr_scheduler = False

        # very last
        self._memory_tracker.stop_and_update_metrics()

        # torch.compile
        if args.torch_compile and not is_torch_compile_available():
            raise RuntimeError("Using torch.compile requires PyTorch 2.0 or higher.")

        self.is_fsdp_xla_v2_enabled = args.fsdp_config.get("xla_fsdp_v2", False)
        if self.is_fsdp_xla_v2_enabled:
            if not IS_XLA_FSDPV2_POST_2_2:
                raise ValueError("FSDPv2 requires `torch_xla` 2.2 or higher.")
            # Prepare the SPMD mesh that is going to be used by the data loader and the FSDPv2 wrapper.
            # Tensor axis is just a placeholder where it will not be used in FSDPv2.
            num_devices = xr.global_runtime_device_count()
            xs.set_global_mesh(xs.Mesh(np.array(range(num_devices)), (num_devices, 1), axis_names=("fsdp", "tensor")))
        self.is_fsdp_xla_v1_enabled = self.is_fsdp_xla_enabled and not self.is_fsdp_xla_v2_enabled

    @property
    def tokenizer(self) -> Optional[PreTrainedTokenizerBase]:
        logger.warning("Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.")
        return self.processing_class

    @tokenizer.setter
    def tokenizer(self, processing_class) -> None:
        logger.warning(
            "Trainer.tokenizer is now deprecated. You should use `Trainer.processing_class = processing_class` instead."
        )
        self.processing_class = processing_class

    def _activate_neftune(self, model):
        r"""
        Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper:
        https://arxiv.org/abs/2310.05914
        """
        unwrapped_model = self.accelerator.unwrap_model(model)

        if _is_peft_model(unwrapped_model):
            embeddings = unwrapped_model.base_model.model.get_input_embeddings()
        else:
            embeddings = unwrapped_model.get_input_embeddings()

        del unwrapped_model

        embeddings.neftune_noise_alpha = self.neftune_noise_alpha
        hook_handle = embeddings.register_forward_hook(neftune_post_forward_hook)
        self.neftune_hook_handle = hook_handle
        return model

    def _deactivate_neftune(self, model):
        """
        Deactivates the neftune method. Make sure to call `_activate_neftune` first.
        """
        if not hasattr(self, "neftune_hook_handle"):
            raise ValueError("Neftune is not activated make sure to call `trainer._activate_neftune()` first")

        unwrapped_model = self.accelerator.unwrap_model(model)

        if _is_peft_model(unwrapped_model):
            embeddings = unwrapped_model.base_model.model.get_input_embeddings()
        else:
            embeddings = unwrapped_model.get_input_embeddings()

        self.neftune_hook_handle.remove()
        del embeddings.neftune_noise_alpha, unwrapped_model

    def add_callback(self, callback):
        """
        Add a callback to the current list of [`~transformers.TrainerCallback`].

        Args:
           callback (`type` or [`~transformers.TrainerCallback]`):
               A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
               first case, will instantiate a member of that class.
        """
        self.callback_handler.add_callback(callback)

    def pop_callback(self, callback):
        """
        Remove a callback from the current list of [`~transformers.TrainerCallback`] and returns it.

        If the callback is not found, returns `None` (and no error is raised).

        Args:
           callback (`type` or [`~transformers.TrainerCallback]`):
               A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
               first case, will pop the first member of that class found in the list of callbacks.

        Returns:
            [`~transformers.TrainerCallback`]: The callback removed, if found.
        """
        return self.callback_handler.pop_callback(callback)

    def remove_callback(self, callback):
        """
        Remove a callback from the current list of [`~transformers.TrainerCallback`].

        Args:
           callback (`type` or [`~transformers.TrainerCallback]`):
               A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the
               first case, will remove the first member of that class found in the list of callbacks.
        """
        self.callback_handler.remove_callback(callback)

    def _move_model_to_device(self, model, device):
        model = model.to(device)
        # Moving a model to an XLA device disconnects the tied weights, so we have to retie them.
        if self.args.parallel_mode == ParallelMode.TPU and hasattr(model, "tie_weights"):
            model.tie_weights()

    def _set_signature_columns_if_needed(self):
        if self._signature_columns is None:
            # Inspect model forward signature to keep only the arguments it accepts.
            model_to_inspect = self.model
            if _is_peft_model(self.model):
                if hasattr(self.model, "get_base_model"):
                    model_to_inspect = self.model.get_base_model()
                else:
                    # PeftMixedModel do not provide a `get_base_model` method
                    model_to_inspect = self.model.base_model.model
            signature = inspect.signature(model_to_inspect.forward)
            self._signature_columns = list(signature.parameters.keys())
            # Labels may be named label or label_ids, the default data collator handles that.
            self._signature_columns += list(set(["label", "label_ids"] + self.label_names))

    def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None):
        if not self.args.remove_unused_columns:
            return dataset
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        ignored_columns = list(set(dataset.column_names) - set(signature_columns))
        if len(ignored_columns) > 0:
            dset_description = "" if description is None else f"in the {description} set"
            logger.info(
                f"The following columns {dset_description} don't have a corresponding argument in "
                f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
                f" If {', '.join(ignored_columns)} are not expected by `{self.model.__class__.__name__}.forward`, "
                " you can safely ignore this message."
            )

        columns = [k for k in signature_columns if k in dataset.column_names]
        if len(columns) == 0:
            raise ValueError(
                "No columns in the dataset match the model's forward method signature. "
                f"The following columns have been ignored: [{', '.join(ignored_columns)}]. "
                "Please check the dataset and model. You may need to set `remove_unused_columns=False` in `TrainingArguments`."
            )

        if version.parse(datasets.__version__) < version.parse("1.4.0"):
            dataset.set_format(
                type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"]
            )
            return dataset
        else:
            return dataset.remove_columns(ignored_columns)

    def _get_collator_with_removed_columns(
        self, data_collator: Callable, description: Optional[str] = None
    ) -> Callable:
        """Wrap the data collator in a callable removing unused columns."""
        if not self.args.remove_unused_columns:
            return data_collator
        self._set_signature_columns_if_needed()
        signature_columns = self._signature_columns

        remove_columns_collator = RemoveColumnsCollator(
            data_collator=data_collator,
            signature_columns=signature_columns,
            logger=logger,
            description=description,
            model_name=self.model.__class__.__name__,
        )
        return remove_columns_collator

    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
        if self.train_dataset is None or not has_length(self.train_dataset):
            return None

        # Build the sampler.
        if self.args.group_by_length:
            if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
                lengths = (
                    self.train_dataset[self.args.length_column_name]
                    if self.args.length_column_name in self.train_dataset.column_names
                    else None
                )
            else:
                lengths = None
            model_input_name = (
                self.processing_class.model_input_names[0] if self.processing_class is not None else None
            )
            return LengthGroupedSampler(
                self.args.train_batch_size * self.args.gradient_accumulation_steps,
                dataset=self.train_dataset,
                lengths=lengths,
                model_input_name=model_input_name,
            )

        else:
            return RandomSampler(self.train_dataset)

    def get_train_dataloader(self) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        train_dataset = self.train_dataset
        data_collator = self.data_collator
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="training")

        dataloader_params = {
            "batch_size": self._train_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(train_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_train_sampler()
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor

        return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))

    def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.Sampler]:
        if eval_dataset is None or not has_length(eval_dataset):
            return None
        # Build the sampler.

        # Deprecated code
        if self.args.use_legacy_prediction_loop:
            if is_torch_xla_available():
                return SequentialDistributedSampler(
                    eval_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal()
                )
            elif is_sagemaker_mp_enabled():
                return SequentialDistributedSampler(
                    eval_dataset,
                    num_replicas=smp.dp_size(),
                    rank=smp.dp_rank(),
                    batch_size=self.args.per_device_eval_batch_size,
                )
            else:
                return SequentialSampler(eval_dataset)

        if self.args.group_by_length:
            if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
                lengths = (
                    eval_dataset[self.args.length_column_name]
                    if self.args.length_column_name in eval_dataset.column_names
                    else None
                )
            else:
                lengths = None
            model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
            return LengthGroupedSampler(
                self.args.eval_batch_size,
                dataset=eval_dataset,
                lengths=lengths,
                model_input_name=model_input_name,
            )

        if self.args.world_size <= 1:
            return SequentialSampler(eval_dataset)
        else:
            return None

    def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None) -> DataLoader:
        """
        Returns the evaluation [`~torch.utils.data.DataLoader`].

        Subclass and override this method if you want to inject some custom behavior.

        Args:
            eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
                If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")

        # If we have persistent workers, don't do a fork bomb especially as eval datasets
        # don't change during training
        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
        if (
            hasattr(self, "_eval_dataloaders")
            and dataloader_key in self._eval_dataloaders
            and self.args.dataloader_persistent_workers
        ):
            return self.accelerator.prepare(self._eval_dataloaders[dataloader_key])

        eval_dataset = (
            self.eval_dataset[eval_dataset]
            if isinstance(eval_dataset, str)
            else eval_dataset
            if eval_dataset is not None
            else self.eval_dataset
        )
        data_collator = self.data_collator

        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="evaluation")

        dataloader_params = {
            "batch_size": self.args.eval_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(eval_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor

        # accelerator.free_memory() will destroy the references, so
        # we need to store the non-prepared version
        eval_dataloader = DataLoader(eval_dataset, **dataloader_params)
        if self.args.dataloader_persistent_workers:
            if hasattr(self, "_eval_dataloaders"):
                self._eval_dataloaders[dataloader_key] = eval_dataloader
            else:
                self._eval_dataloaders = {dataloader_key: eval_dataloader}

        return self.accelerator.prepare(eval_dataloader)

    def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
        """
        Returns the test [`~torch.utils.data.DataLoader`].

        Subclass and override this method if you want to inject some custom behavior.

        Args:
            test_dataset (`torch.utils.data.Dataset`, *optional*):
                The test dataset to use. If it is a [`~datasets.Dataset`], columns not accepted by the
                `model.forward()` method are automatically removed. It must implement `__len__`.
        """
        data_collator = self.data_collator

        if is_datasets_available() and isinstance(test_dataset, datasets.Dataset):
            test_dataset = self._remove_unused_columns(test_dataset, description="test")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="test")

        dataloader_params = {
            "batch_size": self.args.eval_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
            "persistent_workers": self.args.dataloader_persistent_workers,
        }

        if not isinstance(test_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_eval_sampler(test_dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor

        # We use the same batch_size as for eval.
        return self.accelerator.prepare(DataLoader(test_dataset, **dataloader_params))

    def create_optimizer_and_scheduler(self, num_training_steps: int):
        """
        Setup the optimizer and the learning rate scheduler.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method (or `create_optimizer` and/or
        `create_scheduler`) in a subclass.
        """
        self.create_optimizer()
        if IS_SAGEMAKER_MP_POST_1_10 and smp.state.cfg.fp16:
            # If smp >= 1.10 and fp16 is enabled, we unwrap the optimizer
            optimizer = self.optimizer.optimizer
        else:
            optimizer = self.optimizer
        self.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer)

    def get_decay_parameter_names(self, model) -> List[str]:
        """
        Get all parameter names that weight decay will be applied to

        Note that some models implement their own layernorm instead of calling nn.LayerNorm, weight decay could still
        apply to those modules since this function only filter out instance of nn.LayerNorm
        """
        decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS)
        decay_parameters = [name for name in decay_parameters if "bias" not in name]
        return decay_parameters

    def create_optimizer(self):
        """
        Setup the optimizer.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
        """
        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model

        if self.optimizer is None:
            decay_parameters = self.get_decay_parameter_names(opt_model)
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": self.args.weight_decay,
                },
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": 0.0,
                },
            ]

            if self.optimizer_cls_and_kwargs is not None:
                optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs
            else:
                optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model)

            # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs`
            # e.g. for GaLore optimizer.
            if "params" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("params")

            # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs`
            # e.g. for LOMO optimizer.
            if "model" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("model")

            # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict`
            # to avoid arguments conflicts.
            if "optimizer_dict" in optimizer_kwargs:
                optimizer_grouped_parameters = optimizer_kwargs.pop("optimizer_dict")

            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

            if optimizer_cls.__name__ == "Adam8bit":
                import bitsandbytes

                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()

                skipped = 0
                for module in opt_model.modules():
                    if isinstance(module, nn.Embedding):
                        skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values())
                        logger.info(f"skipped {module}: {skipped/2**20}M params")
                        manager.register_module_override(module, "weight", {"optim_bits": 32})
                        logger.debug(f"bitsandbytes: will optimize {module} in fp32")
                logger.info(f"skipped: {skipped/2**20}M params")

        if is_sagemaker_mp_enabled():
            self.optimizer = smp.DistributedOptimizer(self.optimizer)

        return self.optimizer

    def get_num_trainable_parameters(self):
        """
        Get the number of trainable parameters.
        """
        return sum(p.numel() for p in self.model.parameters() if p.requires_grad)

    def get_learning_rates(self):
        """
        Returns the learning rate of each parameter from self.optimizer.
        """
        if self.optimizer is None:
            raise ValueError("Trainer optimizer is None, please make sure you have setup the optimizer before.")
        return [group["lr"] for group in self.optimizer.param_groups]

    def get_optimizer_group(self, param: Optional[Union[str, torch.nn.parameter.Parameter]] = None):
        """
        Returns optimizer group for a parameter if given, else returns all optimizer groups for params.

        Args:
            param (`str` or `torch.nn.parameter.Parameter`, *optional*):
                The parameter for which optimizer group needs to be returned.
        """
        if self.optimizer is None:
            raise ValueError("Trainer optimizer is None, please make sure you have setup the optimizer before.")
        if param is not None:
            for group in self.optimizer.param_groups:
                if param in group["params"]:
                    return group
        return [group["params"] for group in self.optimizer.param_groups]

    @staticmethod
    def get_optimizer_cls_and_kwargs(
        args: TrainingArguments, model: Optional[PreTrainedModel] = None
    ) -> Tuple[Any, Any]:
        """
        Returns the optimizer class and optimizer parameters based on the training arguments.

        Args:
            args (`transformers.training_args.TrainingArguments`):
                The training arguments for the training session.

        """

        # parse args.optim_args
        optim_args = {}
        if args.optim_args:
            for mapping in args.optim_args.replace(" ", "").split(","):
                key, value = mapping.split("=")
                optim_args[key] = value

        optimizer_kwargs = {"lr": args.learning_rate}

        adam_kwargs = {
            "betas": (args.adam_beta1, args.adam_beta2),
            "eps": args.adam_epsilon,
        }
        if args.optim == OptimizerNames.ADAFACTOR:
            optimizer_cls = Adafactor
            optimizer_kwargs.update({"scale_parameter": False, "relative_step": False})
        elif args.optim == OptimizerNames.ADAMW_HF:
            from .optimization import AdamW

            optimizer_cls = AdamW
            optimizer_kwargs.update(adam_kwargs)
        elif args.optim in [OptimizerNames.ADAMW_TORCH, OptimizerNames.ADAMW_TORCH_FUSED]:
            from torch.optim import AdamW

            optimizer_cls = AdamW
            optimizer_kwargs.update(adam_kwargs)
            if args.optim == OptimizerNames.ADAMW_TORCH_FUSED:
                optimizer_kwargs.update({"fused": True})
        elif args.optim == OptimizerNames.ADAMW_TORCH_XLA:
            try:
                from torch_xla.amp.syncfree import AdamW

                optimizer_cls = AdamW
                optimizer_kwargs.update(adam_kwargs)
            except ImportError:
                raise ValueError("Trainer failed to import syncfree AdamW from torch_xla.")
        elif args.optim == OptimizerNames.ADAMW_TORCH_NPU_FUSED:
            try:
                from torch_npu.optim import NpuFusedAdamW

                optimizer_cls = NpuFusedAdamW
                optimizer_kwargs.update(adam_kwargs)
            except ImportError:
                raise ValueError("Trainer failed to import FusedAdamW from torch_npu.")
        elif args.optim == OptimizerNames.ADAMW_APEX_FUSED:
            try:
                from apex.optimizers import FusedAdam

                optimizer_cls = FusedAdam
                optimizer_kwargs.update(adam_kwargs)
            except ImportError:
                raise ValueError("Trainer tried to instantiate apex FusedAdam but apex is not installed!")
        elif args.optim in [
            OptimizerNames.ADAMW_BNB,
            OptimizerNames.ADAMW_8BIT,
            OptimizerNames.PAGED_ADAMW,
            OptimizerNames.PAGED_ADAMW_8BIT,
            OptimizerNames.ADEMAMIX,
            OptimizerNames.ADEMAMIX_8BIT,
            OptimizerNames.PAGED_ADEMAMIX,
            OptimizerNames.PAGED_ADEMAMIX_8BIT,
            OptimizerNames.LION,
            OptimizerNames.LION_8BIT,
            OptimizerNames.PAGED_LION,
            OptimizerNames.PAGED_LION_8BIT,
            OptimizerNames.RMSPROP_BNB,
            OptimizerNames.RMSPROP_8BIT,
            OptimizerNames.RMSPROP_32BIT,
        ]:
            try:
                from bitsandbytes.optim import AdamW, Lion, RMSprop

                is_paged = False
                optim_bits = 32
                optimizer_cls = None
                additional_optim_kwargs = adam_kwargs
                if "paged" in args.optim:
                    is_paged = True
                if "8bit" in args.optim:
                    optim_bits = 8
                if "adam" in args.optim:
                    optimizer_cls = AdamW
                elif "lion" in args.optim:
                    optimizer_cls = Lion
                    additional_optim_kwargs = {"betas": (args.adam_beta1, args.adam_beta2)}
                elif "rmsprop" in args.optim:
                    optimizer_cls = RMSprop
                    # Above we pass all `adam_kwargs` to the optimizer, here
                    # we only pass `optim_args` which can be passed by the user.
                    additional_optim_kwargs = optim_args
                elif "ademamix" in args.optim:
                    if is_bitsandbytes_available() and version.parse(
                        importlib.metadata.version("bitsandbytes")
                    ) < version.parse("0.44.0"):
                        raise ValueError(
                            "The AdEMAMix optimizer is not supported by your current version of `bitsandbytes`. "
                            "Please install `bitsandbytes` >= 0.44.0."
                        )

                    from bitsandbytes.optim import AdEMAMix

                    optimizer_cls = AdEMAMix
                    additional_optim_kwargs = {
                        "betas": (
                            float(optim_args.get("beta1", args.adam_beta1)),
                            float(optim_args.get("beta2", args.adam_beta2)),
                            float(optim_args.get("beta3", 0.9999)),
                        ),
                        "alpha": float(optim_args.get("alpha", 5.0)),
                        "eps": float(optim_args.get("eps", args.adam_epsilon)),
                    }

                    if "t_alpha" in optim_args:
                        additional_optim_kwargs["t_alpha"] = int(optim_args["t_alpha"])

                    if "t_beta3" in optim_args:
                        additional_optim_kwargs["t_beta3"] = int(optim_args["t_beta3"])

                bnb_kwargs = {"optim_bits": optim_bits}
                if "rmsprop" not in args.optim:
                    bnb_kwargs["is_paged"] = is_paged

                optimizer_kwargs.update(additional_optim_kwargs)
                optimizer_kwargs.update(bnb_kwargs)
            except ImportError:
                raise ValueError("Trainer tried to instantiate bnb optimizer but `bitsandbytes` is not installed!")
            if is_bitsandbytes_available() and version.parse(
                importlib.metadata.version("bitsandbytes")
            ) < version.parse("0.41.1"):
                logger.warning(
                    "You are using 8-bit optimizers with a version of `bitsandbytes` < 0.41.1. "
                    "It is recommended to update your version as a major bug has been fixed in 8-bit optimizers."
                )
        elif args.optim == OptimizerNames.ADAMW_ANYPRECISION:
            try:
                from torchdistx.optimizers import AnyPrecisionAdamW

                optimizer_cls = AnyPrecisionAdamW
                optimizer_kwargs.update(adam_kwargs)

                # TODO Change dtypes back to M=FP32, Var = BF16, Kahan = False once they can be cast together in torchdistx.
                optimizer_kwargs.update(
                    {
                        "use_kahan_summation": strtobool(optim_args.get("use_kahan_summation", "False")),
                        "momentum_dtype": getattr(torch, optim_args.get("momentum_dtype", "float32")),
                        "variance_dtype": getattr(torch, optim_args.get("variance_dtype", "float32")),
                        "compensation_buffer_dtype": getattr(
                            torch, optim_args.get("compensation_buffer_dtype", "bfloat16")
                        ),
                    }
                )
            except ImportError:
                raise ValueError("Please install https://github.com/pytorch/torchdistx")
        elif args.optim == OptimizerNames.SGD:
            optimizer_cls = torch.optim.SGD
        elif args.optim == OptimizerNames.ADAGRAD:
            optimizer_cls = torch.optim.Adagrad
        elif args.optim == OptimizerNames.RMSPROP:
            optimizer_cls = torch.optim.RMSprop
        elif args.optim in [
            OptimizerNames.GALORE_ADAMW,
            OptimizerNames.GALORE_ADAMW_8BIT,
            OptimizerNames.GALORE_ADAFACTOR,
            OptimizerNames.GALORE_ADAMW_LAYERWISE,
            OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE,
            OptimizerNames.GALORE_ADAFACTOR_LAYERWISE,
        ]:
            if not is_galore_torch_available():
                raise ImportError(
                    "You need to install `galore_torch` in order to use GaLore optimizers"
                    " install it with `pip install git+https://github.com/jiaweizzhao/GaLore`"
                )
            from galore_torch import GaLoreAdafactor, GaLoreAdamW, GaLoreAdamW8bit

            is_layerwise = args.optim.lower().endswith("layerwise")
            if is_layerwise and args.parallel_mode == ParallelMode.DISTRIBUTED:
                raise NotImplementedError("Layer-wise GaLore does not support DDP at this time")

            optimizer_mapping = {
                OptimizerNames.GALORE_ADAMW: GaLoreAdamW,
                OptimizerNames.GALORE_ADAMW_8BIT: GaLoreAdamW8bit,
                OptimizerNames.GALORE_ADAFACTOR: GaLoreAdafactor,
                OptimizerNames.GALORE_ADAMW_LAYERWISE: GaLoreAdamW,
                OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE: GaLoreAdamW8bit,
                OptimizerNames.GALORE_ADAFACTOR_LAYERWISE: GaLoreAdafactor,
            }

            optimizer_cls = optimizer_mapping[args.optim]

            if args.optim_target_modules is None:
                raise ValueError(
                    "You need to define a `optim_target_modules` in order to properly use GaLore optimizers"
                )

            if not isinstance(args.optim_target_modules, (list, str)):
                raise ValueError(
                    f"`optim_target_modules` has to be a list of strings, a string corresponding to a regex, or a specific module or 'all-linear', you passed {args.optim_target_modules}"
                )

            if model is None:
                raise ValueError("You need to pass a model in order to correctly initialize a GaLore optimizer.")

            logger.warning(
                "Activated GaLoRE fine-tuning, depending on your model size and hardware, the training might take a while before starting. Please be patient !"
            )

            all_linear = (
                isinstance(args.optim_target_modules, str)
                and args.optim_target_modules.replace("_", "-") == "all-linear"
            )

            galore_params = []
            galore_params_names = []
            for module_name, module in model.named_modules():
                target_module_exists, is_regex = check_target_module_exists(
                    args.optim_target_modules, module_name, return_is_regex=True
                )

                if not isinstance(module, nn.Linear):
                    # Warn in case we match but it's not a linear layer
                    if target_module_exists and not is_regex:
                        logger.warning(
                            f"{module_name} has been matched but ignored as GaLore only supports linear layers. Please double check your `optim_target_modules`!"
                        )

                    continue

                if not target_module_exists and not all_linear:
                    continue

                galore_params.append(module.weight)
                galore_params_names.append(module_name + ".weight")

            if len(galore_params) == 0:
                raise ValueError(
                    f"None of the target modules were found! ({args.optim_target_modules}). Please make sure to pass a valid `target_modules`."
                )

            non_galore_params = [p for n, p in model.named_parameters() if n not in galore_params_names]

            galore_optim_kwargs = {
                "rank": int(optim_args.pop("rank", 128)),
                "update_proj_gap": int(optim_args.pop("update_proj_gap", 200)),
                "scale": float(optim_args.pop("scale", 0.25)),
                "proj_type": optim_args.pop("proj_type", "std"),
            }

            # The default args are from the official repository: https://github.com/jiaweizzhao/GaLore
            param_groups = [
                {"params": non_galore_params},
                {"params": galore_params, **galore_optim_kwargs},
            ]

            if is_layerwise:
                # For layer-wise optimizers, the optimization step is done through post accumulation
                # gradient hooks. The trick is to first attach these hooks to the model parameters then
                # create a dummy optimizer that will perform no-ops in the Trainer.
                # See the original implementation or the nice implementation from @hiyouga
                # here: https://github.com/hiyouga/LLaMA-Factory/commit/8664262cde3919e10eaecbd66e8c5d356856362e#diff-ebe08ab14496dfb9e06075f0fdd36799ef6d1535cc4dd4715b74c4e3e06fe3ba
                if args.gradient_accumulation_steps != 1:
                    raise ValueError("Layerwise GaLoRE optimizer do not support gradient accumulation !")

                optimizer_dict = {}
                for param in non_galore_params:
                    param_groups = [{"params": [param]}]
                    optimizer_dict[param] = optimizer_cls(param_groups, **optimizer_kwargs)
                for param in galore_params:
                    param_groups = [{"params": [param], **galore_optim_kwargs}]
                    optimizer_dict[param] = optimizer_cls(param_groups, **optimizer_kwargs)

                def optimizer_hook(param):
                    if param.grad is not None:
                        optimizer_dict[param].step()
                        optimizer_dict[param].zero_grad()

                for param in model.parameters():
                    if param.requires_grad:
                        param.register_post_accumulate_grad_hook(optimizer_hook)

                optimizer_cls = LayerWiseDummyOptimizer
                optimizer_kwargs.update({"optimizer_dict": optimizer_dict})

            optimizer_kwargs.update({"params": param_groups})

            if args.optim == OptimizerNames.GALORE_ADAFACTOR:
                optimizer_kwargs.update({"scale_parameter": False, "relative_step": False})
        elif args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
            if not is_lomo_available():
                raise ImportError(
                    "You need to install `lomo_optim` in order to use LOMO optimizers"
                    " install it with `pip install lomo-optim`"
                )
            if not is_accelerate_available("0.30.0"):
                raise ImportError("You need to have `accelerate>=0.30.0` to be able to use LOMO optimizers")

            if model is None:
                raise ValueError("You need to pass a `model` in order to correctly initialize a LOMO optimizer.")

            from lomo_optim import AdaLomo, Lomo

            if "ada" in args.optim:
                optimizer_cls = AdaLomo
            else:
                optimizer_cls = Lomo

            optimizer_kwargs.update({"model": model})
        elif args.optim == OptimizerNames.GROKADAMW:
            if not is_grokadamw_available():
                raise ValueError("Please install grokadamw with `pip install grokadamw`")

            from grokadamw import GrokAdamW

            optimizer_cls = GrokAdamW
            optimizer_kwargs.update(
                {
                    "alpha_init": float(optim_args.get("alpha_init", 0.98)),
                    "lamb": float(optim_args.get("lamb", 2.0)),
                    "gamma": float(optim_args.get("gamma", 0.1)),
                    "grokking_signal_decay_rate": float(optim_args.get("grokking_signal_decay_rate", 0.1)),
                    "gradient_clipping": float(optim_args.get("gradient_clipping", 1.0)),
                }
            )
        elif args.optim == OptimizerNames.ADAMW_TORCH_4BIT:
            if not is_torchao_available() or version.parse(importlib.metadata.version("torchao")) < version.parse(
                "0.4.0"
            ):
                raise ImportError(
                    "You need to have `torchao>=0.4.0` in order to use torch 4-bit optimizers."
                    "Install it with `pip install torchao` or follow the instructions here: https://github.com/pytorch/ao"
                )
            if version.parse(importlib.metadata.version("torch")) <= version.parse("2.4"):
                raise ImportError(
                    "You need to have `torch>2.4` in order to use torch 4-bit optimizers. "
                    "Install it with `pip install --upgrade torch` it is available on pipy. Otherwise, you need to install torch nightly."
                )
            from torchao.prototype.low_bit_optim import AdamW4bit

            optimizer_cls = AdamW4bit
            optimizer_kwargs.update(adam_kwargs)
        elif args.optim in [
            OptimizerNames.SCHEDULE_FREE_ADAMW,
            OptimizerNames.SCHEDULE_FREE_SGD,
        ]:
            if not is_schedulefree_available():
                raise ImportError(
                    "You need to install `schedulefree` in order to use schedulefree optimizers"
                    " install it with `pip install schedulefree`"
                )
            if not is_accelerate_available("0.30.0"):
                raise ImportError("You need to have `accelerate>=0.30.0` to be able to use schedulefree optimizers")
            from schedulefree import AdamWScheduleFree, SGDScheduleFree

            additional_optim_kwargs = {}
            if args.optim == OptimizerNames.SCHEDULE_FREE_ADAMW:
                optimizer_cls = AdamWScheduleFree
                additional_optim_kwargs = adam_kwargs
            elif args.optim == OptimizerNames.SCHEDULE_FREE_SGD:
                optimizer_cls = SGDScheduleFree
            else:
                raise ValueError("Invalid schedulefree optimizer")
            additional_optim_kwargs["weight_decay"] = args.weight_decay
            additional_optim_kwargs["warmup_steps"] = args.warmup_steps
            additional_optim_kwargs.update(
                {
                    "weight_lr_power": float(optim_args.get("weight_lr_power", 2.0)),
                    "r": float(optim_args.get("r", 0.0)),
                }
            )
            optimizer_kwargs.update(additional_optim_kwargs)
        else:
            raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}")
        return optimizer_cls, optimizer_kwargs

    def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None):
        """
        Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or
        passed as an argument.

        Args:
            num_training_steps (int): The number of training steps to do.
        """
        if self.lr_scheduler is None:
            self.lr_scheduler = get_scheduler(
                self.args.lr_scheduler_type,
                optimizer=self.optimizer if optimizer is None else optimizer,
                num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
                num_training_steps=num_training_steps,
                scheduler_specific_kwargs=self.args.lr_scheduler_kwargs,
            )
            self._created_lr_scheduler = True
        return self.lr_scheduler

    def num_examples(self, dataloader: DataLoader) -> int:
        """
        Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset. When
        dataloader.dataset does not exist or has no length, estimates as best it can
        """
        try:
            dataset = dataloader.dataset
            # Special case for IterableDatasetShard, we need to dig deeper
            if isinstance(dataset, IterableDatasetShard):
                return len(dataloader.dataset.dataset)
            return len(dataloader.dataset)
        except (NameError, AttributeError, TypeError):  # no dataset or length, estimate by length of dataloader
            return len(dataloader) * self.args.per_device_train_batch_size

    @staticmethod
    def num_tokens(train_dl: DataLoader, max_steps: Optional[int] = None) -> int:
        """
        Helper to get number of tokens in a [`~torch.utils.data.DataLoader`] by enumerating dataloader.
        """
        train_tokens = 0
        try:
            for batch in train_dl:
                tokens = batch["input_ids"].numel()
                if max_steps is not None:
                    return tokens * max_steps
                train_tokens += tokens
        except KeyError:
            logger.warning("Cannot get num_tokens from dataloader")
        return train_tokens

    def _hp_search_setup(self, trial: Union["optuna.Trial", Dict[str, Any]]):
        """HP search setup code"""
        self._trial = trial

        if self.hp_search_backend is None or trial is None:
            return
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            params = self.hp_space(trial)
        elif self.hp_search_backend == HPSearchBackend.RAY:
            params = trial
            params.pop("wandb", None)
        elif self.hp_search_backend == HPSearchBackend.SIGOPT:
            params = {k: int(v) if isinstance(v, str) else v for k, v in trial.assignments.items()}
        elif self.hp_search_backend == HPSearchBackend.WANDB:
            params = trial

        for key, value in params.items():
            if not hasattr(self.args, key):
                logger.warning(
                    f"Trying to set {key} in the hyperparameter search but there is no corresponding field in"
                    " `TrainingArguments`."
                )
                continue
            old_attr = getattr(self.args, key, None)
            # Casting value to the proper type
            if old_attr is not None:
                value = type(old_attr)(value)

            setattr(self.args, key, value)
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            logger.info(f"Trial: {trial.params}")
        if self.hp_search_backend == HPSearchBackend.SIGOPT:
            logger.info(f"SigOpt Assignments: {trial.assignments}")
        if self.hp_search_backend == HPSearchBackend.WANDB:
            logger.info(f"W&B Sweep parameters: {trial}")
        if self.is_deepspeed_enabled:
            if self.args.deepspeed is None:
                raise ValueError("For sweeps with deepspeed, `args.deepspeed` must be set")

            self.accelerator.free_memory()

            # Rebuild the deepspeed config to reflect the updated training parameters
            from accelerate.utils import DeepSpeedPlugin

            from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig

            self.args.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.args.deepspeed)
            self.args.hf_deepspeed_config.trainer_config_process(self.args)
            self.args.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.args.hf_deepspeed_config)

            # From 1.0 on, we need to fully wipe the DS plugin when doing sweeps.
            # Simply calling `_reset_state` is enough and doesn't need a version pin.
            AcceleratorState()._reset_state()

        self.create_accelerator_and_postprocess()

    def _report_to_hp_search(self, trial: Union["optuna.Trial", Dict[str, Any]], step: int, metrics: Dict[str, float]):
        if self.hp_search_backend is None or trial is None:
            return
        metrics = metrics.copy()
        self.objective = self.compute_objective(metrics)
        if self.hp_search_backend == HPSearchBackend.OPTUNA:
            import optuna

            if hasattr(trial, "study") and not trial.study._is_multi_objective():
                trial.report(self.objective, step)
                if trial.should_prune():
                    self.callback_handler.on_train_end(self.args, self.state, self.control)
                    raise optuna.TrialPruned()
        elif self.hp_search_backend == HPSearchBackend.RAY:
            import ray.train

            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
                checkpoint = None
                if self.control.should_save:
                    self._tune_save_checkpoint(checkpoint_dir=temp_checkpoint_dir)
                    checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
                metrics["objective"] = self.objective
                ray.train.report(metrics, checkpoint=checkpoint)

    def _tune_save_checkpoint(self, checkpoint_dir: str):
        output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}")
        self.save_model(output_dir, _internal_call=True)
        if self.args.should_save:
            # Update the `TrainerControl` state to where we are currently
            self.state.stateful_callbacks["TrainerControl"] = self.control.state()
            self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
            torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))

    def call_model_init(self, trial=None):
        model_init_argcount = number_of_arguments(self.model_init)
        if model_init_argcount == 0:
            model = self.model_init()
        elif model_init_argcount == 1:
            model = self.model_init(trial)
        else:
            raise RuntimeError("model_init should have 0 or 1 argument.")

        if model is None:
            raise RuntimeError("model_init should not return None.")

        return model

    def torch_jit_model_eval(self, model, dataloader, training=False):
        if not training:
            if dataloader is None:
                logger.warning("failed to use PyTorch jit mode due to current dataloader is none.")
                return model
            example_batch = next(iter(dataloader))
            example_batch = self._prepare_inputs(example_batch)
            try:
                jit_model = copy.copy(model)
                jit_model.eval()
                original_forward = jit_model.__dict__.pop("_original_forward", None)
                # remove mixed precision hooks from the model
                if original_forward:
                    jit_model.forward = original_forward
                with self.accelerator.autocast(cache_enabled=False), torch.no_grad():
                    if version.parse(version.parse(torch.__version__).base_version) >= version.parse("2.0.0"):
                        if isinstance(example_batch, dict):
                            jit_model = torch.jit.trace(jit_model, example_kwarg_inputs=example_batch, strict=False)
                        else:
                            jit_model = torch.jit.trace(
                                jit_model,
                                example_kwarg_inputs={key: example_batch[key] for key in example_batch},
                                strict=False,
                            )
                    else:
                        jit_inputs = []
                        for key in example_batch:
                            example_tensor = torch.ones_like(example_batch[key])
                            jit_inputs.append(example_tensor)
                        jit_inputs = tuple(jit_inputs)
                        jit_model = torch.jit.trace(jit_model, jit_inputs, strict=False)
                jit_model = torch.jit.freeze(jit_model)
                with torch.no_grad():
                    jit_model(**example_batch)
                    jit_model(**example_batch)
                model = jit_model
                self.use_cpu_amp = False
            except (RuntimeError, TypeError, ValueError, NameError, IndexError) as e:
                logger.warning(f"failed to use PyTorch jit mode due to: {e}.")

        return model

    def ipex_optimize_model(self, model, training=False, dtype=torch.float32):
        if not is_ipex_available():
            raise ImportError(
                "Using IPEX but IPEX is not installed or IPEX's version does not match current PyTorch, please refer"
                " to https://github.com/intel/intel-extension-for-pytorch."
            )

        import intel_extension_for_pytorch as ipex

        if not training:
            model.eval()
            dtype = torch.bfloat16 if not self.is_in_train and self.args.bf16_full_eval else dtype
            # conv_bn_folding is disabled as it fails in symbolic tracing, resulting in ipex warnings
            model = ipex.optimize(model, dtype=dtype, level="O1", conv_bn_folding=False, inplace=not self.is_in_train)
        else:
            if not model.training:
                model.train()
            model, self.optimizer = ipex.optimize(
                model, dtype=dtype, optimizer=self.optimizer, inplace=True, level="O1"
            )

        return model

    def compare_trainer_and_checkpoint_args(self, training_args, trainer_state):
        attributes_map = {
            "logging_steps": "logging_steps",
            "eval_steps": "eval_steps",
            "save_steps": "save_steps",
        }

        has_warning = False
        warning_str = "Warning: The following arguments do not match the ones in the `trainer_state.json` within the checkpoint directory: "
        for arg_attr, state_attr in attributes_map.items():
            arg_value = getattr(training_args, arg_attr, None)
            state_value = getattr(trainer_state, state_attr, None)

            if arg_value is not None and state_value is not None and arg_value != state_value:
                warning_str += f"\n\t{arg_attr}: {arg_value} (from args) != {state_value} (from trainer_state.json)"
                has_warning = True

        # train bs is special as we need to account for multi-GPU
        train_bs_args = training_args.per_device_train_batch_size
        train_bs_state = trainer_state.train_batch_size // max(1, training_args.n_gpu)

        if train_bs_args != train_bs_state:
            warning_str += f"\n\tper_device_train_batch_size: {train_bs_args} (from args) != {train_bs_state} (from trainer_state.json)"
            has_warning = True

        if has_warning:
            logger.warning_once(warning_str)

    def _wrap_model(self, model, training=True, dataloader=None):
        if self.args.use_ipex:
            dtype = torch.bfloat16 if self.use_cpu_amp else torch.float32
            model = self.ipex_optimize_model(model, training, dtype=dtype)

        if is_sagemaker_mp_enabled():
            # Wrapping the base model twice in a DistributedModel will raise an error.
            if isinstance(self.model_wrapped, smp.model.DistributedModel):
                return self.model_wrapped
            return smp.DistributedModel(model, backward_passes_per_step=self.args.gradient_accumulation_steps)

        # train/eval could be run multiple-times - if already wrapped, don't re-wrap it again
        if self.accelerator.unwrap_model(model) is not model:
            return model

        # Mixed precision training with apex (torch < 1.6)
        if self.use_apex and training:
            model, self.optimizer = amp.initialize(model, self.optimizer, opt_level=self.args.fp16_opt_level)

        # Multi-gpu training (should be after apex fp16 initialization) / 8bit models does not support DDP
        if self.args.n_gpu > 1 and not getattr(model, "is_loaded_in_8bit", False):
            model = nn.DataParallel(model)

        if self.args.jit_mode_eval:
            start_time = time.time()
            model = self.torch_jit_model_eval(model, dataloader, training)
            self.jit_compilation_time = round(time.time() - start_time, 4)

        # Note: in torch.distributed mode, there's no point in wrapping the model
        # inside a DistributedDataParallel as we'll be under `no_grad` anyways.
        if not training:
            return model

        # Distributed training (should be after apex fp16 initialization)
        # Distributed training using PyTorch FSDP
        if self.is_fsdp_xla_enabled:
            try:
                from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as FSDP
                from torch_xla.distributed.fsdp import checkpoint_module
                from torch_xla.distributed.fsdp.wrap import (
                    size_based_auto_wrap_policy,
                    transformer_auto_wrap_policy,
                )

                if self.is_fsdp_xla_v2_enabled:
                    from torch_xla.experimental.spmd_fully_sharded_data_parallel import (
                        SpmdFullyShardedDataParallel as FSDPv2,
                    )
            except ImportError:
                raise ImportError("Missing XLA FSDP related module; please make sure to use torch-xla >= 2.0.")
            auto_wrap_policy = None
            auto_wrapper_callable = None
            default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", None)
            fsdp_transformer_layer_cls_to_wrap = self.args.fsdp_config.get(
                "transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap
            )

            if self.args.fsdp_config["min_num_params"] > 0:
                auto_wrap_policy = functools.partial(
                    size_based_auto_wrap_policy, min_num_params=self.args.fsdp_config["min_num_params"]
                )
            elif fsdp_transformer_layer_cls_to_wrap is not None:
                transformer_cls_to_wrap = set()
                for layer_class in fsdp_transformer_layer_cls_to_wrap:
                    transformer_cls = get_module_class_from_name(model, layer_class)
                    if transformer_cls is None:
                        raise Exception("Could not find the transformer layer class to wrap in the model.")
                    else:
                        transformer_cls_to_wrap.add(transformer_cls)

                auto_wrap_policy = functools.partial(
                    transformer_auto_wrap_policy,
                    # Transformer layer class to wrap
                    transformer_layer_cls=transformer_cls_to_wrap,
                )
            fsdp_kwargs = self.args.xla_fsdp_config
            if self.args.fsdp_config["xla_fsdp_grad_ckpt"]:
                if model.config.use_cache:
                    logger.warning_once(
                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
                    )
                    model.config.use_cache = False

                # Apply gradient checkpointing to auto-wrapped sub-modules if specified
                def auto_wrapper_callable(m, *args, **kwargs):
                    target_cls = FSDP if not self.is_fsdp_xla_v2_enabled else FSDPv2
                    return target_cls(checkpoint_module(m), *args, **kwargs)

            # Wrap the base model with an outer FSDP wrapper
            if self.is_fsdp_xla_v2_enabled:

                def shard_output(output, mesh):
                    from .modeling_outputs import CausalLMOutputWithPast

                    real_output = None
                    if isinstance(output, torch.Tensor):
                        real_output = output
                    elif isinstance(output, tuple):
                        real_output = output[0]
                    elif isinstance(output, CausalLMOutputWithPast):
                        real_output = output.logits

                    if real_output is None:
                        raise ValueError("Something went wrong, the output of the model shouldn't be `None`")
                    xs.mark_sharding(real_output, mesh, ("fsdp", None, None))

                self.model = model = FSDPv2(
                    model,
                    shard_output=shard_output,
                    auto_wrap_policy=auto_wrap_policy,
                    auto_wrapper_callable=auto_wrapper_callable,
                )
            else:
                self.model = model = FSDP(
                    model,
                    auto_wrap_policy=auto_wrap_policy,
                    auto_wrapper_callable=auto_wrapper_callable,
                    **fsdp_kwargs,
                )

            # Patch `xm.optimizer_step` should not reduce gradients in this case,
            # as FSDP does not need gradient reduction over sharded parameters.
            def patched_optimizer_step(optimizer, barrier=False, optimizer_args={}):
                loss = optimizer.step(**optimizer_args)
                if barrier:
                    xm.mark_step()
                return loss

            xm.optimizer_step = patched_optimizer_step
        elif is_sagemaker_dp_enabled():
            model = nn.parallel.DistributedDataParallel(
                model, device_ids=[int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))]
            )
        elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            if is_torch_neuroncore_available():
                return model
            kwargs = {}
            if self.args.ddp_find_unused_parameters is not None:
                kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters
            elif isinstance(model, PreTrainedModel):
                # find_unused_parameters breaks checkpointing as per
                # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021
                kwargs["find_unused_parameters"] = not model.is_gradient_checkpointing
            else:
                kwargs["find_unused_parameters"] = True

            if self.args.ddp_bucket_cap_mb is not None:
                kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb

            if self.args.ddp_broadcast_buffers is not None:
                kwargs["broadcast_buffers"] = self.args.ddp_broadcast_buffers

            self.accelerator.ddp_handler = DistributedDataParallelKwargs(**kwargs)

        return model

    def train(
        self,
        resume_from_checkpoint: Optional[Union[str, bool]] = None,
        trial: Union["optuna.Trial", Dict[str, Any]] = None,
        ignore_keys_for_eval: Optional[List[str]] = None,
        **kwargs,
    ):
        """
        Main training entry point.

        Args:
            resume_from_checkpoint (`str` or `bool`, *optional*):
                If a `str`, local path to a saved checkpoint as saved by a previous instance of [`Trainer`]. If a
                `bool` and equals `True`, load the last checkpoint in *args.output_dir* as saved by a previous instance
                of [`Trainer`]. If present, training will resume from the model/optimizer/scheduler states loaded here.
            trial (`optuna.Trial` or `Dict[str, Any]`, *optional*):
                The trial run or the hyperparameter dictionary for hyperparameter search.
            ignore_keys_for_eval (`List[str]`, *optional*)
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions for evaluation during the training.
            kwargs (`Dict[str, Any]`, *optional*):
                Additional keyword arguments used to hide deprecated arguments
        """
        if resume_from_checkpoint is False:
            resume_from_checkpoint = None

        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

        args = self.args

        self.is_in_train = True

        # Attach NEFTune hooks if necessary
        if self.neftune_noise_alpha is not None:
            self.model = self._activate_neftune(self.model)

        # do_train is not a reliable argument, as it might not be set and .train() still called, so
        # the following is a workaround:
        if (args.fp16_full_eval or args.bf16_full_eval) and not args.do_train and not self.is_model_parallel:
            self._move_model_to_device(self.model, args.device)

        if "model_path" in kwargs:
            resume_from_checkpoint = kwargs.pop("model_path")
            warnings.warn(
                "`model_path` is deprecated and will be removed in a future version. Use `resume_from_checkpoint` "
                "instead.",
                FutureWarning,
            )
        if len(kwargs) > 0:
            raise TypeError(f"train() got unexpected keyword arguments: {', '.join(list(kwargs.keys()))}.")
        # This might change the seed so needs to run first.
        self._hp_search_setup(trial)
        self._train_batch_size = self.args.train_batch_size

        # Model re-init
        model_reloaded = False
        if self.model_init is not None:
            # Seed must be set before instantiating the model when using model_init.
            enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed)
            self.model = self.call_model_init(trial)
            model_reloaded = True
            # Reinitializes optimizer and scheduler
            self.optimizer, self.lr_scheduler = None, None

        # Load potential model checkpoint
        if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint:
            resume_from_checkpoint = get_last_checkpoint(args.output_dir)
            if resume_from_checkpoint is None:
                raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})")

        if resume_from_checkpoint is not None:
            if not is_sagemaker_mp_enabled() and not self.is_deepspeed_enabled and not self.is_fsdp_enabled:
                self._load_from_checkpoint(resume_from_checkpoint)
            # In case of repeating the find_executable_batch_size, set `self._train_batch_size` properly
            state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            if state.train_batch_size is not None:
                self._train_batch_size = state.train_batch_size

        # If model was re-initialized, put it on the right device and update self.model_wrapped
        if model_reloaded:
            if self.place_model_on_device:
                self._move_model_to_device(self.model, args.device)
            self.model_wrapped = self.model

        inner_training_loop = find_executable_batch_size(
            self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size
        )
        if args.push_to_hub:
            try:
                # Disable progress bars when uploading models during checkpoints to avoid polluting stdout
                hf_hub_utils.disable_progress_bars()
                return inner_training_loop(
                    args=args,
                    resume_from_checkpoint=resume_from_checkpoint,
                    trial=trial,
                    ignore_keys_for_eval=ignore_keys_for_eval,
                )
            finally:
                hf_hub_utils.enable_progress_bars()
        else:
            return inner_training_loop(
                args=args,
                resume_from_checkpoint=resume_from_checkpoint,
                trial=trial,
                ignore_keys_for_eval=ignore_keys_for_eval,
            )

    def _inner_training_loop(
        self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
    ):
        self.accelerator.free_memory()
        self._train_batch_size = batch_size
        if self.args.auto_find_batch_size:
            if self.state.train_batch_size != self._train_batch_size:
                from accelerate.utils import release_memory

                (self.model_wrapped,) = release_memory(self.model_wrapped)
                self.model_wrapped = self.model

                # Check for DeepSpeed *after* the intial pass and modify the config
                if self.is_deepspeed_enabled:
                    # Temporarily unset `self.args.train_batch_size`
                    original_bs = self.args.per_device_train_batch_size
                    self.args.per_device_train_batch_size = self._train_batch_size // max(1, self.args.n_gpu)
                    self.propagate_args_to_deepspeed(True)
                    self.args.per_device_train_batch_size = original_bs
            self.state.train_batch_size = self._train_batch_size
        logger.debug(f"Currently training with a batch size of: {self._train_batch_size}")
        # Data loader and number of training steps
        train_dataloader = self.get_train_dataloader()
        if self.is_fsdp_xla_v2_enabled:
            train_dataloader = tpu_spmd_dataloader(train_dataloader)

        # Setting up training control variables:
        # number of training epochs: num_train_epochs
        # number of training steps per epoch: num_update_steps_per_epoch
        # total number of training steps to execute: max_steps
        total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size

        len_dataloader = None
        num_train_tokens = None
        if has_length(train_dataloader):
            len_dataloader = len(train_dataloader)
            num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
            num_examples = self.num_examples(train_dataloader)
            if args.max_steps > 0:
                max_steps = args.max_steps
                num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
                    args.max_steps % num_update_steps_per_epoch > 0
                )
                # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
                # the best we can do.
                num_train_samples = args.max_steps * total_train_batch_size
                if args.include_tokens_per_second:
                    num_train_tokens = (
                        self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps
                    )
            else:
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                num_train_epochs = math.ceil(args.num_train_epochs)
                num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
                if args.include_tokens_per_second:
                    num_train_tokens = self.num_tokens(train_dataloader) * args.num_train_epochs
        elif args.max_steps > 0:  # Rely on max_steps when dataloader does not have a working size
            max_steps = args.max_steps
            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
            num_train_epochs = sys.maxsize
            num_update_steps_per_epoch = max_steps
            num_examples = total_train_batch_size * args.max_steps
            num_train_samples = args.max_steps * total_train_batch_size
            if args.include_tokens_per_second:
                num_train_tokens = self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps
        else:
            raise ValueError(
                "args.max_steps must be set to a positive value if dataloader does not have a length, was"
                f" {args.max_steps}"
            )

        if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
            if self.args.n_gpu > 1:
                # nn.DataParallel(model) replicates the model, creating new variables and module
                # references registered here no longer work on other gpus, breaking the module
                raise ValueError(
                    "Currently --debug underflow_overflow is not supported under DP. Please use DDP"
                    " (torchrun or torch.distributed.launch (deprecated))."
                )
            else:
                debug_overflow = DebugUnderflowOverflow(self.model)  # noqa

        delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled

        # We need to reset the scheduler, as its parameters may be different on subsequent calls
        if self._created_lr_scheduler:
            self.lr_scheduler = None
            self._created_lr_scheduler = False

        if self.is_deepspeed_enabled:
            self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)

        if not delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        self.state = TrainerState(
            stateful_callbacks=[
                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
            ]
        )
        self.state.is_hyper_param_search = trial is not None
        self.state.train_batch_size = self._train_batch_size

        # Compute absolute values for logging, eval, and save if given as ratio
        if args.logging_steps is not None:
            if args.logging_steps < 1:
                self.state.logging_steps = math.ceil(max_steps * args.logging_steps)
            else:
                self.state.logging_steps = args.logging_steps
        if args.eval_steps is not None:
            if args.eval_steps < 1:
                self.state.eval_steps = math.ceil(max_steps * args.eval_steps)
            else:
                self.state.eval_steps = args.eval_steps
        if args.save_steps is not None:
            if args.save_steps < 1:
                self.state.save_steps = math.ceil(max_steps * args.save_steps)
            else:
                self.state.save_steps = args.save_steps

        # Activate gradient checkpointing if needed
        if args.gradient_checkpointing:
            self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=args.gradient_checkpointing_kwargs)

        model = self._wrap_model(self.model_wrapped)

        # as the model is wrapped, don't use `accelerator.prepare`
        # this is for unhandled cases such as
        # FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX
        use_accelerator_prepare = True if model is self.model else False

        if use_accelerator_prepare and self.is_fsdp_enabled:
            # In case of auto_find_batch_size=True
            # Remove FSDP wrapping from sub-models.
            self.model = unwrap_model(self.model, recursive=True)

        if delay_optimizer_creation:
            if use_accelerator_prepare:
                # configure fsdp plugin for qlora if any
                self._fsdp_qlora_plugin_updates()
                if self.accelerator.mixed_precision != "fp8":
                    self.model = self.accelerator.prepare(self.model)
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        # prepare using `accelerator` prepare
        if use_accelerator_prepare:
            self.model.train()
            if hasattr(self.lr_scheduler, "step"):
                if self.use_apex:
                    model = self.accelerator.prepare(self.model)
                else:
                    model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
            else:
                # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config.
                model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
                    self.model, self.optimizer, self.lr_scheduler
                )
        elif self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
            # In this case we are in DDP + LOMO, which should be supported
            self.optimizer = self.accelerator.prepare(self.optimizer)

        if self.is_fsdp_enabled:
            self.model = self.model_wrapped = model

        # for the rest of this function `model` is the outside model, whether it was wrapped or not
        if model is not self.model:
            self.model_wrapped = model

        # backward compatibility
        if self.is_deepspeed_enabled:
            self.deepspeed = self.model_wrapped

        # ckpt loading
        if resume_from_checkpoint is not None:
            if self.is_deepspeed_enabled:
                deepspeed_load_checkpoint(
                    self.model_wrapped, resume_from_checkpoint, load_module_strict=not _is_peft_model(self.model)
                )
            elif is_sagemaker_mp_enabled() or self.is_fsdp_enabled:
                self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)

        # Check if saved optimizer or scheduler states exist
        self._load_optimizer_and_scheduler(resume_from_checkpoint)

        # important: at this point:
        # self.model         is the Transformers Model
        # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model),
        # FSDP(Transformers Model), Dynamo Optimized Module(Transformers Model) etc.

        # Train!
        logger.info("***** Running training *****")
        logger.info(f"  Num examples = {num_examples:,}")
        logger.info(f"  Num Epochs = {num_train_epochs:,}")
        logger.info(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}")
        if self.args.per_device_train_batch_size != self._train_batch_size:
            logger.info(f"  Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}")
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}")
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        logger.info(f"  Total optimization steps = {max_steps:,}")
        logger.info(f"  Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}")

        self.state.epoch = 0
        start_time = time.time()
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
        steps_trained_progress_bar = None

        # Check if continuing training from a checkpoint
        if resume_from_checkpoint is not None and os.path.isfile(
            os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
        ):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            self.compare_trainer_and_checkpoint_args(self.args, self.state)
            self._load_callback_state()
            epochs_trained = int(self.state.global_step // num_update_steps_per_epoch)
            if not args.ignore_data_skip:
                steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
            else:
                steps_trained_in_current_epoch = 0

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
            if not args.ignore_data_skip:
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first"
                    f" {steps_trained_in_current_epoch} batches in the first epoch."
                )

        # Update the references
        self.callback_handler.model = self.model
        self.callback_handler.optimizer = self.optimizer
        self.callback_handler.lr_scheduler = self.lr_scheduler
        self.callback_handler.train_dataloader = train_dataloader
        if self.hp_name is not None and self._trial is not None:
            # use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial
            # parameter to Train when using DDP.
            self.state.trial_name = self.hp_name(self._trial)
        if trial is not None:
            assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial
            self.state.trial_params = hp_params(assignments)
        else:
            self.state.trial_params = None
        # This should be the same if the state has been saved but in case the training arguments changed, it's safer
        # to set this after the load.
        self.state.max_steps = max_steps
        self.state.num_train_epochs = num_train_epochs
        self.state.is_local_process_zero = self.is_local_process_zero()
        self.state.is_world_process_zero = self.is_world_process_zero()

        # tr_loss is a tensor to avoid synchronization of TPUs through .item()
        tr_loss = torch.tensor(0.0).to(args.device)
        # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses
        self._total_loss_scalar = 0.0
        self._globalstep_last_logged = self.state.global_step
        model.zero_grad()
        grad_norm: Optional[float] = None
        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)

        if args.eval_on_start:
            self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True)

        for epoch in range(epochs_trained, num_train_epochs):
            epoch_dataloader = train_dataloader
            if hasattr(epoch_dataloader, "set_epoch"):
                epoch_dataloader.set_epoch(epoch)

            # Reset the past mems state at the beginning of each epoch if necessary.
            if args.past_index >= 0:
                self._past = None

            steps_in_epoch = (
                len(epoch_dataloader)
                if len_dataloader is not None
                else args.max_steps * args.gradient_accumulation_steps
            )
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)

            if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
                self._load_rng_state(resume_from_checkpoint)

            rng_to_sync = False
            steps_skipped = 0
            if steps_trained_in_current_epoch > 0:
                epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch)
                steps_skipped = steps_trained_in_current_epoch
                steps_trained_in_current_epoch = 0
                rng_to_sync = True

            step = -1
            epoch_iterator = iter(epoch_dataloader)
            # We chunkify the epoch iterator into gradient accumulation steps `n` batches
            remainder = num_examples % args.gradient_accumulation_steps
            if remainder == 0:
                remainder = args.gradient_accumulation_steps
            update_step = -1
            total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1
            for _ in range(total_updates):
                update_step += 1
                num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder
                batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches)
                for i, inputs in enumerate(batch_samples):
                    step += 1
                    do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch
                    # Since we perform prefetching, we need to manually set sync_gradients
                    if not do_sync_step:
                        self.accelerator.gradient_state._set_sync_gradients(False)
                    else:
                        self.accelerator.gradient_state._set_sync_gradients(True)

                    if self.args.include_num_input_tokens_seen:
                        main_input_name = getattr(self.model, "main_input_name", "input_ids")
                        if main_input_name not in inputs:
                            logger.warning(
                                "Tried to track the number of tokens seen, however the current model is "
                                "not configured properly to know what item is the input. To fix this, add "
                                "a `main_input_name` attribute to the model class you are using."
                            )
                        else:
                            input_tokens = inputs[main_input_name].numel()
                            input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64)
                            self.state.num_input_tokens_seen += (
                                self.accelerator.gather(input_tokens).sum().cpu().item()
                            )
                    if rng_to_sync:
                        self._load_rng_state(resume_from_checkpoint)
                        rng_to_sync = False

                    # Skip past any already trained steps if resuming training
                    if steps_trained_in_current_epoch > 0:
                        steps_trained_in_current_epoch -= 1
                        if steps_trained_progress_bar is not None:
                            steps_trained_progress_bar.update(1)
                        if steps_trained_in_current_epoch == 0:
                            self._load_rng_state(resume_from_checkpoint)
                        continue
                    elif steps_trained_progress_bar is not None:
                        steps_trained_progress_bar.close()
                        steps_trained_progress_bar = None

                    if step % args.gradient_accumulation_steps == 0:
                        self.control = self.callback_handler.on_step_begin(args, self.state, self.control)

                    # We explicitly want to avoid relying on `accelerator.accumulate` for generation training
                    context = (
                        functools.partial(self.accelerator.no_sync, model=model)
                        if i != len(batch_samples) - 1
                        and self.accelerator.distributed_type != DistributedType.DEEPSPEED
                        else contextlib.nullcontext
                    )
                    with context():
                        tr_loss_step = self.training_step(model, inputs, num_items_in_batch)

                    if (
                        args.logging_nan_inf_filter
                        and not is_torch_xla_available()
                        and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
                    ):
                        # if loss is nan or inf simply add the average of previous logged losses
                        tr_loss = tr_loss + tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
                    else:
                        if tr_loss.device != tr_loss_step.device:
                            raise ValueError(
                                f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}"
                            )
                        tr_loss = tr_loss + tr_loss_step

                    self.current_flos += float(self.floating_point_ops(inputs))

                    if do_sync_step:
                        # Since we perform prefetching, we need to manually set sync_gradients to True
                        self.accelerator.gradient_state._set_sync_gradients(True)

                        # Gradient clipping
                        if args.max_grad_norm is not None and args.max_grad_norm > 0:
                            # deepspeed does its own clipping

                            if is_sagemaker_mp_enabled() and args.fp16:
                                _grad_norm = self.optimizer.clip_master_grads(args.max_grad_norm)
                            elif self.use_apex:
                                # Revert to normal clipping otherwise, handling Apex or full precision
                                _grad_norm = nn.utils.clip_grad_norm_(
                                    amp.master_params(self.optimizer),
                                    args.max_grad_norm,
                                )
                            else:
                                _grad_norm = self.accelerator.clip_grad_norm_(
                                    model.parameters(),
                                    args.max_grad_norm,
                                )

                            if (
                                is_accelerate_available()
                                and self.accelerator.distributed_type == DistributedType.DEEPSPEED
                            ):
                                grad_norm = model.get_global_grad_norm()
                                # In some cases the grad norm may not return a float
                                if hasattr(grad_norm, "item"):
                                    grad_norm = grad_norm.item()
                            else:
                                grad_norm = _grad_norm

                        self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control)

                        self.optimizer.step()

                        self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control)

                        optimizer_was_run = not self.accelerator.optimizer_step_was_skipped
                        if optimizer_was_run:
                            # Delay optimizer scheduling until metrics are generated
                            if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                                self.lr_scheduler.step()

                        model.zero_grad()
                        self.state.global_step += 1
                        self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
                        self.control = self.callback_handler.on_step_end(args, self.state, self.control)
                        self._maybe_log_save_evaluate(
                            tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time
                        )
                    else:
                        self.control = self.callback_handler.on_substep_end(args, self.state, self.control)

                    # PyTorch/XLA relies on the data loader to insert the mark_step for
                    # each step. Since we are breaking the loop early, we need to manually
                    # insert the mark_step here.
                    if self.control.should_epoch_stop or self.control.should_training_stop:
                        if is_torch_xla_available():
                            xm.mark_step()
                        break
                # We also need to break out of the nested loop
                if self.control.should_epoch_stop or self.control.should_training_stop:
                    if is_torch_xla_available():
                        xm.mark_step()
                    break
            if step < 0:
                logger.warning(
                    "There seems not to be a single sample in your epoch_iterator, stopping training at step"
                    f" {self.state.global_step}! This is expected if you're using an IterableDataset and set"
                    f" num_steps ({max_steps}) higher than the number of available samples."
                )
                self.control.should_training_stop = True

            self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
            self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time)

            if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
                if is_torch_xla_available():
                    # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                    xm.master_print(met.metrics_report())
                else:
                    logger.warning(
                        "You enabled PyTorch/XLA debug metrics but you don't have a TPU "
                        "configured. Check your training configuration if this is unexpected."
                    )
            if self.control.should_training_stop:
                break

        if args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of training
            delattr(self, "_past")

        logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
        if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
            # Wait for everyone to get here so we are sure the model has been saved by process 0.
            if is_torch_xla_available():
                xm.rendezvous("load_best_model_at_end")
            elif args.parallel_mode == ParallelMode.DISTRIBUTED:
                dist.barrier()
            elif is_sagemaker_mp_enabled():
                smp.barrier()

            self._load_best_model()

        # add remaining tr_loss
        self._total_loss_scalar += tr_loss.item()
        effective_global_step = max(self.state.global_step, 0.001)  # Avoid ZeroDivisionError
        train_loss = self._total_loss_scalar / effective_global_step

        metrics = speed_metrics(
            "train",
            start_time,
            num_samples=num_train_samples,
            num_steps=self.state.max_steps,
            num_tokens=num_train_tokens,
        )
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
        metrics["train_loss"] = train_loss

        self.is_in_train = False

        self._memory_tracker.stop_and_update_metrics(metrics)

        self.log(metrics)

        run_dir = self._get_output_dir(trial)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir)

        # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save.
        if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
            for checkpoint in checkpoints_sorted:
                if not os.path.samefile(checkpoint, self.state.best_model_checkpoint):
                    logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
                    shutil.rmtree(checkpoint, ignore_errors=True)

        self.control = self.callback_handler.on_train_end(args, self.state, self.control)

        # Wait for the checkpoint to be uploaded.
        self._finish_current_push()

        # After training we make sure to retrieve back the original forward pass method
        # for the embedding layer by removing the forward post hook.
        if self.neftune_noise_alpha is not None:
            self._deactivate_neftune(self.model)

        return TrainOutput(self.state.global_step, train_loss, metrics)

    def _get_output_dir(self, trial):
        if self.hp_search_backend is not None and trial is not None:
            if self.hp_search_backend == HPSearchBackend.OPTUNA:
                run_id = trial.number
            elif self.hp_search_backend == HPSearchBackend.RAY:
                import ray.train

                run_id = ray.train.get_context().get_trial_id()
            elif self.hp_search_backend == HPSearchBackend.SIGOPT:
                run_id = trial.id
            elif self.hp_search_backend == HPSearchBackend.WANDB:
                import wandb

                run_id = wandb.run.id
            run_name = self.hp_name(trial) if self.hp_name is not None else f"run-{run_id}"
            run_dir = os.path.join(self.args.output_dir, run_name)
        else:
            run_dir = self.args.output_dir
        return run_dir

    def _load_from_checkpoint(self, resume_from_checkpoint, model=None):
        if model is None:
            model = self.model

        config_file = os.path.join(resume_from_checkpoint, CONFIG_NAME)
        adapter_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_WEIGHTS_NAME)
        adapter_safe_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)
        weights_file = os.path.join(resume_from_checkpoint, WEIGHTS_NAME)
        weights_index_file = os.path.join(resume_from_checkpoint, WEIGHTS_INDEX_NAME)
        safe_weights_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_NAME)
        safe_weights_index_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_INDEX_NAME)
        is_fsdp_ckpt = os.path.isdir(resume_from_checkpoint) and (
            # this checks the FSDP state dict when `SHARDED_STATE_DICT` is used
            any(
                FSDP_MODEL_NAME in folder_name
                for folder_name in os.listdir(resume_from_checkpoint)
                if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name))
            )
            # this checks the FSDP state dict when `FULL_STATE_DICT` is used
            or os.path.isfile(os.path.join(resume_from_checkpoint, f"{FSDP_MODEL_NAME}.bin"))
        )
        # if multiple adapters exist, they get saved in sub directories
        adapter_subdirs = (
            [
                folder_name
                for folder_name in os.listdir(resume_from_checkpoint)
                if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name))
                and (
                    os.path.isfile(os.path.join(resume_from_checkpoint, folder_name, ADAPTER_WEIGHTS_NAME))
                    or os.path.isfile(os.path.join(resume_from_checkpoint, folder_name, ADAPTER_SAFE_WEIGHTS_NAME))
                )
            ]
            if os.path.isdir(resume_from_checkpoint)
            else []
        )

        if is_fsdp_ckpt and not self.is_fsdp_enabled:
            raise ValueError(f"Checkpoint found at {resume_from_checkpoint} is only supported when using PyTorch FSDP")

        if not (
            any(
                os.path.isfile(f)
                for f in [
                    weights_file,
                    safe_weights_file,
                    weights_index_file,
                    safe_weights_index_file,
                    adapter_weights_file,
                    adapter_safe_weights_file,
                ]
            )
            or is_fsdp_ckpt
            or adapter_subdirs
        ):
            raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}")

        logger.info(f"Loading model from {resume_from_checkpoint}.")

        if os.path.isfile(config_file):
            config = PretrainedConfig.from_json_file(config_file)
            checkpoint_version = config.transformers_version
            if checkpoint_version is not None and checkpoint_version != __version__:
                logger.warning(
                    f"You are resuming training from a checkpoint trained with {checkpoint_version} of "
                    f"Transformers but your current version is {__version__}. This is not recommended and could "
                    "yield to errors or unwanted behaviors."
                )

        if os.path.isfile(weights_file) or os.path.isfile(safe_weights_file) or is_fsdp_ckpt:
            weights_only_kwarg = {"weights_only": True}
            # If the model is on the GPU, it still works!
            if is_sagemaker_mp_enabled():
                if os.path.isfile(os.path.join(resume_from_checkpoint, "user_content.pt")):
                    # If the 'user_content.pt' file exists, load with the new smp api.
                    # Checkpoint must have been saved with the new smp api.
                    smp.resume_from_checkpoint(
                        path=resume_from_checkpoint, tag=WEIGHTS_NAME, partial=False, load_optimizer=False
                    )
                else:
                    # If the 'user_content.pt' file does NOT exist, load with the old smp api.
                    # Checkpoint must have been saved with the old smp api.
                    if hasattr(self.args, "fp16") and self.args.fp16 is True:
                        logger.warning(
                            "Enabling FP16 and loading from smp < 1.10 checkpoint together is not suppported."
                        )
                    state_dict = torch.load(
                        weights_file,
                        map_location="cpu",
                        **weights_only_kwarg,
                    )
                    # Required for smp to not auto-translate state_dict from hf to smp (is already smp).
                    state_dict["_smp_is_partial"] = False
                    load_result = model.load_state_dict(state_dict, strict=True)
                    # release memory
                    del state_dict
            elif self.is_fsdp_enabled:
                load_fsdp_model(
                    self.accelerator.state.fsdp_plugin,
                    self.accelerator,
                    model,
                    resume_from_checkpoint,
                    **_get_fsdp_ckpt_kwargs(),
                )
            else:
                # We load the model state dict on the CPU to avoid an OOM error.
                if self.args.save_safetensors and os.path.isfile(safe_weights_file):
                    state_dict = safetensors.torch.load_file(safe_weights_file, device="cpu")
                else:
                    state_dict = torch.load(
                        weights_file,
                        map_location="cpu",
                        **weights_only_kwarg,
                    )

                # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963
                # which takes *args instead of **kwargs
                load_result = model.load_state_dict(state_dict, False)
                # release memory
                del state_dict
                self._issue_warnings_after_load(load_result)

        # Load adapters following PR # 24096
        elif _is_peft_model(model):
            # If train a model using PEFT & LoRA, assume that adapter have been saved properly.
            # TODO: in the future support only specific min PEFT versions
            if (hasattr(model, "active_adapter") or hasattr(model, "active_adapters")) and hasattr(
                model, "load_adapter"
            ):
                if os.path.exists(resume_from_checkpoint):
                    # For BC for older PEFT versions
                    if hasattr(model, "active_adapters"):
                        active_adapters = model.active_adapters
                        if len(active_adapters) > 1:
                            logger.warning("Multiple active adapters detected will only consider the first adapter")
                        active_adapter = active_adapters[0]
                    else:
                        active_adapter = model.active_adapter

                    if adapter_subdirs:
                        for subdir_name in adapter_subdirs:
                            peft_id = os.path.join(resume_from_checkpoint, subdir_name)
                            model.load_adapter(peft_id, subdir_name, is_trainable=(subdir_name == active_adapter))
                        model.set_adapter(active_adapter)
                    else:
                        model.load_adapter(resume_from_checkpoint, active_adapter, is_trainable=True)
                else:
                    logger.warning(
                        "The intermediate checkpoints of PEFT may not be saved correctly, "
                        f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. "
                        "Check some examples here: https://github.com/huggingface/peft/issues/96"
                    )
            else:
                logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed")
        else:
            # We load the sharded checkpoint
            load_result = load_sharded_checkpoint(
                model, resume_from_checkpoint, strict=is_sagemaker_mp_enabled(), prefer_safe=self.args.save_safetensors
            )
            if not is_sagemaker_mp_enabled():
                self._issue_warnings_after_load(load_result)

    def _load_best_model(self):
        logger.info(f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric}).")
        best_model_path = os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME)
        best_safe_model_path = os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_NAME)
        best_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_WEIGHTS_NAME)
        best_safe_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME)

        model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
        if self.is_deepspeed_enabled:
            deepspeed_load_checkpoint(
                self.model_wrapped,
                self.state.best_model_checkpoint,
                load_module_strict=not _is_peft_model(self.model),
            )
        elif self.is_fsdp_enabled:
            load_result = load_fsdp_model(
                self.accelerator.state.fsdp_plugin,
                self.accelerator,
                model,
                self.state.best_model_checkpoint,
                **_get_fsdp_ckpt_kwargs(),
            )
        elif (
            os.path.exists(best_model_path)
            or os.path.exists(best_safe_model_path)
            or os.path.exists(best_adapter_model_path)
            or os.path.exists(best_safe_adapter_model_path)
        ):
            has_been_loaded = True
            weights_only_kwarg = {"weights_only": True}
            if is_sagemaker_mp_enabled():
                if os.path.isfile(os.path.join(self.state.best_model_checkpoint, "user_content.pt")):
                    # If the 'user_content.pt' file exists, load with the new smp api.
                    # Checkpoint must have been saved with the new smp api.
                    smp.resume_from_checkpoint(
                        path=self.state.best_model_checkpoint,
                        tag=WEIGHTS_NAME,
                        partial=False,
                        load_optimizer=False,
                    )
                else:
                    # If the 'user_content.pt' file does NOT exist, load with the old smp api.
                    # Checkpoint must have been saved with the old smp api.
                    if self.args.save_safetensors and os.path.isfile(best_safe_model_path):
                        state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu")
                    else:
                        state_dict = torch.load(
                            best_model_path,
                            map_location="cpu",
                            **weights_only_kwarg,
                        )

                    state_dict["_smp_is_partial"] = False
                    load_result = model.load_state_dict(state_dict, strict=True)
            else:
                if _is_peft_model(model):
                    # If train a model using PEFT & LoRA, assume that adapter have been saved properly.
                    # TODO: in the future support only specific min PEFT versions
                    if (hasattr(model, "active_adapter") or hasattr(model, "active_adapters")) and hasattr(
                        model, "load_adapter"
                    ):
                        # For BC for older PEFT versions
                        if hasattr(model, "active_adapters"):
                            active_adapter = model.active_adapters[0]
                            if len(model.active_adapters) > 1:
                                logger.warning("Detected multiple active adapters, will only consider the first one")
                        else:
                            active_adapter = model.active_adapter

                        if os.path.exists(best_adapter_model_path) or os.path.exists(best_safe_adapter_model_path):
                            try:
                                model.load_adapter(self.state.best_model_checkpoint, active_adapter)
                            except RuntimeError as exc:
                                if model.peft_config[active_adapter].is_prompt_learning:
                                    # for context: https://github.com/huggingface/peft/issues/2256
                                    msg = (
                                        "When using prompt learning PEFT methods such as "
                                        f"{model.peft_config[active_adapter].peft_type.value}, setting "
                                        "load_best_model_at_end=True can lead to errors, it is recommended "
                                        "to set this to False and to load the model manually from the checkpoint "
                                        "directory using PeftModel.from_pretrained(base_model, <path>) after training "
                                        "has finished."
                                    )
                                    raise RuntimeError(msg) from exc
                                else:
                                    raise
                            # Load_adapter has no return value present, modify it when appropriate.
                            from torch.nn.modules.module import _IncompatibleKeys

                            load_result = _IncompatibleKeys([], [])
                        else:
                            logger.warning(
                                "The intermediate checkpoints of PEFT may not be saved correctly, "
                                f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. "
                                "Check some examples here: https://github.com/huggingface/peft/issues/96"
                            )
                            has_been_loaded = False
                    else:
                        logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed")
                        has_been_loaded = False
                else:
                    # We load the model state dict on the CPU to avoid an OOM error.
                    if self.args.save_safetensors and os.path.isfile(best_safe_model_path):
                        state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu")
                    else:
                        state_dict = torch.load(
                            best_model_path,
                            map_location="cpu",
                            **weights_only_kwarg,
                        )

                    # If the model is on the GPU, it still works!
                    # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963
                    # which takes *args instead of **kwargs
                    load_result = model.load_state_dict(state_dict, False)
                if not is_sagemaker_mp_enabled() and has_been_loaded:
                    self._issue_warnings_after_load(load_result)
        elif os.path.exists(os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_INDEX_NAME)) or os.path.exists(
            os.path.join(self.state.best_model_checkpoint, WEIGHTS_INDEX_NAME)
        ):
            load_result = load_sharded_checkpoint(
                model, self.state.best_model_checkpoint, strict=is_sagemaker_mp_enabled()
            )
            if not is_sagemaker_mp_enabled():
                self._issue_warnings_after_load(load_result)
        else:
            logger.warning(
                f"Could not locate the best model at {best_model_path}, if you are running a distributed training "
                "on multiple nodes, you should activate `--save_on_each_node`."
            )

    def _issue_warnings_after_load(self, load_result):
        if len(load_result.missing_keys) != 0:
            if self.model._keys_to_ignore_on_save is not None and set(load_result.missing_keys) == set(
                self.model._keys_to_ignore_on_save
            ):
                self.model.tie_weights()
            else:
                logger.warning(f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}.")
        if len(load_result.unexpected_keys) != 0:
            logger.warning(
                f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}."
            )

    def _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler=False):
        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
        self._report_to_hp_search(trial, self.state.global_step, metrics)

        # Run delayed LR scheduler now that metrics are populated
        if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) and not skip_scheduler:
            metric_to_check = self.args.metric_for_best_model
            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"
            try:
                self.lr_scheduler.step(metrics[metric_to_check])
            except KeyError as exc:
                raise KeyError(
                    f"The `metric_for_best_model` training argument is set to '{metric_to_check}', "
                    f"which is not found in the evaluation metrics. "
                    f"The available evaluation metrics are: {list(metrics.keys())}. "
                    f"Please ensure that the `compute_metrics` function returns a dictionary that includes '{metric_to_check}' or "
                    f"consider changing the `metric_for_best_model` via the TrainingArguments."
                ) from exc
        return metrics

    def _maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time):
        if self.control.should_log and self.state.global_step > self._globalstep_last_logged:
            if is_torch_xla_available():
                xm.mark_step()

            logs: Dict[str, float] = {}

            # all_gather + mean() to get average loss over all processes
            tr_loss_scalar = self._nested_gather(tr_loss).mean().item()

            # reset tr_loss to zero
            tr_loss -= tr_loss

            logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4)
            if grad_norm is not None:
                logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm
            logs["learning_rate"] = self._get_learning_rate()

            self._total_loss_scalar += tr_loss_scalar
            self._globalstep_last_logged = self.state.global_step
            self.store_flos()

            self.log(logs, start_time)

        metrics = None
        if self.control.should_evaluate:
            metrics = self._evaluate(trial, ignore_keys_for_eval)
            is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial)

            if self.args.save_strategy == SaveStrategy.BEST:
                self.control.should_save = is_new_best_metric

        if self.control.should_save:
            self._save_checkpoint(model, trial)
            self.control = self.callback_handler.on_save(self.args, self.state, self.control)

    def _load_rng_state(self, checkpoint):
        # Load RNG states from `checkpoint`
        if checkpoint is None:
            return

        if self.args.world_size > 1:
            process_index = self.args.process_index
            rng_file = os.path.join(checkpoint, f"rng_state_{process_index}.pth")
            if not os.path.isfile(rng_file):
                logger.info(
                    f"Didn't find an RNG file for process {process_index}, if you are resuming a training that "
                    "wasn't launched in a distributed fashion, reproducibility is not guaranteed."
                )
                return
        else:
            rng_file = os.path.join(checkpoint, "rng_state.pth")
            if not os.path.isfile(rng_file):
                logger.info(
                    "Didn't find an RNG file, if you are resuming a training that was launched in a distributed "
                    "fashion, reproducibility is not guaranteed."
                )
                return

        with safe_globals():
            checkpoint_rng_state = torch.load(rng_file)
        random.setstate(checkpoint_rng_state["python"])
        np.random.set_state(checkpoint_rng_state["numpy"])
        torch.random.set_rng_state(checkpoint_rng_state["cpu"])
        if torch.cuda.is_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                torch.cuda.random.set_rng_state_all(checkpoint_rng_state["cuda"])
            else:
                try:
                    torch.cuda.random.set_rng_state(checkpoint_rng_state["cuda"])
                except Exception as e:
                    logger.info(
                        f"Didn't manage to set back the RNG states of the GPU because of the following error:\n {e}"
                        "\nThis won't yield the same results as if the training had not been interrupted."
                    )
        if is_torch_xla_available():
            xm.set_rng_state(checkpoint_rng_state["xla"])
        if is_torch_npu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                torch.npu.random.set_rng_state_all(checkpoint_rng_state["npu"])
            else:
                try:
                    torch.npu.random.set_rng_state(checkpoint_rng_state["npu"])
                except Exception as e:
                    logger.info(
                        f"Didn't manage to set back the RNG states of the NPU because of the following error:\n {e}"
                        "\nThis won't yield the same results as if the training had not been interrupted."
                    )
        if is_torch_mlu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                torch.mlu.random.set_rng_state_all(checkpoint_rng_state["mlu"])
            else:
                try:
                    torch.mlu.random.set_rng_state(checkpoint_rng_state["mlu"])
                except Exception as e:
                    logger.info(
                        f"Didn't manage to set back the RNG states of the MLU because of the following error:\n {e}"
                        "\nThis won't yield the same results as if the training had not been interrupted."
                    )
        if is_torch_musa_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                torch.musa.set_rng_state_all(checkpoint_rng_state["musa"])
            else:
                try:
                    torch.musa.set_rng_state(checkpoint_rng_state["musa"])
                except Exception as e:
                    logger.info(
                        f"Didn't manage to set back the RNG states of the MUSA because of the following error:\n {e}"
                        "\nThis won't yield the same results as if the training had not been interrupted."
                    )

    def _determine_best_metric(self, metrics, trial):
        """
        Determine if the model should be saved based on the evaluation metrics.

        Returns:
            bool: True if a new best metric was found, else False
        """
        is_new_best_metric = False

        if self.args.metric_for_best_model is not None:
            metric_to_check = self.args.metric_for_best_model

            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"

            try:
                metric_value = metrics[metric_to_check]
            except KeyError as exc:
                raise KeyError(
                    f"The `metric_for_best_model` training argument is set to '{metric_to_check}', which is not found in the evaluation metrics. "
                    f"The available evaluation metrics are: {list(metrics.keys())}. Consider changing the `metric_for_best_model` via the TrainingArguments."
                ) from exc

            operator = np.greater if self.args.greater_is_better else np.less

            if self.state.best_metric is None:
                self.state.best_metric = float("-inf") if self.args.greater_is_better else float("inf")

            if operator(metric_value, self.state.best_metric):
                run_dir = self._get_output_dir(trial=trial)
                checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
                output_dir = os.path.join(run_dir, checkpoint_folder)

                self.state.best_metric = metric_value
                self.state.best_model_checkpoint = output_dir

                is_new_best_metric = True

        return is_new_best_metric

    def _save_checkpoint(self, model, trial):
        # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we
        # want to save except FullyShardedDDP.
        # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"

        # Save model checkpoint
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"

        if self.hp_search_backend is None and trial is None:
            self.store_flos()

        run_dir = self._get_output_dir(trial=trial)
        output_dir = os.path.join(run_dir, checkpoint_folder)
        self.save_model(output_dir, _internal_call=True)

        if not self.args.save_only_model:
            # Save optimizer and scheduler
            self._save_optimizer_and_scheduler(output_dir)
            # Save RNG state
            self._save_rng_state(output_dir)

        # Save the Trainer state
        if self.args.should_save:
            # Update `ExportableState` callbacks and `TrainerControl` state to where we are currently
            for cb in [
                cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
            ]:
                cb_name = cb.__class__.__name__
                cb_state = cb.state()
                if isinstance(self.state.stateful_callbacks[cb_name], list):
                    self.state.stateful_callbacks[cb_name].append(cb_state)
                else:
                    self.state.stateful_callbacks[cb_name] = cb_state
            self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))

        if self.args.push_to_hub:
            self._push_from_checkpoint(output_dir)

        # Maybe delete some older checkpoints.
        if self.args.should_save:
            # Solely rely on numerical checkpoint id for rotation.
            # mtime is not reliable especially on some fuse fs in cloud environments.
            self._rotate_checkpoints(use_mtime=False, output_dir=run_dir)

    def _save_rng_state(self, output_dir):
        # Save RNG state in non-distributed training
        rng_states = {
            "python": random.getstate(),
            "numpy": np.random.get_state(),
            "cpu": torch.random.get_rng_state(),
        }
        if torch.cuda.is_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                # In non distributed, we save the global CUDA RNG state (will take care of DataParallel)
                rng_states["cuda"] = torch.cuda.random.get_rng_state_all()
            else:
                rng_states["cuda"] = torch.cuda.random.get_rng_state()

        if is_torch_xla_available():
            rng_states["xla"] = xm.get_rng_state()

        if is_torch_npu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["npu"] = torch.npu.random.get_rng_state_all()
            else:
                rng_states["npu"] = torch.npu.random.get_rng_state()

        if is_torch_mlu_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["mlu"] = torch.mlu.random.get_rng_state_all()
            else:
                rng_states["mlu"] = torch.mlu.random.get_rng_state()

        if is_torch_musa_available():
            if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
                rng_states["musa"] = torch.musa.get_rng_state_all()
            else:
                rng_states["musa"] = torch.musa.get_rng_state()

        # A process can arrive here before the process 0 has a chance to save the model, in which case output_dir may
        # not yet exist.
        os.makedirs(output_dir, exist_ok=True)

        if self.args.world_size <= 1:
            torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
        else:
            torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth"))

    def _save_optimizer_and_scheduler(self, output_dir):
        if is_torch_xla_available():
            xm.rendezvous("saving_optimizer_states")
            if self.is_fsdp_xla_v1_enabled:
                optm = {
                    "optimizer": self.optimizer.state_dict(),
                    "shard_metadata": self.model.get_shard_metadata(),
                }
                xm.save(
                    optm,
                    os.path.join(
                        output_dir, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}"
                    ),
                    master_only=False,
                )
            else:
                xm.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
            with warnings.catch_warnings(record=True) as caught_warnings:
                xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
                reissue_pt_warnings(caught_warnings)
        elif is_sagemaker_mp_enabled():
            opt_state_dict = self.optimizer.local_state_dict(gather_if_shard=False)
            smp.barrier()
            if smp.rdp_rank() == 0 or smp.state.cfg.shard_optimizer_state:
                smp.save(
                    opt_state_dict,
                    os.path.join(output_dir, OPTIMIZER_NAME),
                    partial=True,
                    v3=smp.state.cfg.shard_optimizer_state,
                )
        elif self.is_deepspeed_enabled:
            # under zero3 model file itself doesn't get saved since it's bogus! Unless deepspeed
            # config `stage3_gather_16bit_weights_on_model_save` is True
            accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set(
                inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys()
            )
            if accept_exclude_frozen_parameters and _is_peft_model(self.model):
                self.model_wrapped.save_checkpoint(output_dir, exclude_frozen_parameters=True)
            else:
                self.model_wrapped.save_checkpoint(output_dir)
        elif self.is_fsdp_enabled:
            # save fsdp specific ckpt for resuming from ckpt
            save_fsdp_model(
                self.accelerator.state.fsdp_plugin, self.accelerator, self.model, output_dir, **_get_fsdp_ckpt_kwargs()
            )
            save_fsdp_optimizer(
                self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir
            )
        elif self.args.should_save:
            # deepspeed.save_checkpoint above saves model/optim/sched
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))

        # Save SCHEDULER & SCALER
        is_deepspeed_custom_scheduler = self.is_deepspeed_enabled and not isinstance(
            self.lr_scheduler, DeepSpeedSchedulerWrapper
        )
        if (
            self.args.should_save
            and (not self.is_deepspeed_enabled or is_deepspeed_custom_scheduler)
            and not is_torch_xla_available()
        ):
            with warnings.catch_warnings(record=True) as caught_warnings:
                torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
            reissue_pt_warnings(caught_warnings)

    def _load_optimizer_and_scheduler(self, checkpoint):
        """If optimizer and scheduler states exist, load them."""
        if checkpoint is None:
            return

        if self.is_deepspeed_enabled:
            # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init
            if not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper):
                with warnings.catch_warnings(record=True) as caught_warnings:
                    self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, SCHEDULER_NAME)))
                reissue_pt_warnings(caught_warnings)
            return

        checkpoint_file_exists = (
            glob.glob(os.path.join(checkpoint, OPTIMIZER_NAME) + "_*")
            if is_sagemaker_mp_enabled()
            else (
                os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME))
                or os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME_BIN))
                or (
                    os.path.isdir(checkpoint)
                    and any(
                        OPTIMIZER_NAME_BIN.split(".")[0] in folder_name
                        for folder_name in os.listdir(checkpoint)
                        if os.path.isdir(os.path.join(checkpoint, folder_name))
                    )
                )
            )
        )
        checkpoint_file_exists = (
            glob.glob(os.path.join(checkpoint, f"rank*-of-{self.args.world_size}-{OPTIMIZER_NAME}"))
            if self.is_fsdp_xla_v1_enabled
            else checkpoint_file_exists
        )
        if checkpoint_file_exists and os.path.isfile(os.path.join(checkpoint, SCHEDULER_NAME)):
            # Load in optimizer and scheduler states
            if is_torch_xla_available():
                # On TPU we have to take some extra precautions to properly load the states on the right device.
                if self.is_fsdp_xla_v1_enabled:
                    optimizer_state = torch.load(
                        os.path.join(
                            checkpoint, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}"
                        ),
                        map_location="cpu",
                    )
                    # We only need `optimizer` when resuming from checkpoint
                    optimizer_state = optimizer_state["optimizer"]
                else:
                    optimizer_state = torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location="cpu")
                with warnings.catch_warnings(record=True) as caught_warnings:
                    lr_scheduler_state = torch.load(os.path.join(checkpoint, SCHEDULER_NAME), map_location="cpu")
                reissue_pt_warnings(caught_warnings)

                xm.send_cpu_data_to_device(optimizer_state, self.args.device)
                xm.send_cpu_data_to_device(lr_scheduler_state, self.args.device)

                self.optimizer.load_state_dict(optimizer_state)
                self.lr_scheduler.load_state_dict(lr_scheduler_state)
            else:
                if is_sagemaker_mp_enabled():
                    if os.path.isfile(os.path.join(checkpoint, "user_content.pt")):
                        # Optimizer checkpoint was saved with smp >= 1.10
                        def opt_load_hook(mod, opt):
                            opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True))

                    else:
                        # Optimizer checkpoint was saved with smp < 1.10
                        def opt_load_hook(mod, opt):
                            if IS_SAGEMAKER_MP_POST_1_10:
                                opt.load_state_dict(
                                    smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True, back_compat=True)
                                )
                            else:
                                opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True))

                    self.model_wrapped.register_post_step_hook(opt_load_hook)
                else:
                    # We use the CPU when training on one GPU to avoid OOM for GPU RAM when training big models.
                    # In distributed training however, we load directly on each GPU and risk the GPU OOM as it's more
                    # likely to get OOM on CPU (since we load num_gpu times the optimizer state
                    map_location = self.args.device if self.args.world_size > 1 else "cpu"
                    if self.is_fsdp_enabled:
                        load_fsdp_optimizer(
                            self.accelerator.state.fsdp_plugin,
                            self.accelerator,
                            self.optimizer,
                            self.model,
                            checkpoint,
                            **_get_fsdp_ckpt_kwargs(),
                        )
                    else:
                        self.optimizer.load_state_dict(
                            torch.load(os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location)
                        )
                with warnings.catch_warnings(record=True) as caught_warnings:
                    self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, SCHEDULER_NAME)))
                reissue_pt_warnings(caught_warnings)

    def _load_callback_state(self):
        """If callback states exist and were passed in, restore their states if enabled"""
        if not self.args.restore_callback_states_from_checkpoint:
            return
        # Callback states are stored in stateful_callbacks
        not_found = []
        new_callbacks = []
        original_callbacks = self.callback_handler.callbacks + [self.control]
        for stored_callback, data in self.state.stateful_callbacks.items():
            if not isinstance(data, list):
                data = [data]
            if any(callback.__class__.__name__ == stored_callback for callback in original_callbacks):
                # We can load/restore from multiple callbacks of the same type.
                duplicates = [
                    callback for callback in original_callbacks if callback.__class__.__name__ == stored_callback
                ]
                for callback, callback_data in zip(duplicates, data):
                    args = callback_data.get("args", {})
                    attributes = callback_data.get("attributes", {})
                    new_callback = type(callback)(**args)
                    for attribute, value in attributes.items():
                        setattr(new_callback, attribute, value)
                    if isinstance(callback, TrainerControl):
                        # Specifically for restoring the `control` state
                        self.control = new_callback
                    else:
                        new_callbacks.append(new_callback)
                    # We remove the existing callback and add it to the list of new callbacks
                    self.callback_handler.remove_callback(type(new_callback))
                logger.info("Continuing training from checkpoint, restoring any callbacks that were passed in")
            else:
                not_found.append(stored_callback)
        if len(not_found) > 0:
            logger.warning(
                f"Checkpoint included callbacks not included in current configuration. Ignoring. ({', '.join(not_found)})"
            )
        for callback in new_callbacks:
            self.callback_handler.add_callback(callback)

    def hyperparameter_search(
        self,
        hp_space: Optional[Callable[["optuna.Trial"], Dict[str, float]]] = None,
        compute_objective: Optional[Callable[[Dict[str, float]], float]] = None,
        n_trials: int = 20,
        direction: Union[str, List[str]] = "minimize",
        backend: Optional[Union["str", HPSearchBackend]] = None,
        hp_name: Optional[Callable[["optuna.Trial"], str]] = None,
        **kwargs,
    ) -> Union[BestRun, List[BestRun]]:
        """
        Launch an hyperparameter search using `optuna` or `Ray Tune` or `SigOpt`. The optimized quantity is determined
        by `compute_objective`, which defaults to a function returning the evaluation loss when no metric is provided,
        the sum of all metrics otherwise.

        <Tip warning={true}>

        To use this method, you need to have provided a `model_init` when initializing your [`Trainer`]: we need to
        reinitialize the model at each new run. This is incompatible with the `optimizers` argument, so you need to
        subclass [`Trainer`] and override the method [`~Trainer.create_optimizer_and_scheduler`] for custom
        optimizer/scheduler.

        </Tip>

        Args:
            hp_space (`Callable[["optuna.Trial"], Dict[str, float]]`, *optional*):
                A function that defines the hyperparameter search space. Will default to
                [`~trainer_utils.default_hp_space_optuna`] or [`~trainer_utils.default_hp_space_ray`] or
                [`~trainer_utils.default_hp_space_sigopt`] depending on your backend.
            compute_objective (`Callable[[Dict[str, float]], float]`, *optional*):
                A function computing the objective to minimize or maximize from the metrics returned by the `evaluate`
                method. Will default to [`~trainer_utils.default_compute_objective`].
            n_trials (`int`, *optional*, defaults to 100):
                The number of trial runs to test.
            direction (`str` or `List[str]`, *optional*, defaults to `"minimize"`):
                If it's single objective optimization, direction is `str`, can be `"minimize"` or `"maximize"`, you
                should pick `"minimize"` when optimizing the validation loss, `"maximize"` when optimizing one or
                several metrics. If it's multi objectives optimization, direction is `List[str]`, can be List of
                `"minimize"` and `"maximize"`, you should pick `"minimize"` when optimizing the validation loss,
                `"maximize"` when optimizing one or several metrics.
            backend (`str` or [`~training_utils.HPSearchBackend`], *optional*):
                The backend to use for hyperparameter search. Will default to optuna or Ray Tune or SigOpt, depending
                on which one is installed. If all are installed, will default to optuna.
            hp_name (`Callable[["optuna.Trial"], str]]`, *optional*):
                A function that defines the trial/run name. Will default to None.
            kwargs (`Dict[str, Any]`, *optional*):
                Additional keyword arguments for each backend:

                - `optuna`: parameters from
                  [optuna.study.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html)
                  and also the parameters `timeout`, `n_jobs` and `gc_after_trial` from
                  [optuna.study.Study.optimize](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize)
                - `ray`: parameters from [tune.run](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run).
                  If `resources_per_trial` is not set in the `kwargs`, it defaults to 1 CPU core and 1 GPU (if available).
                  If `progress_reporter` is not set in the `kwargs`,
                  [ray.tune.CLIReporter](https://docs.ray.io/en/latest/tune/api/doc/ray.tune.CLIReporter.html) is used.
                - `sigopt`: the parameter `proxies` from
                  [sigopt.Connection.set_proxies](https://docs.sigopt.com/support/faq#how-do-i-use-sigopt-with-a-proxy).

        Returns:
            [`trainer_utils.BestRun` or `List[trainer_utils.BestRun]`]: All the information about the best run or best
            runs for multi-objective optimization. Experiment summary can be found in `run_summary` attribute for Ray
            backend.
        """
        if backend is None:
            backend = default_hp_search_backend()
        backend = HPSearchBackend(backend)
        backend_obj = ALL_HYPERPARAMETER_SEARCH_BACKENDS[backend]()
        backend_obj.ensure_available()
        self.hp_search_backend = backend
        if self.model_init is None:
            raise RuntimeError(
                "To use hyperparameter search, you need to pass your model through a model_init function."
            )

        self.hp_space = backend_obj.default_hp_space if hp_space is None else hp_space
        self.hp_name = hp_name
        self.compute_objective = default_compute_objective if compute_objective is None else compute_objective

        best_run = backend_obj.run(self, n_trials, direction, **kwargs)

        self.hp_search_backend = None
        return best_run

    def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None:
        """
        Log `logs` on the various objects watching training.

        Subclass and override this method to inject custom behavior.

        Args:
            logs (`Dict[str, float]`):
                The values to log.
            start_time (`Optional[float]`):
                The start of training.
        """
        if self.state.epoch is not None:
            logs["epoch"] = self.state.epoch
        if self.args.include_num_input_tokens_seen:
            logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen
            if start_time is not None:
                speed_metrics("train", start_time, num_tokens=self.state.num_input_tokens_seen)

        output = {**logs, **{"step": self.state.global_step}}
        self.state.log_history.append(output)
        self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)

    def _prepare_input(self, data: Union[torch.Tensor, Any]) -> Union[torch.Tensor, Any]:
        """
        Prepares one `data` before feeding it to the model, be it a tensor or a nested list/dictionary of tensors.
        """
        if isinstance(data, Mapping):
            return type(data)({k: self._prepare_input(v) for k, v in data.items()})
        elif isinstance(data, (tuple, list)):
            return type(data)(self._prepare_input(v) for v in data)
        elif isinstance(data, torch.Tensor):
            kwargs = {"device": self.args.device}
            if self.is_deepspeed_enabled and (torch.is_floating_point(data) or torch.is_complex(data)):
                # NLP models inputs are int/uint and those get adjusted to the right dtype of the
                # embedding. Other models such as wav2vec2's inputs are already float and thus
                # may need special handling to match the dtypes of the model
                kwargs.update({"dtype": self.accelerator.state.deepspeed_plugin.hf_ds_config.dtype()})
            return data.to(**kwargs)
        return data

    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
        """
        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
        handling potential state.
        """
        inputs = self._prepare_input(inputs)
        if len(inputs) == 0:
            raise ValueError(
                "The batch received was empty, your model won't be able to train on it. Double-check that your "
                f"training dataset contains keys expected by the model: {','.join(self._signature_columns)}."
            )
        if self.args.past_index >= 0 and self._past is not None:
            inputs["mems"] = self._past

        return inputs

    def compute_loss_context_manager(self):
        """
        A helper wrapper to group together context managers.
        """
        return self.autocast_smart_context_manager()

    def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True):
        """
        A helper wrapper that creates an appropriate context manager for `autocast` while feeding it the desired
        arguments, depending on the situation.
        """
        if self.use_cpu_amp:
            ctx_manager = torch.cpu.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype)
        else:
            ctx_manager = contextlib.nullcontext()

        return ctx_manager

    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None
    ) -> torch.Tensor:
        """
        Perform a training step on a batch of inputs.

        Subclass and override to inject custom behavior.

        Args:
            model (`nn.Module`):
                The model to train.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.

        Return:
            `torch.Tensor`: The tensor with training loss on this batch.
        """
        model.train()
        if hasattr(self.optimizer, "train") and callable(self.optimizer.train):
            self.optimizer.train()

        inputs = self._prepare_inputs(inputs)
        if is_sagemaker_mp_enabled():
            loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
            return loss_mb.reduce_mean().detach().to(self.args.device)

        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)

        del inputs
        if (
            self.args.torch_empty_cache_steps is not None
            and self.state.global_step % self.args.torch_empty_cache_steps == 0
        ):
            if is_torch_xpu_available():
                torch.xpu.empty_cache()
            elif is_torch_mlu_available():
                torch.mlu.empty_cache()
            elif is_torch_musa_available():
                torch.musa.empty_cache()
            elif is_torch_npu_available():
                torch.npu.empty_cache()
            elif is_torch_mps_available(min_version="2.0"):
                torch.mps.empty_cache()
            else:
                torch.cuda.empty_cache()

        kwargs = {}

        # For LOMO optimizers you need to explicitly use the learnign rate
        if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]:
            kwargs["learning_rate"] = self._get_learning_rate()

        if self.args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training

        if self.use_apex:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            # Finally we need to normalize the loss for reporting
            if not self.model_accepts_loss_kwargs and self.compute_loss_func is None:
                loss = loss / self.args.gradient_accumulation_steps

            self.accelerator.backward(loss, **kwargs)

            return loss.detach()

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """
        if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
        if self.model_accepts_loss_kwargs:
            loss_kwargs = {}
            if num_items_in_batch is not None:
                loss_kwargs["num_items_in_batch"] = num_items_in_batch
            inputs = {**inputs, **loss_kwargs}
        outputs = model(**inputs)
        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        if labels is not None:
            unwrapped_model = self.accelerator.unwrap_model(model)
            if _is_peft_model(unwrapped_model):
                model_name = unwrapped_model.base_model.model._get_name()
            else:
                model_name = unwrapped_model._get_name()
            # User-defined compute_loss function
            if self.compute_loss_func is not None:
                loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch)
            elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                loss = self.label_smoother(outputs, labels, shift_labels=True)
            else:
                loss = self.label_smoother(outputs, labels)
        else:
            if isinstance(outputs, dict) and "loss" not in outputs:
                raise ValueError(
                    "The model did not return a loss from the inputs, only the following keys: "
                    f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
                )
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        if self.args.average_tokens_across_devices and self.model_accepts_loss_kwargs:
            loss *= self.accelerator.num_processes

        return (loss, outputs) if return_outputs else loss

    def is_local_process_zero(self) -> bool:
        """
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
        machines) main process.
        """
        return self.args.local_process_index == 0

    def is_world_process_zero(self) -> bool:
        """
        Whether or not this process is the global main process (when training in a distributed fashion on several
        machines, this is only going to be `True` for one process).
        """
        # Special case for SageMaker ModelParallel since there process_index is dp_process_index, not the global
        # process index.
        if is_sagemaker_mp_enabled():
            return smp.rank() == 0
        else:
            return self.args.process_index == 0

    def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
        """
        Will save the model, so you can reload it using `from_pretrained()`.

        Will only save from the main process.
        """

        if output_dir is None:
            output_dir = self.args.output_dir

        if is_torch_xla_available():
            self._save_tpu(output_dir)
        elif is_sagemaker_mp_enabled():
            # Calling the state_dict needs to be done on the wrapped model and on all processes.
            os.makedirs(output_dir, exist_ok=True)
            state_dict = self.model_wrapped.state_dict()
            if self.args.should_save:
                self._save(output_dir, state_dict=state_dict)
            if IS_SAGEMAKER_MP_POST_1_10:
                # 'user_content.pt' indicates model state_dict saved with smp >= 1.10
                Path(os.path.join(output_dir, "user_content.pt")).touch()
        elif self.is_fsdp_enabled:
            if ("FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)) and (
                version.parse(accelerate_version) > version.parse("0.24.1")
            ):
                state_dict = self.accelerator.get_state_dict(self.model)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)
        elif self.is_deepspeed_enabled:
            try:
                state_dict = self.accelerator.get_state_dict(self.deepspeed)
                if self.args.should_save:
                    self._save(output_dir, state_dict=state_dict)
            except ValueError:
                logger.warning(
                    " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use"
                    " zero_to_fp32.py to recover weights"
                )
                if self.args.should_save:
                    self._save(output_dir, state_dict={})
                # remove the dummy state_dict
                remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
                self.model_wrapped.save_checkpoint(output_dir)

        elif self.args.should_save:
            self._save(output_dir)

        # Push to the Hub when `save_model` is called by the user.
        if self.args.push_to_hub and not _internal_call:
            self.push_to_hub(commit_message="Model save")

    def _save_tpu(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else self.args.output_dir

        logger.info(f"Saving model checkpoint to {output_dir}")
        model = self.model
        xm.mark_step()

        if xm.is_master_ordinal(local=False):
            os.makedirs(output_dir, exist_ok=True)
            torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        supported_classes = (PushToHubMixin,)
        xm.rendezvous("saving_checkpoint")
        if self.is_fsdp_xla_v1_enabled:
            ckpt = {
                "model": model.state_dict(),
                "shard_metadata": model.get_shard_metadata(),
            }
            ckpt_path = os.path.join(
                output_dir, f"rank{self.args.process_index}-of-{self.args.world_size}-{WEIGHTS_NAME}"
            )
            # All ranks save sharded checkpoint
            xm.save(ckpt, ckpt_path, master_only=False)
            # Make sure all ranks have saved checkpoints
            xm.rendezvous("save_full_checkpoints")
            # Master save full checkpoint
            if self.args.should_save:
                from torch_xla.distributed.fsdp import consolidate_sharded_model_checkpoints

                full_state_dict, _ = consolidate_sharded_model_checkpoints(
                    ckpt_prefix=os.path.join(output_dir, ""),
                    ckpt_suffix=f"rank*-of-*-{WEIGHTS_NAME}",
                    save_model=False,
                )
                model = model.module.module
                unwrapped_model = self.accelerator.unwrap_model(model)
                if isinstance(unwrapped_model, supported_classes):
                    unwrapped_model.save_pretrained(
                        output_dir,
                        state_dict=full_state_dict,
                        save_function=xm.save,
                        safe_serialization=self.args.save_safetensors,
                    )
                else:
                    logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                    xm.save(full_state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        elif not isinstance(model, supported_classes):
            if isinstance(self.accelerator.unwrap_model(model), supported_classes):
                self.accelerator.unwrap_model(model).save_pretrained(
                    output_dir,
                    is_main_process=self.args.should_save,
                    state_dict=xm._maybe_convert_to_cpu(model.state_dict()),
                    save_function=xm.save,
                    safe_serialization=self.args.save_safetensors,
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                state_dict = xm._maybe_convert_to_cpu(model.state_dict())
                xm.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            model.save_pretrained(
                output_dir,
                is_main_process=self.args.should_save,
                save_function=xm.save,
                safe_serialization=self.args.save_safetensors,
                state_dict=xm._maybe_convert_to_cpu(model.state_dict()),
            )
        if self.processing_class is not None and self.args.should_save:
            self.processing_class.save_pretrained(output_dir)

    def _save(self, output_dir: Optional[str] = None, state_dict=None):
        # If we are executing this function, we are the process zero, so we don't check for that.
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")

        supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel)
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        if not isinstance(self.model, supported_classes):
            if state_dict is None:
                state_dict = self.model.state_dict()

            if isinstance(self.accelerator.unwrap_model(self.model), supported_classes):
                self.accelerator.unwrap_model(self.model).save_pretrained(
                    output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
                )
            else:
                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
                if self.args.save_safetensors:
                    safetensors.torch.save_file(
                        state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"}
                    )
                else:
                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
        else:
            self.model.save_pretrained(
                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
            )

        if self.processing_class is not None:
            self.processing_class.save_pretrained(output_dir)

        # Good practice: save your training arguments together with the trained model
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

    def store_flos(self):
        # Storing the number of floating-point operations that went into the model
        if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            self.state.total_flos += (
                distributed_broadcast_scalars([self.current_flos], device=self.args.device).sum().item()
            )
            self.current_flos = 0
        else:
            self.state.total_flos += self.current_flos
            self.current_flos = 0

    def _sorted_checkpoints(
        self, output_dir=None, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False
    ) -> List[str]:
        ordering_and_checkpoint_path = []

        glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)]

        for path in glob_checkpoints:
            if use_mtime:
                ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
            else:
                regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path)
                if regex_match is not None and regex_match.groups() is not None:
                    ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))

        checkpoints_sorted = sorted(ordering_and_checkpoint_path)
        checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
        # Make sure we don't delete the best model.
        if (
            self.state.best_model_checkpoint is not None
            and str(Path(self.state.best_model_checkpoint)) in checkpoints_sorted
        ):
            best_model_index = checkpoints_sorted.index(str(Path(self.state.best_model_checkpoint)))
            for i in range(best_model_index, len(checkpoints_sorted) - 2):
                checkpoints_sorted[i], checkpoints_sorted[i + 1] = checkpoints_sorted[i + 1], checkpoints_sorted[i]
        return checkpoints_sorted

    def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
        if self.args.save_total_limit is None or self.args.save_total_limit <= 0:
            return

        # Check if we should delete older checkpoint(s)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir)
        if len(checkpoints_sorted) <= self.args.save_total_limit:
            return

        # If save_total_limit=1 with load_best_model_at_end=True, we could end up deleting the last checkpoint, which
        # we don't do to allow resuming.
        save_total_limit = self.args.save_total_limit
        if (
            self.state.best_model_checkpoint is not None
            and self.args.save_total_limit == 1
            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
        ):
            save_total_limit = 2

        number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
        checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
        for checkpoint in checkpoints_to_be_deleted:
            logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
            shutil.rmtree(checkpoint, ignore_errors=True)

    def evaluate(
        self,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> Dict[str, float]:
        """
        Run evaluation and returns metrics.

        The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
        (pass it to the init `compute_metrics` argument).

        You can also subclass and override this method to inject custom behavior.

        Args:
            eval_dataset (Union[`Dataset`, Dict[str, `Dataset`]), *optional*):
                Pass a dataset if you wish to override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns
                not accepted by the `model.forward()` method are automatically removed. If it is a dictionary, it will
                evaluate on each dataset, prepending the dictionary key to the metric name. Datasets must implement the
                `__len__` method.

                <Tip>

                If you pass a dictionary with names of datasets as keys and datasets as values, evaluate will run
                separate evaluations on each dataset. This can be useful to monitor how training affects other
                datasets or simply to get a more fine-grained evaluation.
                When used with `load_best_model_at_end`, make sure `metric_for_best_model` references exactly one
                of the datasets. If you, for example, pass in `{"data1": data1, "data2": data2}` for two datasets
                `data1` and `data2`, you could specify `metric_for_best_model="eval_data1_loss"` for using the
                loss on `data1` and `metric_for_best_model="eval_data2_loss"` for the loss on `data2`.

                </Tip>

            ignore_keys (`List[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
            metric_key_prefix (`str`, *optional*, defaults to `"eval"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "eval_bleu" if the prefix is "eval" (default)

        Returns:
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
            dictionary also contains the epoch number which comes from the training state.
        """
        # handle multipe eval datasets
        override = eval_dataset is not None
        eval_dataset = eval_dataset if override else self.eval_dataset
        if isinstance(eval_dataset, dict):
            metrics = {}
            for eval_dataset_name, _eval_dataset in eval_dataset.items():
                dataset_metrics = self.evaluate(
                    eval_dataset=_eval_dataset if override else eval_dataset_name,
                    ignore_keys=ignore_keys,
                    metric_key_prefix=f"{metric_key_prefix}_{eval_dataset_name}",
                )
                metrics.update(dataset_metrics)
            return metrics

        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        if self.is_fsdp_xla_v2_enabled:
            eval_dataloader = tpu_spmd_dataloader(eval_dataloader)

        start_time = time.time()

        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
            eval_dataloader,
            description="Evaluation",
            # No point gathering the predictions if there are no metrics, otherwise we defer to
            # self.args.prediction_loss_only
            prediction_loss_only=True if self.compute_metrics is None else None,
            ignore_keys=ignore_keys,
            metric_key_prefix=metric_key_prefix,
        )

        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
        if f"{metric_key_prefix}_model_preparation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        self.log(output.metrics)

        if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
            # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
            xm.master_print(met.metrics_report())

        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)

        self._memory_tracker.stop_and_update_metrics(output.metrics)

        return output.metrics

    def predict(
        self, test_dataset: Dataset, ignore_keys: Optional[List[str]] = None, metric_key_prefix: str = "test"
    ) -> PredictionOutput:
        """
        Run prediction and returns predictions and potential metrics.

        Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method
        will also return metrics, like in `evaluate()`.

        Args:
            test_dataset (`Dataset`):
                Dataset to run the predictions on. If it is an `datasets.Dataset`, columns not accepted by the
                `model.forward()` method are automatically removed. Has to implement the method `__len__`
            ignore_keys (`List[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
            metric_key_prefix (`str`, *optional*, defaults to `"test"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "test_bleu" if the prefix is "test" (default)

        <Tip>

        If your predictions or labels have different sequence length (for instance because you're doing dynamic padding
        in a token classification task) the predictions will be padded (on the right) to allow for concatenation into
        one array. The padding index is -100.

        </Tip>

        Returns: *NamedTuple* A namedtuple with the following keys:

            - predictions (`np.ndarray`): The predictions on `test_dataset`.
            - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some).
            - metrics (`Dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained
              labels).
        """
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()

        test_dataloader = self.get_test_dataloader(test_dataset)
        start_time = time.time()

        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        output = eval_loop(
            test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix
        )
        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
        if f"{metric_key_prefix}_model_preparation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        self.control = self.callback_handler.on_predict(self.args, self.state, self.control, output.metrics)
        self._memory_tracker.stop_and_update_metrics(output.metrics)

        return PredictionOutput(predictions=output.predictions, label_ids=output.label_ids, metrics=output.metrics)

    def evaluation_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> EvalLoopOutput:
        """
        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.

        Works both with or without labels.
        """
        args = self.args

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only

        # if eval is called w/o train, handle model prep here
        if self.is_deepspeed_enabled and self.deepspeed is None:
            _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

        model = self._wrap_model(self.model, training=False, dataloader=dataloader)

        if len(self.accelerator._models) == 0 and model is self.model:
            start_time = time.time()
            model = (
                self.accelerator.prepare(model)
                if self.is_deepspeed_enabled or (self.is_fsdp_enabled and self.accelerator.mixed_precision != "fp8")
                else self.accelerator.prepare_model(model, evaluation_mode=True)
            )
            self.model_preparation_time = round(time.time() - start_time, 4)

            if self.is_fsdp_enabled:
                self.model = model

            # for the rest of this function `model` is the outside model, whether it was wrapped or not
            if model is not self.model:
                self.model_wrapped = model

            # backward compatibility
            if self.is_deepspeed_enabled:
                self.deepspeed = self.model_wrapped

        # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called
        # while ``train`` is running, cast it to the right dtype first and then put on device
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = self.args.eval_batch_size

        logger.info(f"\n***** Running {description} *****")
        if has_length(dataloader):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
        logger.info(f"  Batch size = {batch_size}")

        model.eval()
        if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval):
            self.optimizer.eval()

        self.callback_handler.eval_dataloader = dataloader
        # Do this before wrapping.
        eval_dataset = getattr(dataloader, "dataset", None)

        if args.past_index >= 0:
            self._past = None

        # Initialize containers
        all_losses = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
        all_preds = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
        all_labels = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)
        all_inputs = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100)

        metrics = None
        eval_set_kwargs = {}

        # Will be useful when we have an iterable dataset so don't know its length.
        observed_num_examples = 0

        # Main evaluation loop
        for step, inputs in enumerate(dataloader):
            # Update the observed num examples
            observed_batch_size = find_batch_size(inputs)
            if observed_batch_size is not None:
                observed_num_examples += observed_batch_size
                # For batch samplers, batch_size is not known by the dataloader in advance.
                if batch_size is None:
                    batch_size = observed_batch_size

            # Prediction step
            losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
            main_input_name = getattr(self.model, "main_input_name", "input_ids")
            inputs_decode = (
                self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None
            )

            if is_torch_xla_available():
                xm.mark_step()

            # Update containers
            if losses is not None:
                losses = self.gather_function((losses.repeat(batch_size)))
                all_losses.add(losses)
            if inputs_decode is not None:
                inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100)
                inputs_decode = self.gather_function((inputs_decode))
                if not self.args.batch_eval_metrics or description == "Prediction":
                    all_inputs.add(inputs_decode)
            if labels is not None:
                # Pad labels here, preparing for preprocess_logits_for_metrics in next logits block.
                labels = self.accelerator.pad_across_processes(labels, dim=1, pad_index=-100)
            if logits is not None:
                logits = self.accelerator.pad_across_processes(logits, dim=1, pad_index=-100)
                if self.preprocess_logits_for_metrics is not None:
                    logits = self.preprocess_logits_for_metrics(logits, labels)
                logits = self.gather_function((logits))
                if not self.args.batch_eval_metrics or description == "Prediction":
                    all_preds.add(logits)
            if labels is not None:
                labels = self.gather_function((labels))
                if not self.args.batch_eval_metrics or description == "Prediction":
                    all_labels.add(labels)

            self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

            if self.args.batch_eval_metrics:
                if self.compute_metrics is not None and logits is not None and labels is not None:
                    is_last_step = self.accelerator.gradient_state.end_of_dataloader
                    batch_kwargs = {}
                    batch_kwargs["losses"] = losses if "loss" in args.include_for_metrics else None
                    batch_kwargs["inputs"] = inputs if "inputs" in args.include_for_metrics else None
                    metrics = self.compute_metrics(
                        EvalPrediction(predictions=logits, label_ids=labels, **batch_kwargs),
                        compute_result=is_last_step,
                    )

                del losses, logits, labels, inputs
                torch.cuda.empty_cache()

            # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
            elif args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0:
                all_losses.to_cpu_and_numpy()
                all_preds.to_cpu_and_numpy()
                all_labels.to_cpu_and_numpy()
                all_inputs.to_cpu_and_numpy()

                del losses, logits, labels, inputs
                torch.cuda.empty_cache()

        # After all calls to `.gather_function`, reset to `gather_for_metrics`:
        self.gather_function = self.accelerator.gather_for_metrics
        if args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of the evaluation loop
            delattr(self, "_past")

        # Gather all remaining tensors and put them back on the CPU
        all_losses = all_losses.get_arrays()
        all_preds = all_preds.get_arrays()
        all_labels = all_labels.get_arrays()
        all_inputs = all_inputs.get_arrays()

        # Number of samples
        if has_length(eval_dataset):
            num_samples = len(eval_dataset)
        # The instance check is weird and does not actually check for the type, but whether the dataset has the right
        # methods. Therefore we need to make sure it also has the attribute.
        elif isinstance(eval_dataset, IterableDatasetShard) and getattr(eval_dataset, "num_examples", 0) > 0:
            num_samples = eval_dataset.num_examples
        else:
            if has_length(dataloader):
                num_samples = self.num_examples(dataloader)
            else:  # both len(dataloader.dataset) and len(dataloader) fail
                num_samples = observed_num_examples
        if num_samples == 0 and observed_num_examples > 0:
            num_samples = observed_num_examples

        # Metrics!
        if (
            self.compute_metrics is not None
            and all_preds is not None
            and all_labels is not None
            and not self.args.batch_eval_metrics
        ):
            eval_set_kwargs["losses"] = all_losses if "loss" in args.include_for_metrics else None
            eval_set_kwargs["inputs"] = all_inputs if "inputs" in args.include_for_metrics else None
            metrics = self.compute_metrics(
                EvalPrediction(predictions=all_preds, label_ids=all_labels, **eval_set_kwargs)
            )
        elif metrics is None:
            metrics = {}

        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

        if isinstance(all_losses, list) and all_losses:
            metrics[f"{metric_key_prefix}_loss"] = np.concatenate(all_losses).mean().item()
        elif isinstance(all_losses, np.ndarray):
            metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()
        if hasattr(self, "jit_compilation_time"):
            metrics[f"{metric_key_prefix}_jit_compilation_time"] = self.jit_compilation_time
        if hasattr(self, "model_preparation_time"):
            metrics[f"{metric_key_prefix}_model_preparation_time"] = self.model_preparation_time

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples)

    def _nested_gather(self, tensors, name=None):
        """
        Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before
        concatenating them to `gathered`
        """
        if tensors is None:
            return
        if is_torch_xla_available():
            if name is None:
                name = "nested_gather"
            tensors = nested_xla_mesh_reduce(tensors, name)
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
        elif (self.args.distributed_state is not None and self.args.distributed_state.distributed_type != "NO") or (
            self.args.distributed_state is None and self.args.local_rank != -1
        ):
            tensors = distributed_concat(tensors)
        return tensors

    def prediction_step(
        self,
        model: nn.Module,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        """
        Perform an evaluation step on `model` using `inputs`.

        Subclass and override to inject custom behavior.

        Args:
            model (`nn.Module`):
                The model to evaluate.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.
            prediction_loss_only (`bool`):
                Whether or not to return the loss only.
            ignore_keys (`List[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.

        Return:
            Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
            logits and labels (each being optional).
        """
        has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names)
        # For CLIP-like models capable of returning loss values.
        # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss`
        # is `True` in `model.forward`.
        return_loss = inputs.get("return_loss", None)
        if return_loss is None:
            return_loss = self.can_return_loss
        loss_without_labels = True if len(self.label_names) == 0 and return_loss else False

        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []

        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
        if has_labels or loss_without_labels:
            labels = nested_detach(tuple(inputs.get(name) for name in self.label_names))
            if len(labels) == 1:
                labels = labels[0]
        else:
            labels = None

        with torch.no_grad():
            if is_sagemaker_mp_enabled():
                raw_outputs = smp_forward_only(model, inputs)
                if has_labels or loss_without_labels:
                    if isinstance(raw_outputs, dict):
                        loss_mb = raw_outputs["loss"]
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        loss_mb = raw_outputs[0]
                        logits_mb = raw_outputs[1:]

                    loss = loss_mb.reduce_mean().detach().cpu()
                    logits = smp_nested_concat(logits_mb)
                else:
                    loss = None
                    if isinstance(raw_outputs, dict):
                        logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys)
                    else:
                        logits_mb = raw_outputs
                    logits = smp_nested_concat(logits_mb)
            else:
                if has_labels or loss_without_labels:
                    with self.compute_loss_context_manager():
                        loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
                    loss = loss.mean().detach()

                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
                    else:
                        logits = outputs[1:]
                else:
                    loss = None
                    with self.compute_loss_context_manager():
                        outputs = model(**inputs)
                    if isinstance(outputs, dict):
                        logits = tuple(v for k, v in outputs.items() if k not in ignore_keys)
                    else:
                        logits = outputs
                    # TODO: this needs to be fixed and made cleaner later.
                    if self.args.past_index >= 0:
                        self._past = outputs[self.args.past_index - 1]

        if prediction_loss_only:
            return (loss, None, None)

        logits = nested_detach(logits)
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)

    def floating_point_ops(self, inputs: Dict[str, Union[torch.Tensor, Any]]):
        """
        For models that inherit from [`PreTrainedModel`], uses that method to compute the number of floating point
        operations for every backward + forward pass. If using another model, either implement such a method in the
        model or subclass and override this method.

        Args:
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

        Returns:
            `int`: The number of floating-point operations.
        """
        if hasattr(self.model, "floating_point_ops"):
            return self.model.floating_point_ops(inputs)
        else:
            return 0

    def init_hf_repo(self, token: Optional[str] = None):
        """
        Initializes a git repo in `self.args.hub_model_id`.
        """
        # Only on process zero
        if not self.is_world_process_zero():
            return

        if self.args.hub_model_id is None:
            repo_name = Path(self.args.output_dir).absolute().name
        else:
            repo_name = self.args.hub_model_id

        token = token if token is not None else self.args.hub_token
        repo_url = create_repo(repo_name, token=token, private=self.args.hub_private_repo, exist_ok=True)
        self.hub_model_id = repo_url.repo_id
        self.push_in_progress = None

    def create_model_card(
        self,
        language: Optional[str] = None,
        license: Optional[str] = None,
        tags: Union[str, List[str], None] = None,
        model_name: Optional[str] = None,
        finetuned_from: Optional[str] = None,
        tasks: Union[str, List[str], None] = None,
        dataset_tags: Union[str, List[str], None] = None,
        dataset: Union[str, List[str], None] = None,
        dataset_args: Union[str, List[str], None] = None,
    ):
        """
        Creates a draft of a model card using the information available to the `Trainer`.

        Args:
            language (`str`, *optional*):
                The language of the model (if applicable)
            license (`str`, *optional*):
                The license of the model. Will default to the license of the pretrained model used, if the original
                model given to the `Trainer` comes from a repo on the Hub.
            tags (`str` or `List[str]`, *optional*):
                Some tags to be included in the metadata of the model card.
            model_name (`str`, *optional*):
                The name of the model.
            finetuned_from (`str`, *optional*):
                The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo
                of the original model given to the `Trainer` (if it comes from the Hub).
            tasks (`str` or `List[str]`, *optional*):
                One or several task identifiers, to be included in the metadata of the model card.
            dataset_tags (`str` or `List[str]`, *optional*):
                One or several dataset tags, to be included in the metadata of the model card.
            dataset (`str` or `List[str]`, *optional*):
                One or several dataset identifiers, to be included in the metadata of the model card.
            dataset_args (`str` or `List[str]`, *optional*):
               One or several dataset arguments, to be included in the metadata of the model card.
        """
        if not self.is_world_process_zero():
            return

        model_card_filepath = os.path.join(self.args.output_dir, "README.md")
        is_peft_library = False
        if os.path.exists(model_card_filepath):
            library_name = ModelCard.load(model_card_filepath).data.get("library_name")
            is_peft_library = library_name == "peft"

            # Append existing tags in `tags`
            existing_tags = ModelCard.load(model_card_filepath).data.tags
            if tags is not None and existing_tags is not None:
                if isinstance(tags, str):
                    tags = [tags]
                for tag in existing_tags:
                    if tag not in tags:
                        tags.append(tag)

        training_summary = TrainingSummary.from_trainer(
            self,
            language=language,
            license=license,
            tags=tags,
            model_name=model_name,
            finetuned_from=finetuned_from,
            tasks=tasks,
            dataset_tags=dataset_tags,
            dataset=dataset,
            dataset_args=dataset_args,
        )
        model_card = training_summary.to_model_card()
        with open(model_card_filepath, "w") as f:
            f.write(model_card)

        if is_peft_library:
            self.accelerator.unwrap_model(self.model).create_or_update_model_card(self.args.output_dir)

    def _push_from_checkpoint(self, checkpoint_folder):
        # Only push from one node.
        if not self.is_world_process_zero() or self.args.hub_strategy == HubStrategy.END:
            return
        # If we haven't finished the last push, we don't do this one unless args.hub_always_push=True.
        if not self.args.hub_always_push and self.push_in_progress is not None and not self.push_in_progress.is_done():
            return

        output_dir = self.args.output_dir
        # To avoid a new synchronization of all model weights, we just copy the file from the checkpoint folder
        modeling_files = [CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME]
        #  Add sharded checkpoints if we have an index
        for index_file in [WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]:
            index_path = os.path.join(checkpoint_folder, index_file)
            if os.path.isfile(index_path):
                modeling_files.append(index_file)
                with open(index_path) as f:
                    index = json.loads(f.read())
                shard_files = list(set(index["weight_map"].values()))
                modeling_files.extend(shard_files)
        if is_peft_available():
            modeling_files.extend([ADAPTER_CONFIG_NAME, ADAPTER_WEIGHTS_NAME, ADAPTER_SAFE_WEIGHTS_NAME])
        for modeling_file in modeling_files:
            if os.path.isfile(os.path.join(checkpoint_folder, modeling_file)):
                shutil.copy(os.path.join(checkpoint_folder, modeling_file), os.path.join(output_dir, modeling_file))
        # Saving the processing class is fast and we don't know how many files it may have spawned, so we resave it to be sure.
        if self.processing_class is not None:
            self.processing_class.save_pretrained(output_dir)
        # Same for the training arguments
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

        if self.args.save_strategy == SaveStrategy.STEPS:
            commit_message = f"Training in progress, step {self.state.global_step}"
        else:
            commit_message = f"Training in progress, epoch {int(self.state.epoch)}"

        model_push_job = upload_folder(
            repo_id=self.hub_model_id,
            folder_path=output_dir,
            commit_message=commit_message,
            token=self.args.hub_token,
            run_as_future=True,
            ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
        )

        push_jobs = [model_push_job]

        if self.args.hub_strategy in [HubStrategy.CHECKPOINT, HubStrategy.ALL_CHECKPOINTS]:
            path_in_repo = (
                "last-checkpoint" if self.args.hub_strategy == HubStrategy.CHECKPOINT else Path(checkpoint_folder).name
            )
            checkpoint_push = upload_folder(
                repo_id=self.hub_model_id,
                folder_path=checkpoint_folder,
                path_in_repo=path_in_repo,
                commit_message=commit_message + ", checkpoint",
                token=self.args.hub_token,
                run_as_future=True,
            )
            push_jobs.append(checkpoint_push)

        if self.push_in_progress is None or self.push_in_progress.is_done():
            self.push_in_progress = PushInProgress(push_jobs)
        else:
            self.push_in_progress.jobs.extend(push_jobs)

    def _finish_current_push(self):
        if not hasattr(self, "push_in_progress"):
            return
        if self.push_in_progress is not None and not self.push_in_progress.is_done():
            logger.info("Waiting for the current checkpoint push to be finished, this might take a couple of minutes.")
            self.push_in_progress.wait_until_done()

    def push_to_hub(
        self,
        commit_message: Optional[str] = "End of training",
        blocking: bool = True,
        token: Optional[str] = None,
        revision: Optional[str] = None,
        **kwargs,
    ) -> str:
        """
        Upload `self.model` and `self.processing_class` to the 🤗 model hub on the repo `self.args.hub_model_id`.

        Parameters:
            commit_message (`str`, *optional*, defaults to `"End of training"`):
                Message to commit while pushing.
            blocking (`bool`, *optional*, defaults to `True`):
                Whether the function should return only when the `git push` has finished.
            token (`str`, *optional*, defaults to `None`):
                Token with write permission to overwrite Trainer's original args.
            revision (`str`, *optional*):
                The git revision to commit from. Defaults to the head of the "main" branch.
            kwargs (`Dict[str, Any]`, *optional*):
                Additional keyword arguments passed along to [`~Trainer.create_model_card`].

        Returns:
            The URL of the repository where the model was pushed if `blocking=False`, or a `Future` object tracking the
            progress of the commit if `blocking=True`.
        """
        model_name = kwargs.pop("model_name", None)
        if model_name is None and self.args.should_save:
            if self.args.hub_model_id is None:
                model_name = Path(self.args.output_dir).name
            else:
                model_name = self.args.hub_model_id.split("/")[-1]
        token = token if token is not None else self.args.hub_token

        # In case the user calls this method with args.push_to_hub = False
        if self.hub_model_id is None:
            self.init_hf_repo(token=token)

        # Needs to be executed on all processes for TPU training, but will only save on the processed determined by
        # self.args.should_save.
        self.save_model(_internal_call=True)

        # Only push from one node.
        if not self.is_world_process_zero():
            return

        # Add additional tags in the case the model has already some tags and users pass
        # "tags" argument to `push_to_hub` so that trainer automatically handles internal tags
        # from all models since Trainer does not call `model.push_to_hub`.
        if getattr(self.model, "model_tags", None) is not None:
            if "tags" not in kwargs:
                kwargs["tags"] = []

            # If it is a string, convert it to a list
            if isinstance(kwargs["tags"], str):
                kwargs["tags"] = [kwargs["tags"]]

            for model_tag in self.model.model_tags:
                if model_tag not in kwargs["tags"]:
                    kwargs["tags"].append(model_tag)

        self.create_model_card(model_name=model_name, **kwargs)

        # Wait for the current upload to be finished.
        self._finish_current_push()
        return upload_folder(
            repo_id=self.hub_model_id,
            folder_path=self.args.output_dir,
            commit_message=commit_message,
            token=token,
            run_as_future=not blocking,
            ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"],
            revision=revision,
        )

    #
    # Deprecated code
    #

    def prediction_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> EvalLoopOutput:
        """
        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.

        Works both with or without labels.
        """
        args = self.args

        if not has_length(dataloader):
            raise ValueError("dataloader must implement a working __len__")

        prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only

        # if eval is called w/o train, handle model prep here
        if self.is_deepspeed_enabled and self.deepspeed is None:
            _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

        model = self._wrap_model(self.model, training=False, dataloader=dataloader)

        if len(self.accelerator._models) == 0 and model is self.model:
            model = (
                self.accelerator.prepare(model)
                if self.is_deepspeed_enabled or self.is_fsdp_enabled
                else self.accelerator.prepare_model(model, evaluation_mode=True)
            )

            if self.is_fsdp_enabled:
                self.model = model

            # for the rest of this function `model` is the outside model, whether it was wrapped or not
            if model is not self.model:
                self.model_wrapped = model

            # backward compatibility
            if self.is_deepspeed_enabled:
                self.deepspeed = self.model_wrapped

        # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called
        # while ``train`` is running, cast it to the right dtype first and then put on device
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = (
            dataloader.total_batch_size
            if getattr(dataloader, "_is_accelerate_prepared", False)
            else dataloader.batch_size
        )

        if batch_size is None:
            raise ValueError(
                "Batch size cannot be None. Ensure the dataloader has a valid batch_size or total_batch_size."
            )

        num_examples = self.num_examples(dataloader)
        logger.info(f"\n***** Running {description} *****")
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Batch size = {batch_size}")

        losses_host: torch.Tensor = None
        preds_host: Union[torch.Tensor, List[torch.Tensor]] = None
        labels_host: Union[torch.Tensor, List[torch.Tensor]] = None
        inputs_host: Union[torch.Tensor, List[torch.Tensor]] = None
        metrics: Optional[dict] = None
        eval_set_kwargs: dict = {}

        world_size = max(1, args.world_size)

        eval_losses_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=batch_size)
        if not prediction_loss_only:
            # The actual number of eval_sample can be greater than num_examples in distributed settings (when we pass
            # a batch size to the sampler)
            make_multiple_of = None
            if hasattr(dataloader, "sampler") and isinstance(dataloader.sampler, SequentialDistributedSampler):
                make_multiple_of = dataloader.sampler.batch_size
            preds_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)
            labels_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)
            inputs_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of)

        model.eval()
        if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval):
            self.optimizer.eval()

        if args.past_index >= 0:
            self._past = None

        self.callback_handler.eval_dataloader = dataloader

        for step, inputs in enumerate(dataloader):
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
            main_input_name = getattr(self.model, "main_input_name", "input_ids")
            inputs_decode = (
                self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None
            )

            if loss is not None:
                losses = loss.repeat(batch_size)
                losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0)
            if logits is not None:
                preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100)
            if labels is not None:
                labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100)
            if inputs_decode is not None:
                inputs_host = (
                    inputs_decode
                    if inputs_host is None
                    else nested_concat(inputs_host, inputs_decode, padding_index=-100)
                )
            self.control = self.callback_handler.on_prediction_step(args, self.state, self.control)

            if self.args.batch_eval_metrics:
                if self.compute_metrics is not None and preds_host is not None and labels_host is not None:
                    is_last_step = self.accelerator.gradient_state.end_of_dataloader
                    batch_kwargs = {}
                    batch_kwargs["losses"] = losses_host if "loss" in args.include_for_metrics else None
                    batch_kwargs["inputs"] = inputs_host if "inputs" in args.include_for_metrics else None
                    metrics = self.compute_metrics(
                        EvalPrediction(predictions=preds_host, label_ids=labels_host, **batch_kwargs),
                        compute_result=is_last_step,
                    )

            if self.args.batch_eval_metrics or (
                args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0
            ):
                # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
                eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
                if not prediction_loss_only:
                    preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
                    labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids"))
                    inputs_gatherer.add_arrays(self._gather_and_numpify(inputs_host, "eval_inputs_ids"))

                # Set back to None to begin a new accumulation
                del losses_host, preds_host, labels_host, inputs_host
                torch.cuda.empty_cache()
                losses_host, preds_host, labels_host, inputs_host = None, None, None, None

        if args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of the evaluation loop
            delattr(self, "_past")

        # Gather all remaining tensors and put them back on the CPU
        eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses"))
        if not prediction_loss_only:
            preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds"))
            labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids"))
            inputs_gatherer.add_arrays(self._gather_and_numpify(inputs_host, "eval_inputs_ids"))

        eval_loss = eval_losses_gatherer.finalize()
        preds = preds_gatherer.finalize() if not prediction_loss_only else None
        label_ids = labels_gatherer.finalize() if not prediction_loss_only else None
        inputs_ids = inputs_gatherer.finalize() if not prediction_loss_only else None

        if (
            self.compute_metrics is not None
            and preds is not None
            and label_ids is not None
            and not self.args.batch_eval_metrics
        ):
            eval_set_kwargs["losses"] = eval_loss if "loss" in args.include_for_metrics else None
            eval_set_kwargs["inputs"] = inputs_ids if "inputs" in args.include_for_metrics else None
            metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids, **eval_set_kwargs))
        elif metrics is None:
            metrics = {}

        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

        if eval_loss is not None:
            metrics[f"{metric_key_prefix}_loss"] = eval_loss.mean().item()

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return EvalLoopOutput(predictions=preds, label_ids=label_ids, metrics=metrics, num_samples=num_examples)

    def _gather_and_numpify(self, tensors, name):
        """
        Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before
        concatenating them to `gathered`
        """
        if tensors is None:
            return
        if is_torch_xla_available():
            tensors = nested_xla_mesh_reduce(tensors, name)
        elif is_sagemaker_mp_enabled():
            tensors = smp_gather(tensors)
        elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            tensors = distributed_concat(tensors)

        return nested_numpify(tensors)

    def _add_sm_patterns_to_gitignore(self) -> None:
        """Add SageMaker Checkpointing patterns to .gitignore file."""
        # Make sure we only do this on the main process
        if not self.is_world_process_zero():
            return

        patterns = ["*.sagemaker-uploading", "*.sagemaker-uploaded"]

        # Get current .gitignore content
        if os.path.exists(os.path.join(self.repo.local_dir, ".gitignore")):
            with open(os.path.join(self.repo.local_dir, ".gitignore"), "r") as f:
                current_content = f.read()
        else:
            current_content = ""

        # Add the patterns to .gitignore
        content = current_content
        for pattern in patterns:
            if pattern not in content:
                if content.endswith("\n"):
                    content += pattern
                else:
                    content += f"\n{pattern}"

        # Write the .gitignore file if it has changed
        if content != current_content:
            with open(os.path.join(self.repo.local_dir, ".gitignore"), "w") as f:
                logger.debug(f"Writing .gitignore file. Content: {content}")
                f.write(content)

        self.repo.git_add(".gitignore")

        # avoid race condition with git status
        time.sleep(0.5)

        if not self.repo.is_repo_clean():
            self.repo.git_commit("Add *.sagemaker patterns to .gitignore.")
            self.repo.git_push()

    def create_accelerator_and_postprocess(self):
        # We explicitly don't rely on the `Accelerator` to do gradient accumulation
        grad_acc_kwargs = {}
        if is_accelerate_available("0.28.0") and self.args.accelerator_config.gradient_accumulation_kwargs is not None:
            grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs

        # check if num_steps is attempted to be passed in gradient_accumulation_kwargs
        if "num_steps" in grad_acc_kwargs:
            if self.args.gradient_accumulation_steps > 1:
                # raise because we do not know which setting is intended.
                raise ValueError(
                    "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`"
                    "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`."
                )
            else:
                self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"]

        accelerator_config = self.args.accelerator_config.to_dict()

        if is_accelerate_available("0.28.0"):
            dataloader_config = DataLoaderConfiguration(
                split_batches=accelerator_config.pop("split_batches"),
                dispatch_batches=accelerator_config.pop("dispatch_batches"),
                even_batches=accelerator_config.pop("even_batches"),
                use_seedable_sampler=accelerator_config.pop("use_seedable_sampler"),
            )
            if is_accelerate_available("1.1.0"):
                dataloader_config.data_seed = self.args.data_seed

        non_blocking = accelerator_config.pop("non_blocking")
        if not is_accelerate_available("0.30.0"):
            if non_blocking:
                raise ImportError(
                    "`non_blocking` is only supported in accelerate v0.30.0 and above. Please upgrade accelerate to use this feature."
                )
        else:
            if non_blocking and not self.args.dataloader_pin_memory:
                logger.warning(
                    "`non_blocking` is enabled but `dataloader_pin_memory` is not. For the best performance, it's recommended to enable both."
                )
            dataloader_config.non_blocking = non_blocking
        # this would have been updated above, no need for it anymore
        accelerator_config.pop("gradient_accumulation_kwargs")

        args = {
            "deepspeed_plugin": self.args.deepspeed_plugin,
        }
        if is_accelerate_available("0.28.0"):
            args["dataloader_config"] = dataloader_config
        else:
            args.update(accelerator_config)

        # create accelerator object
        self.accelerator = Accelerator(**args)
        # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag
        self.gather_function = self.accelerator.gather_for_metrics

        if "use_gather_object" in inspect.signature(self.gather_function).parameters.keys():
            self.gather_function = functools.partial(
                self.gather_function, use_gather_object=self.args.eval_use_gather_object
            )

        # deepspeed and accelerate flags covering both trainer args and accelerate launcher
        self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
        self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None

        # post accelerator creation setup
        if self.is_fsdp_enabled:
            fsdp_plugin = self.accelerator.state.fsdp_plugin
            fsdp_plugin.limit_all_gathers = self.args.fsdp_config.get(
                "limit_all_gathers", fsdp_plugin.limit_all_gathers
            )
            fsdp_plugin.activation_checkpointing = self.args.fsdp_config.get(
                "activation_checkpointing", fsdp_plugin.activation_checkpointing
            )
            if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing:
                raise ValueError(
                    "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg "
                    "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic "
                    "when using FSDP."
                )

        if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None:
            self.propagate_args_to_deepspeed()

        # `save_only_model` can't be used with DeepSpeed/FSDP along with `load_best_model_at_end`
        if (
            self.args.save_only_model
            and (self.is_deepspeed_enabled or self.is_fsdp_enabled)
            and self.args.load_best_model_at_end
        ):
            wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP"
            raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.")

        # `auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3
        if (
            self.is_deepspeed_enabled
            and self.accelerator.state.deepspeed_plugin.zero_stage == 3
            and self.args.auto_find_batch_size
        ):
            raise ValueError(
                "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP"
            )

    def propagate_args_to_deepspeed(self, auto_find_batch_size=False):
        """
        Sets values in the deepspeed plugin based on the Trainer args
        """
        from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig

        ds_plugin = self.accelerator.state.deepspeed_plugin

        ds_plugin.hf_ds_config = HfTrainerDeepSpeedConfig(ds_plugin.hf_ds_config.config)
        ds_plugin.deepspeed_config = ds_plugin.hf_ds_config.config
        ds_plugin.hf_ds_config.trainer_config_process(self.args, auto_find_batch_size)

    def _fsdp_qlora_plugin_updates(self):
        if self.is_fsdp_enabled and _is_peft_model(self.model):
            from peft import LoraConfig
            from peft.utils.other import fsdp_auto_wrap_policy

            if isinstance(self.model.active_peft_config, LoraConfig):
                fsdp_plugin = self.accelerator.state.fsdp_plugin
                fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(self.model)
            if (
                getattr(self.model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES
                and self.model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point
                and version.parse(accelerate_version) > version.parse("0.27.0")
            ):
                fsdp_plugin.set_mixed_precision(
                    self.model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True
                )

    def get_batch_samples(self, epoch_iterator, num_batches):
        batch_samples = []
        num_items_in_batch = None
        for _ in range(num_batches):
            try:
                batch_samples += [next(epoch_iterator)]
            except StopIteration:
                break

        if len(batch_samples) > 0 and "labels" in batch_samples[0]:
            # For now we don't support object detection
            try:
                num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples])
            except (TypeError, AttributeError):
                pass

        if self.args.average_tokens_across_devices and num_items_in_batch is not None:
            num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item()

        if torch.is_tensor(num_items_in_batch):
            num_items_in_batch = num_items_in_batch.item()



SUMMARY:
The provided code snippet appears to be part of the Hugging Face Transformers library, specifically within the `Trainer` class. The purpose of this code is to initialize and configure various parameters and plugins for model training using DeepSpeed-Fused (FSDP) or DeepSpeed Zero-3.

Here's a concise summary:

* This code configures various trainer arguments and plugins.
* It checks if the `auto_find_batch_size`, `fsdp_config`, `hf_deepspeed_config`, and other FSDP-related flags are present for configuration purposes.
* The code initializes FSDFP, FSDP Zero-3, or DeepSpeed Zero with a specific plugin depending on the chosen model's parameters.
* The `propagate_args_to_deepspeed` function applies certain arguments to the DeepSpeed plugin during initialization.
* There is also an `_fsdp_qlora_plugin_updates` method that updates FSDFP plugins based on model-related attributes, such as `QuantizationMethod`.
* Throughout the code, there are various checks and validation mechanisms to ensure compatibility with DeepSpeed Zero-3.

Key takeaways:

* The Trainer class uses various flags and plugins for configuration and plugin management.
* The provided code snippet caters specifically to DeepSpeed-Fused (FSDP) and its extensions (e.g., FSDFP Zero-3).
* There are checks for specific model-related attributes and compatibility with different DeepSpeed versions.

This code serves as a foundation for configuring various trainer arguments, models, and plugins for the Transformers library.


