# embedding_metadata

key: chroma:document
string_value:
CODE:

class PoolerStartLogits(nn.Module):
    """
    Compute SQuAD start logits from sequence hidden states.

    Args:
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model.
    """

    def __init__(self, config: PretrainedConfig):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, 1)

    def forward(
        self, hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor] = None
    ) -> torch.FloatTensor:
        """
        Args:
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
                The final hidden states of the model.
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.

        Returns:
            `torch.FloatTensor`: The start logits for SQuAD.
        """
        x = self.dense(hidden_states).squeeze(-1)

        if p_mask is not None:
            if get_parameter_dtype(self) == torch.float16:
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask



SUMMARY:
**PoolerStartLogits Class Summary**

### Purpose and Functionality

The `PoolerStartLogits` class computes SQuAD start logits from sequence hidden states. It is a PyTorch nn.Module that captures the key behavior of a pooler in a BERT-like model.

### Main Methods and Signatures

* `__init__(config: PretrainedConfig)`: Initializes the module with a linear dense layer for computing start logits.
	+ **Parameters**: `config` (PretrainedConfig)
* `forward(hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor] = None) -> torch.FloatTensor`: Computes start logits from sequence hidden states.
	+ **Parameters**:
		- `hidden_states` (torch.FloatTensor of shape `(batch_size, seq_len,.hidden_size)`): Sequence hidden states
		- `p_mask` (Optional[torch.FloatTensor] of shape `(batch_size, seq_len)`)
	+ **Returns**: `start_logits` (torch.FloatTensor)

### Key Implementation Details

* The module uses a linear dense layer to compute start logits from sequence hidden states.
* The input mask (`p_mask`) is applied element-wise to the output by modifying the values according to the model's data type.

### Dependencies and Requirements

* `PretrainedConfig`: The configuration used by the model (inherited interface)
* `nn.Linear` and `torch.FloatTensor`: PyTorch nn modules and tensor types
* A 16-bit floating point model (as indicated by `get_parameter_dtype(self) == torch.float16`) or a 32-bit floating point model


key: chroma:document
string_value:
CODE:

class PoolerEndLogits(nn.Module):
    """
    Compute SQuAD end logits from sequence hidden states.

    Args:
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
            to use.
    """

    def __init__(self, config: PretrainedConfig):
        super().__init__()
        self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
        self.activation = nn.Tanh()
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dense_1 = nn.Linear(config.hidden_size, 1)

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_states: Optional[torch.FloatTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        p_mask: Optional[torch.FloatTensor] = None,
    ) -> torch.FloatTensor:
        """
        Args:
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
                The final hidden states of the model.
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
                The hidden states of the first tokens for the labeled span.
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                The position of the first token for the labeled span.
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.

        <Tip>

        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.

        </Tip>

        Returns:
            `torch.FloatTensor`: The end logits for SQuAD.
        """
        assert (
            start_states is not None or start_positions is not None
        ), "One of start_states, start_positions should be not None"
        if start_positions is not None:
            slen, hsz = hidden_states.shape[-2:]
            start_positions = start_positions[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
            start_states = hidden_states.gather(-2, start_positions)  # shape (bsz, 1, hsz)
            start_states = start_states.expand(-1, slen, -1)  # shape (bsz, slen, hsz)

        x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1))
        x = self.activation(x)
        x = self.LayerNorm(x)
        x = self.dense_1(x).squeeze(-1)

        if p_mask is not None:
            if get_parameter_dtype(self) == torch.float16:
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask



SUMMARY:
**PoolerEndLogits Class Summary**

The `PoolerEndLogits` class is a neural network module that computes SQuAD end logits from sequence hidden states. It takes in the following inputs:

* `hidden_states`: The final hidden states of the model, of shape `(batch_size, seq_len, hidden_size)`
* `start_states`: The hidden states of the first tokens for the labeled span, optional
* `start_positions`: The position of the first token for the labeled span, optional
* `p_mask`: Mask for tokens at invalid position, optional

The class applies the following operations:

1. Concatenates `hidden_states` and `start_states` using a linear layer (`dense_0`)
2. Applies the tanh activation function to the concatenated output
3. Normalizes the output using layer normalization (`LayerNorm`)
4. Applies another linear layer (`dense_1`) to produce the end logits

The class also applies masking to the output using `p_mask` if provided.

**Key Implementation Details**

* Use of concatenated concatenation and layer normalization for feature extraction.
* Tanh activation function applied to the output.
* Masking of output using `p_mask`.

**Notable Dependencies and Requirements**

* `nn.Module`: The base class for neural network modules.
* `PretrainedConfig`: A config configuration object used by the model.
* `torch.FloatTensor` and `torch.LongTensor`: Tensor data types used for input/output.
* `get_parameter_dtype` function: A utility function to determine the data type of a parameter.

**Interface and Method Signatures**

* `__init__(config: PretrainedConfig)`: Initializes the module with the given config.
* `forward(hidden_states: torch.FloatTensor, ...)`: Performs the forward pass on the inputs, returning the end logits.


key: chroma:document
string_value:
CODE:

class PoolerAnswerClass(nn.Module):
    """
    Compute SQuAD 2.0 answer class from classification and start tokens hidden states.

    Args:
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model.
    """

    def __init__(self, config):
        super().__init__()
        self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
        self.activation = nn.Tanh()
        self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False)

    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_states: Optional[torch.FloatTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        cls_index: Optional[torch.LongTensor] = None,
    ) -> torch.FloatTensor:
        """
        Args:
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
                The final hidden states of the model.
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
                The hidden states of the first tokens for the labeled span.
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                The position of the first token for the labeled span.
            cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Position of the CLS token for each sentence in the batch. If `None`, takes the last token.

        <Tip>

        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.

        </Tip>

        Returns:
            `torch.FloatTensor`: The SQuAD 2.0 answer class.
        """
        # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample.
        hsz = hidden_states.shape[-1]
        assert (
            start_states is not None or start_positions is not None
        ), "One of start_states, start_positions should be not None"
        if start_positions is not None:
            start_positions = start_positions[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
            start_states = hidden_states.gather(-2, start_positions).squeeze(-2)  # shape (bsz, hsz)

        if cls_index is not None:
            cls_index = cls_index[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
            cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, hsz)
        else:
            cls_token_state = hidden_states[:, -1, :]  # shape (bsz, hsz)

        x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1))
        x = self.activation(x)
        x = self.dense_1(x).squeeze(-1)



SUMMARY:
**Summary: Pooler Answer Class**

The `PoolerAnswerClass` is a PyTorch neural network module that computes the SQuAD 2.0 answer class from classification and start tokens' hidden states.

**Main Purpose and Functionality:**
The primary function of this class is to take in input (`hidden_states`, `start_states`, `start_positions`, and `cls_index`) and produce an output (`torch.FloatTensor`) representing the SQuAD 2.0 answer class.

**Key Implementation Details:**

* The model uses two linear layers (`dense_0` and `dense_1`) with ReLU activation (`activation`) in between.
* The input to `dense_0` is a concatenation of `start_states` and `cls_token_state`, where `cls_token_state` is obtained by gathering the hidden states along the sequence dimension with respect to either `start_positions` or `cls_index`.
* The output of `dense_1` is used as the final answer class.

**Notable Dependencies and Requirements:**

* `PretrainedConfig`: a configuration object required for model initialization.
* `torch.FloatTensor`, `Optional[torch.FloatTensor]`, `torch.LongTensor`: input/output data types.
* `nn.Module`, `nn.Linear`: PyTorch neural network module and linear layer classes.

Note that the original code uses some mathematical operations not extensively discussed here.


key: chroma:document
string_value:
CODE:

class SQuADHead(nn.Module):
    r"""
    A SQuAD head inspired by XLNet.

    Args:
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
            to use.
    """

    def __init__(self, config):
        super().__init__()
        self.start_n_top = config.start_n_top
        self.end_n_top = config.end_n_top

        self.start_logits = PoolerStartLogits(config)
        self.end_logits = PoolerEndLogits(config)
        self.answer_class = PoolerAnswerClass(config)

    @replace_return_docstrings(output_type=SquadHeadOutput, config_class=PretrainedConfig)
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_positions: Optional[torch.LongTensor] = None,
        end_positions: Optional[torch.LongTensor] = None,
        cls_index: Optional[torch.LongTensor] = None,
        is_impossible: Optional[torch.LongTensor] = None,
        p_mask: Optional[torch.FloatTensor] = None,
        return_dict: bool = False,
    ) -> Union[SquadHeadOutput, Tuple[torch.FloatTensor]]:
        """
        Args:
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
                Final hidden states of the model on the sequence tokens.
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Positions of the first token for the labeled span.
            end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Positions of the last token for the labeled span.
            cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Position of the CLS token for each sentence in the batch. If `None`, takes the last token.
            is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Whether the question has a possible answer in the paragraph or not.
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.

        Returns:
        """
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)

        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, let's remove the dimension added by batch splitting
            for x in (start_positions, end_positions, cls_index, is_impossible):
                if x is not None and x.dim() > 1:
                    x.squeeze_(-1)

            # during training, compute the end logits based on the ground truth of the start position
            end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask)

            loss_fct = CrossEntropyLoss()
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2

            if cls_index is not None and is_impossible is not None:
                # Predict answerability from the representation of CLS and START
                cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
                loss_fct_cls = nn.BCEWithLogitsLoss()
                cls_loss = loss_fct_cls(cls_logits, is_impossible)

                # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss
                total_loss += cls_loss * 0.5

            return SquadHeadOutput(loss=total_loss) if return_dict else (total_loss,)

        else:
            # during inference, compute the end logits based on beam search
            bsz, slen, hsz = hidden_states.size()
            start_log_probs = nn.functional.softmax(start_logits, dim=-1)  # shape (bsz, slen)

            start_top_log_probs, start_top_index = torch.topk(
                start_log_probs, self.start_n_top, dim=-1
            )  # shape (bsz, start_n_top)
            start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz)  # shape (bsz, start_n_top, hsz)
            start_states = torch.gather(hidden_states, -2, start_top_index_exp)  # shape (bsz, start_n_top, hsz)
            start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1)  # shape (bsz, slen, start_n_top, hsz)

            hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(
                start_states
            )  # shape (bsz, slen, start_n_top, hsz)
            p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
            end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
            end_log_probs = nn.functional.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)

            end_top_log_probs, end_top_index = torch.topk(
                end_log_probs, self.end_n_top, dim=1
            )  # shape (bsz, end_n_top, start_n_top)
            end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
            end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)

            start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs)
            cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index)

            if not return_dict:
                return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits)
            else:
                return SquadHeadOutput(
                    start_top_log_probs=start_top_log_probs,
                    start_top_index=start_top_index,
                    end_top_log_probs=end_top_log_probs,
                    end_top_index=end_top_index,
                    cls_logits=cls_logits,


SUMMARY:
**SQuADHead Class Summary**

The `SQuADHead` class is a PyTorch neural network module that inherits from `nn.Module`. It is designed to work with the SQuAD dataset and provides a head for QA tasks.

**Primary Purpose:**
Provide a custom neural network head for SQuAD data, capable of generating start and end logits for labeled spans.

**Key Implementation Details:**

* `forward()` method:
	+ Computes start logits using `start_logits`
	+ Computes end logits based on ground truth start positions or beam search
	+ Predicts answerability from CLS token representation
* Utilizes multiple losses, including:
	+ CrossEntropyLoss for start and end position prediction
	+ BCEWithLogitsLoss for predicting answerability

**Notable Dependencies/Requirements:**

* `PoolerStartLogits`, `PoolerEndLogits`, and `PoolerAnswerClass` classes
* `CrossEntropyLoss` and `BCEWithLogitsLoss` loss functions


key: chroma:document
string_value:
CODE:

class SequenceSummary(nn.Module):
    r"""
    Compute a single vector summary of a sequence hidden states.

    Args:
        config ([`PretrainedConfig`]):
            The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
            config class of your model for the default values it uses):

            - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:

                - `"last"` -- Take the last token hidden state (like XLNet)
                - `"first"` -- Take the first token hidden state (like Bert)
                - `"mean"` -- Take the mean of all tokens hidden states
                - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
                - `"attn"` -- Not implemented now, use multi-head attention

            - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
            - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
              (otherwise to `config.hidden_size`).
            - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
              another string or `None` will add no activation.
            - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
            - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.
    """

    def __init__(self, config: PretrainedConfig):
        super().__init__()

        self.summary_type = getattr(config, "summary_type", "last")
        if self.summary_type == "attn":
            # We should use a standard multi-head attention module with absolute positional embedding for that.
            # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276
            # We can probably just use the multi-head attention module of PyTorch >=1.1.0
            raise NotImplementedError

        self.summary = Identity()
        if hasattr(config, "summary_use_proj") and config.summary_use_proj:
            if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0:
                num_classes = config.num_labels
            else:
                num_classes = config.hidden_size
            self.summary = nn.Linear(config.hidden_size, num_classes)

        activation_string = getattr(config, "summary_activation", None)
        self.activation: Callable = get_activation(activation_string) if activation_string else Identity()

        self.first_dropout = Identity()
        if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
            self.first_dropout = nn.Dropout(config.summary_first_dropout)

        self.last_dropout = Identity()
        if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
            self.last_dropout = nn.Dropout(config.summary_last_dropout)

    def forward(
        self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None
    ) -> torch.FloatTensor:
        """
        Compute a single vector summary of a sequence hidden states.

        Args:
            hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`):
                The hidden states of the last layer.
            cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):
                Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.

        Returns:
            `torch.FloatTensor`: The summary of the sequence hidden states.
        """
        if self.summary_type == "last":
            output = hidden_states[:, -1]
        elif self.summary_type == "first":
            output = hidden_states[:, 0]
        elif self.summary_type == "mean":
            output = hidden_states.mean(dim=1)
        elif self.summary_type == "cls_index":
            if cls_index is None:
                cls_index = torch.full_like(
                    hidden_states[..., :1, :],
                    hidden_states.shape[-2] - 1,
                    dtype=torch.long,
                )
            else:
                cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
                cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
            # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
            output = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, XX, hidden_size)
        elif self.summary_type == "attn":
            raise NotImplementedError

        output = self.first_dropout(output)
        output = self.summary(output)
        output = self.activation(output)
        output = self.last_dropout(output)



SUMMARY:
**Summary**

The `SequenceSummary` class is part of a neural network model that computes a single vector summary of sequence hidden states. It takes in a configuration object and uses it to determine the method for making the summary (e.g., "last", "first", "mean", "cls_index", or "attn").

**Key Implementation Details**

* Methods:
	+ `__init__`: Initializes the model with a given configuration.
	+ `forward`: Computes the vector summary of sequence hidden states based on the chosen method (`summary_type`).
* Activation functions:
	+ Identity (no activation)
	+ tanh (if specified in `config.summary_activation`)
* Droputouts:
	+ First dropout before projection and activation
	+ Last dropout after projection and activation

**Notable Dependencies**

* PyTorch library for neural networks
* `PretrainedConfig` class from the model configuration class (not shown)

**Interfaces**

* Implements a standard multi-head attention module (`nn.Linear`, `torch Dropsout`) when using "attn" summary type is not supported.

Note: The exact implementation of the `get_activation` function is not shown.


key: chroma:document
string_value:
CODE:

def unwrap_model(model: nn.Module, recursive: bool = False) -> nn.Module:
    """
    Recursively unwraps a model from potential containers (as used in distributed training).

    Args:
        model (`torch.nn.Module`): The model to unwrap.
        recursive (`bool`, *optional*, defaults to `False`):
            Whether to recursively extract all cases of `module.module` from `model` as well as unwrap child sublayers
            recursively, not just the top-level distributed containers.
    """
    # Use accelerate implementation if available (should always be the case when using torch)
    # This is for pytorch, as we also have to handle things like dynamo
    if is_accelerate_available():
        kwargs = {}
        if recursive:
            if not is_accelerate_available("0.29.0"):
                raise RuntimeError(
                    "Setting `recursive=True` to `unwrap_model` requires `accelerate` v0.29.0. Please upgrade your version of accelerate"
                )
            else:
                kwargs["recursive"] = recursive
        return extract_model_from_parallel(model, **kwargs)
    else:
        # since there could be multiple levels of wrapping, unwrap recursively
        if hasattr(model, "module"):
            return unwrap_model(model.module)
        else:


SUMMARY:
**Function Summary**

`unwrap_model`: Unwraps a PyTorch model from potential containers used in distributed training.

**Primary Purpose**: Recursively unwraps a model to extract the underlying computation graph.

**Implementation Details**:

* Uses `accelerate` implementation if available (version 0.29.0 or higher recommended).
* If recursive mode enabled, extracts all cases of `module.module` from the model and its sublayers recursively using `extract_model_from_parallel`.

**Notable Dependencies/Requirements**:

* `torch.nn.Module`
* Accelerate library (specifically version 0.29.0 or higher)
* External dependency on `is_accelerate_available` function.

**Method Signatures**:

* `unwrap_model(model: nn.Module, recursive: bool = False) -> nn.Module`

Note that the original code has an additional branch to handle models without a `module` attribute, which is currently returning the model as-is. This branch is not explicitly summarized in this concise version.


key: chroma:document
string_value:
CODE:

def expand_device_map(device_map, param_names, start_prefix):
    """
    Expand a device map to return the correspondance parameter name to device.
    """
    new_device_map = {}
    param_names = [p[len(start_prefix) :] for p in param_names if p.startswith(start_prefix)]
    for module, device in device_map.items():
        new_device_map.update(
            {p: device for p in param_names if p == module or p.startswith(f"{module}.") or module == ""}
        )


SUMMARY:
**Concise Summary**

This function `expand_device_map` expands a device map by mapping parameter names to devices. It:

* Filters parameter names that start with a specified prefix
* Updates the new device map with corresponding device information from the original device map

**Key Implementation Details:**

* Uses list comprehension to filter parameter names
* Iterates over the `device_map` dictionary, updating the `new_device_map`
* Employs the ternary operator for conditional parameter name matching

**Notable Dependencies/Requirements:**

* `device_map`: a dictionary containing module-device pairs
* `param_names`: a list of parameter names with prefix filtering applied
* `start_prefix`: a string used to filter parameter names


key: chroma:document
string_value:
CODE:

def get_disk_only_shard_files(device_map, sharded_metadata, start_prefix):
    """
    Returns the list of shard files containing only weights offloaded to disk.
    """

    weight_map = {
        p[len(start_prefix) :]: v for p, v in sharded_metadata["weight_map"].items() if p.startswith(start_prefix)
    }
    files_content = collections.defaultdict(list)
    for weight_name, filename in weight_map.items():
        while len(weight_name) > 0 and weight_name not in device_map:
            weight_name = ".".join(weight_name.split(".")[:-1])
        files_content[filename].append(device_map[weight_name])



SUMMARY:
**Summary:**

The `get_disk_only_shard_files` function takes three inputs:
- `device_map`: Mapping of weight names to device assignments
- `sharded_metadata`: Metadata containing the `weight_map`
- `start_prefix`: Starting prefix for weights in `weight_map`

It returns a list of files with only disk-attached devices.

**Key Implementation Details:**

* Uses dictionary comprehension to filter `sharded_metadata` and extract relevant weight names.
* Iterates through weight names, updating them by removing device attachments until they are present in `device_map`.
* Creates a collections defaultdict to store files with their contents.

**Notable Dependencies/Requirements:**

None specified, but assumes presence of:
* `collections` module for `defaultdict`
* None external dependencies.


key: chroma:document
string_value:
CODE:

class ExportableState:
    """
    A class for objects that include the ability to have its state
    be saved during `Trainer._save_checkpoint` and loaded back in during
    `Trainer._load_from_checkpoint`.

    These must implement a `state` function that gets called during the respective
    Trainer function call. It should only include parameters and attributes needed to
    recreate the state at a particular time, to avoid utilizing pickle/maintain standard
    file IO writing.

    Example:

    ```python
    class EarlyStoppingCallback(TrainerCallback, ExportableState):
        def __init__(self, early_stopping_patience: int = 1, early_stopping_threshold: Optional[float] = 0.0):
            self.early_stopping_patience = early_stopping_patience
            self.early_stopping_threshold = early_stopping_threshold
            # early_stopping_patience_counter denotes the number of times validation metrics failed to improve.
            self.early_stopping_patience_counter = 0

        def state(self) -> dict:
            return {
                "args": {
                    "early_stopping_patience": self.early_stopping_patience,
                    "early_stopping_threshold": self.early_stopping_threshold,
                },
                "attributes": {
                    "early_stopping_patience_counter": self.early_stopping_patience_counter,
                }
            }
    ```"""

    def state(self) -> dict:
        raise NotImplementedError("You must implement a `state` function to utilize this class.")

    @classmethod
    def from_state(cls, state):
        instance = cls(**state["args"])
        for k, v in state["attributes"].items():
            setattr(instance, k, v)


SUMMARY:
**Class Summary: ExportableState**

### Purpose and Functionality
The `ExportableState` class enables objects to store and retrieve their internal state through the `Trainer._save_checkpoint` and `Trainer._load_from_checkpoint` functions.

**Key Implementation Details**

* The object must implement a `state` function that returns a dictionary containing only necessary parameters and attributes for state recreation (`raises NotImplementedError`)
* A class method `from_state` creates an instance from a provided state dictionary by unpacking argument keywords and setting attribute values (`instance = cls(**state["args"]); setattr(instance, k, v)`)

### Notable Dependencies or Requirements

* The object must be instantiated as part of a `TrainerCallback`
* Inheriting from `ExportableState` requires implementing the `state` method
* External dependency on `Trainer` class using `Trainer._save_checkpoint` and `Trainer._load_from_checkpoint`

### Key Implementation Signatures

* `state(self) -> dict`: Must be implemented by inheriting classes
* `from_state(state: dict) -> instance`: Class method for state dictionary deserialization


key: chroma:document
string_value:
CODE:

class TrainerCallback:
    # no-format
    """
    A class for objects that will inspect the state of the training loop at some events and take some decisions. At
    each of those events the following arguments are available:

    Args:
        args ([`TrainingArguments`]):
            The training arguments used to instantiate the [`Trainer`].
        state ([`TrainerState`]):
            The current state of the [`Trainer`].
        control ([`TrainerControl`]):
            The object that is returned to the [`Trainer`] and can be used to make some decisions.
        model ([`PreTrainedModel`] or `torch.nn.Module`):
            The model being trained.
        tokenizer ([`PreTrainedTokenizer`]):
            The tokenizer used for encoding the data. This is deprecated in favour of `processing_class`.
        processing_class ([`PreTrainedTokenizer` or `BaseImageProcessor` or `ProcessorMixin` or `FeatureExtractionMixin`]):
            The processing class used for encoding the data. Can be a tokenizer, a processor, an image processor or a feature extractor.
        optimizer (`torch.optim.Optimizer`):
            The optimizer used for the training steps.
        lr_scheduler (`torch.optim.lr_scheduler.LambdaLR`):
            The scheduler used for setting the learning rate.
        train_dataloader (`torch.utils.data.DataLoader`, *optional*):
            The current dataloader used for training.
        eval_dataloader (`torch.utils.data.DataLoader`, *optional*):
            The current dataloader used for evaluation.
        metrics (`Dict[str, float]`):
            The metrics computed by the last evaluation phase.

            Those are only accessible in the event `on_evaluate`.
        logs  (`Dict[str, float]`):
            The values to log.

            Those are only accessible in the event `on_log`.

    The `control` object is the only one that can be changed by the callback, in which case the event that changes it
    should return the modified version.

    The argument `args`, `state` and `control` are positionals for all events, all the others are grouped in `kwargs`.
    You can unpack the ones you need in the signature of the event using them. As an example, see the code of the
    simple [`~transformers.PrinterCallback`].

    Example:

    ```python
    class PrinterCallback(TrainerCallback):
        def on_log(self, args, state, control, logs=None, **kwargs):
            _ = logs.pop("total_flos", None)
            if state.is_local_process_zero:
                print(logs)
    ```"""

    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the end of the initialization of the [`Trainer`].
        """
        pass

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the beginning of training.
        """
        pass

    def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the end of training.
        """
        pass

    def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the beginning of an epoch.
        """
        pass

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the end of an epoch.
        """
        pass

    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the beginning of a training step. If using gradient accumulation, one training step might take
        several inputs.
        """
        pass

    def on_pre_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called before the optimizer step but after gradient clipping. Useful for monitoring gradients.
        """
        pass

    def on_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called after the optimizer step but before gradients are zeroed out. Useful for monitoring gradients.
        """
        pass

    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the end of an substep during gradient accumulation.
        """
        pass

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called at the end of a training step. If using gradient accumulation, one training step might take
        several inputs.
        """
        pass

    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called after an evaluation phase.
        """
        pass

    def on_predict(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics, **kwargs):
        """
        Event called after a successful prediction.
        """
        pass

    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called after a checkpoint save.
        """
        pass

    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called after logging the last logs.
        """
        pass

    def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        """
        Event called after a prediction step.
        """


SUMMARY:
**Class Summary: TrainerCallback**

**Main Purpose and Functionality:**
The `TrainerCallback` class is designed to inspect the state of a training loop at various events and take decisions. It provides a callback interface for objects that need to interact with the trainer.

**Key Implementation Details:**

* Provides 8 event callbacks:
	+ `on_init_end`, `on_train_begin`, `on_train_end`, `on_epoch_begin`, `on_epoch_end`, `on_step_begin`, `on_pre_optimizer_step`, and `on_optimizer_step`
* Arguments are available for all events, grouped into positionals (`args`, `state`, `control`) and keyword arguments (`kwargs`)
* Examples of event callback implementations can be found in the code comments

**Notable Dependencies or Requirements:**

* `TrainingArguments` class
* `TrainerState` class
* `TrainerControl` object
* `PreTrainedModel` interface
* `tokenzier` and `processing_class` dependencies have been deprecated in favor of `processing_class`
* Interfaces with external libraries, including:
	+ PyTorch (`torch.optim.Optimizer`, `torch.optim.lr_scheduler.LambdaLR`)
	+ Transformers library (`PreTrainedTokenizer`, `BaseImageProcessor`)


key: chroma:document
string_value:
CODE:

class CallbackHandler(TrainerCallback):
    """Internal class that just calls the list of callbacks in order."""

    def __init__(self, callbacks, model, processing_class, optimizer, lr_scheduler):
        self.callbacks = []
        for cb in callbacks:
            self.add_callback(cb)
        self.model = model
        self.processing_class = processing_class
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.train_dataloader = None
        self.eval_dataloader = None

        if not any(isinstance(cb, DefaultFlowCallback) for cb in self.callbacks):
            logger.warning(
                "The Trainer will not work properly if you don't have a `DefaultFlowCallback` in its callbacks. You\n"
                + "should add one before training with `trainer.add_callback(DefaultFlowCallback). The current list of"
                + "callbacks is\n:"
                + self.callback_list
            )

    def add_callback(self, callback):
        cb = callback() if isinstance(callback, type) else callback
        cb_class = callback if isinstance(callback, type) else callback.__class__
        if cb_class in [c.__class__ for c in self.callbacks]:
            logger.warning(
                f"You are adding a {cb_class} to the callbacks of this Trainer, but there is already one. The current"
                + "list of callbacks is\n:"
                + self.callback_list
            )
        self.callbacks.append(cb)

    def pop_callback(self, callback):
        if isinstance(callback, type):
            for cb in self.callbacks:
                if isinstance(cb, callback):
                    self.callbacks.remove(cb)
                    return cb
        else:
            for cb in self.callbacks:
                if cb == callback:
                    self.callbacks.remove(cb)
                    return cb

    def remove_callback(self, callback):
        if isinstance(callback, type):
            for cb in self.callbacks:
                if isinstance(cb, callback):
                    self.callbacks.remove(cb)
                    return
        else:
            self.callbacks.remove(callback)

    @property
    def callback_list(self):
        return "\n".join(cb.__class__.__name__ for cb in self.callbacks)

    def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_init_end", args, state, control)

    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_training_stop = False
        return self.call_event("on_train_begin", args, state, control)

    def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_train_end", args, state, control)

    def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_epoch_stop = False
        return self.call_event("on_epoch_begin", args, state, control)

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_epoch_end", args, state, control)

    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_log = False
        control.should_evaluate = False
        control.should_save = False
        return self.call_event("on_step_begin", args, state, control)

    def on_pre_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_pre_optimizer_step", args, state, control)

    def on_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_optimizer_step", args, state, control)

    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_substep_end", args, state, control)

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_step_end", args, state, control)

    def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics):
        control.should_evaluate = False
        return self.call_event("on_evaluate", args, state, control, metrics=metrics)

    def on_predict(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics):
        return self.call_event("on_predict", args, state, control, metrics=metrics)

    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        control.should_save = False
        return self.call_event("on_save", args, state, control)

    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs):
        control.should_log = False
        return self.call_event("on_log", args, state, control, logs=logs)

    def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl):
        return self.call_event("on_prediction_step", args, state, control)

    def call_event(self, event, args, state, control, **kwargs):
        for callback in self.callbacks:
            result = getattr(callback, event)(
                args,
                state,
                control,
                model=self.model,
                processing_class=self.processing_class,
                optimizer=self.optimizer,
                lr_scheduler=self.lr_scheduler,
                train_dataloader=self.train_dataloader,
                eval_dataloader=self.eval_dataloader,
                **kwargs,
            )
            # A Callback can skip the return of `control` if it doesn't change it.
            if result is not None:
                control = result


SUMMARY:
**Summary**

The `CallbackHandler` class is an internal utility in TrainerCallback that enables event-driven callbacks to be executed during the training loop. It's responsible for managing a list of callbacks, handling lifecycle events (e.g., training start/stop, epoch begin/end), and propagating control updates to the trainer.

**Key Implementation Details**

*   The `add_callback`, `pop_callback` and `remove_callback` methods manage the callback list, handling duplicate notifications.
*   The `on_\*` methods define event handlers for various Trainer events (e.g., `on_init_end`, `on_train_begin`, `on_epoch_begin`) that delegate to corresponding callbacks. These callbacks can modify the control object before propagating it back.
*   The `call_event` method iterates over the list of callbacks, executing each one and updating the control object accordingly.

**Notable Dependencies/Organizational Requirements**

*   The `TrainerCallback` class is an ancestor that defines common interfaces for callback implementations, likely inherited by custom callbacks.
*   External dependencies:
    *   Logging module (`logger`) using Python's built-in logging capabilities or a third-party logging library (not specified in the provided code snippet).
*   Inherited Interfaces:
    *   `TrainerCallback`: The base class that provides a set of common event handlers and lifecycle methods for Trainer callbacks.
*   Implemented Interfaces:
    *   The internal interface is primarily designed to work seamlessly with the existing Trainer framework.


key: chroma:document
string_value:
CODE:

class DefaultFlowCallback(TrainerCallback):
    """
    A [`TrainerCallback`] that handles the default flow of the training loop for logs, evaluation and checkpoints.
    """

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        # Log
        if state.global_step == 1 and args.logging_first_step:
            control.should_log = True
        if args.logging_strategy == IntervalStrategy.STEPS and state.global_step % state.logging_steps == 0:
            control.should_log = True

        # Evaluate
        if (
            args.eval_strategy == IntervalStrategy.STEPS
            and state.global_step % state.eval_steps == 0
            and args.eval_delay <= state.global_step
        ):
            control.should_evaluate = True

        # Save
        if (
            args.save_strategy == SaveStrategy.STEPS
            and state.save_steps > 0
            and state.global_step % state.save_steps == 0
        ):
            control.should_save = True

        # End training
        if state.global_step >= state.max_steps:
            control.should_training_stop = True
            # Save the model at the end if we have a save strategy
            if args.save_strategy not in [SaveStrategy.NO, SaveStrategy.BEST]:
                control.should_save = True

        return control

    def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        # Log
        if args.logging_strategy == IntervalStrategy.EPOCH:
            control.should_log = True

        # Evaluate
        if args.eval_strategy == IntervalStrategy.EPOCH and args.eval_delay <= state.epoch:
            control.should_evaluate = True

        # Save
        if args.save_strategy == SaveStrategy.EPOCH:
            control.should_save = True



SUMMARY:
**DefaultFlowCallback Class Summary**

The `DefaultFlowCallback` class is a custom callback for training loops in Hugging Face's Transformers, responsible for handling the default flow of the training loop. It provides functionality for logging, evaluation, and checkpointing.

**Primary Purpose:**
Implementing trainer callbacks for managing the workflow of Hugging Face Transformer's training.

**Key Implementation Details:**

*   **on_step_end:** Decides whether to log, evaluate, save model, or stop training based on specified strategies (e.g., `IntervalStrategy.STEPS`, `args.logging_strategy`).
*   **on_epoch_end:** Similar behavior as `on_step_end` but tailored for epoch completion.

**Notable Dependencies and Requirements:**

*   `[TrainerState]`: Accessing global step, save steps, and logging steps.
*   `[TrainerControl]`: Log, evaluate, and save model control.
*   `[TrainingArguments]`: Arguments related to logging and evaluation strategies.


key: chroma:document
string_value:
CODE:

class ProgressCallback(TrainerCallback):
    """
    A [`TrainerCallback`] that displays the progress of training or evaluation.
    You can modify `max_str_len` to control how long strings are truncated when logging.
    """

    def __init__(self, max_str_len: int = 100):
        """
        Initialize the callback with optional max_str_len parameter to control string truncation length.

        Args:
            max_str_len (`int`):
                Maximum length of strings to display in logs.
                Longer strings will be truncated with a message.
        """
        self.training_bar = None
        self.prediction_bar = None
        self.max_str_len = max_str_len

    def on_train_begin(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            self.training_bar = tqdm(total=state.max_steps, dynamic_ncols=True)
        self.current_step = 0

    def on_step_end(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            self.training_bar.update(state.global_step - self.current_step)
            self.current_step = state.global_step

    def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
        if state.is_world_process_zero and has_length(eval_dataloader):
            if self.prediction_bar is None:
                self.prediction_bar = tqdm(
                    total=len(eval_dataloader), leave=self.training_bar is None, dynamic_ncols=True
                )
            self.prediction_bar.update(1)

    def on_evaluate(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            if self.prediction_bar is not None:
                self.prediction_bar.close()
            self.prediction_bar = None

    def on_predict(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            if self.prediction_bar is not None:
                self.prediction_bar.close()
            self.prediction_bar = None

    def on_log(self, args, state, control, logs=None, **kwargs):
        if state.is_world_process_zero and self.training_bar is not None:
            # make a shallow copy of logs so we can mutate the fields copied
            # but avoid doing any value pickling.
            shallow_logs = {}
            for k, v in logs.items():
                if isinstance(v, str) and len(v) > self.max_str_len:
                    shallow_logs[k] = (
                        f"[String too long to display, length: {len(v)} > {self.max_str_len}. "
                        "Consider increasing `max_str_len` if needed.]"
                    )
                else:
                    shallow_logs[k] = v
            _ = shallow_logs.pop("total_flos", None)
            # round numbers so that it looks better in console
            if "epoch" in shallow_logs:
                shallow_logs["epoch"] = round(shallow_logs["epoch"], 2)
            self.training_bar.write(str(shallow_logs))

    def on_train_end(self, args, state, control, **kwargs):
        if state.is_world_process_zero:
            self.training_bar.close()


SUMMARY:
**Summary: ProgressCallback class**

`ProgressCallback` is a custom callback for `TrainerCallback` that displays training or evaluation progress. It utilizes `tqdm` for dynamic progress tracking.

**Key Implementation Details**

* Utilizes `tqdm` for logging and displaying training/evaluation progress.
* Truncates long strings to `max_str_len` in log output.
* Implements methods:
	+ `on_train_begin`: Initializes training progress bar.
	+ `on_step_end`: Updates training progress bar with each step end.
	+ `on_prediction_step`: Initializes prediction progress bar and updates its count during evaluation.
	+ `on_evaluate`, `on_predict`, `on_log`, and `on_train_end`: Handles cleanup of progress bars and truncates log output as needed.

**Notable Dependencies/Requirements**

* Uses `tqdm` for dynamic progress tracking.
* Inheriting from `TrainerCallback`.


key: chroma:document
string_value:
CODE:

class PrinterCallback(TrainerCallback):
    """
    A bare [`TrainerCallback`] that just prints the logs.
    """

    def on_log(self, args, state, control, logs=None, **kwargs):
        _ = logs.pop("total_flos", None)
        if state.is_local_process_zero:


SUMMARY:
**Summary:**

The `PrinterCallback` class is a custom Trainer Callback that prints training logs from the `args`, `state`, and `logs` objects. It removes the "total_flos" key from the logs dictionary when available.

**Key Implementation Details:**

- Prints training logs using `print(logs)` when running on local process 0.
- Removes "total_flos" key from `logs` dictionary using `_ = logs.pop("total_flos", None)`.
- Inherited from `TrainerCallback`.

**Notable Dependencies and Requirements:**

- Extends `TrainerCallback` (from the same package).
- Assumes presence of `args`, `state`, and `logs` objects.
- Does not expose any external dependencies.


key: chroma:document
string_value:
CODE:

class EarlyStoppingCallback(TrainerCallback, ExportableState):
    """
    A [`TrainerCallback`] that handles early stopping.

    Args:
        early_stopping_patience (`int`):
            Use with `metric_for_best_model` to stop training when the specified metric worsens for
            `early_stopping_patience` evaluation calls.
        early_stopping_threshold(`float`, *optional*):
            Use with TrainingArguments `metric_for_best_model` and `early_stopping_patience` to denote how much the
            specified metric must improve to satisfy early stopping conditions. `

    This callback depends on [`TrainingArguments`] argument *load_best_model_at_end* functionality to set best_metric
    in [`TrainerState`]. Note that if the [`TrainingArguments`] argument *save_steps* differs from *eval_steps*, the
    early stopping will not occur until the next save step.
    """

    def __init__(self, early_stopping_patience: int = 1, early_stopping_threshold: Optional[float] = 0.0):
        self.early_stopping_patience = early_stopping_patience
        self.early_stopping_threshold = early_stopping_threshold
        # early_stopping_patience_counter denotes the number of times validation metrics failed to improve.
        self.early_stopping_patience_counter = 0

    def check_metric_value(self, args, state, control, metric_value):
        # best_metric is set by code for load_best_model
        operator = np.greater if args.greater_is_better else np.less
        if state.best_metric is None or (
            operator(metric_value, state.best_metric)
            and abs(metric_value - state.best_metric) > self.early_stopping_threshold
        ):
            self.early_stopping_patience_counter = 0
        else:
            self.early_stopping_patience_counter += 1

    def on_train_begin(self, args, state, control, **kwargs):
        assert args.load_best_model_at_end, "EarlyStoppingCallback requires load_best_model_at_end = True"
        assert (
            args.metric_for_best_model is not None
        ), "EarlyStoppingCallback requires metric_for_best_model is defined"
        assert (
            args.eval_strategy != IntervalStrategy.NO
        ), "EarlyStoppingCallback requires IntervalStrategy of steps or epoch"

    def on_evaluate(self, args, state, control, metrics, **kwargs):
        metric_to_check = args.metric_for_best_model
        if not metric_to_check.startswith("eval_"):
            metric_to_check = f"eval_{metric_to_check}"
        metric_value = metrics.get(metric_to_check)

        if metric_value is None:
            logger.warning(
                f"early stopping required metric_for_best_model, but did not find {metric_to_check} so early stopping"
                " is disabled"
            )
            return

        self.check_metric_value(args, state, control, metric_value)
        if self.early_stopping_patience_counter >= self.early_stopping_patience:
            control.should_training_stop = True

    def state(self) -> dict:
        return {
            "args": {
                "early_stopping_patience": self.early_stopping_patience,
                "early_stopping_threshold": self.early_stopping_threshold,
            },
            "attributes": {
                "early_stopping_patience_counter": self.early_stopping_patience_counter,
            },


SUMMARY:
**EarlyStoppingCallback Class Summary**

The `EarlyStoppingCallback` class is a custom callback for training models, designed to implement early stopping strategies based on pre-defined metrics.

### Main Purpose and Functionality

* The callback tracks the improvement of metric values during training and stops the training process when the specified metric stops improving (early stopping threshold).

### Key Implementation Details

- The callback checks the value of a specific metric (`metric_for_best_model`) at regular intervals.
- It triggers early stopping when the metric's value does not improve after a specified number of evaluation calls (`early_stopping_patience`).
- Customizable `early_stopping.threshold` is introduced to adjust the improvement threshold for early stopping.

### Notable Dependencies or Requirements

* [`TrainingArguments`] argument `load_best_model_at_end`, `metric_for_best_model`, and `eval_strategy`.
* External dependency on [Logging] library (logger) for warning messages.
* Uses NumPy (`np`) for comparison operations.

### Additional Information

* The callback depends on the [`TrainerState`] to access the best metric value during training.
* If different save steps and evaluation steps are set in the `TrainingArguments`, early stopping will occur only after the next scheduled save step.


key: chroma:document
string_value:
CODE:

def import_protobuf(error_message=""):
    if is_sentencepiece_available():
        from sentencepiece import sentencepiece_model_pb2

        return sentencepiece_model_pb2
    if is_protobuf_available():
        import google.protobuf

        if version.parse(google.protobuf.__version__) < version.parse("4.0.0"):
            from transformers.utils import sentencepiece_model_pb2
        else:
            from transformers.utils import sentencepiece_model_pb2_new as sentencepiece_model_pb2
        return sentencepiece_model_pb2
    else:


SUMMARY:
**Summary**

Imports the Protocol Buffers library for natural language processing tasks.

### Main Purpose and Functionality

The `import_protobuf` function imports the necessary libraries based on availability, with specific versions or types of sentencepiece models being used. It returns the imported Protocol Buffers module.

### Key Implementation Details

* Checks for availability of `sentencepiece` and `google.protobuf` libraries.
* Imports `sentencepiece_model_pb2` based on version compatibility:
	+ Old version: `sentencepiece_model_pb2`
	+ Newer version (>= 4.0.0): `sentencepiece_model_pb2_new` with alias `sentencepiece_model_pb2`

### Notable Dependencies or Requirements

* `is_sentencepiece_available()` and `is_protobuf_available()` functions.
* `google.protobuf.__version__`.
* `transformers.utils import sentencepiece_model_pb2_new as sentencepiece_model_pb2`.


key: chroma:document
string_value:
CODE:

def _get_prepend_scheme(add_prefix_space: bool, original_tokenizer) -> str:
    if add_prefix_space:
        prepend_scheme = "always"
        if not getattr(original_tokenizer, "legacy", True):
            prepend_scheme = "first"
    else:
        prepend_scheme = "never"


SUMMARY:
**Summary:**

`_get_prepend_scheme` function determines the scheme for prepending spaces to tokens based on `add_prefix_space` flag and `original_tokenizer`.

**Details:**

* Returns a string indicating whether to always (`"always"`), first only when not in legacy mode (`"first"`), or never (`"never"`).
* Usage depends on `add_prefix_space` flag:
  * If true, scheme is `"always"`.
  * If false, scheme is `"never"`.
* **Note:** The function inspect the presence of a `legacy` attribute on the `original_tokenizer` object to adjust the scheme accordingly.

**Implementation Details:**

* `_get_prepend_scheme` takes two parameters:
  + `add_prefix_space` (bool): Flag indicating whether to add prefix space.
  + `original_tokenizer`: An object with methods/attributes used for tokenizer configuration.


key: chroma:document
string_value:
CODE:

def generate_merges(vocab, vocab_scores):
    reverse = vocab_scores is not None
    vocab_scores = dict(vocab_scores) if reverse else vocab

    merges = []
    for merge, piece_score in vocab_scores.items():
        local = []
        for index in range(1, len(merge)):
            piece_l, piece_r = merge[:index], merge[index:]
            if piece_l in vocab and piece_r in vocab:
                local.append((piece_l, piece_r, piece_score))
        local = sorted(local, key=lambda x: (vocab[x[0]], vocab[x[1]]))
        merges.extend(local)

    merges = sorted(merges, key=lambda val: (val[2], len(val[0]), len(val[1])), reverse=reverse)
    merges = [(val[0], val[1]) for val in merges]


SUMMARY:
**Function Summary**

`generate_merges(vocab, vocab_scores)` generates lists of merge operations between vocabulary words based on their scores.

**Primary Purpose**
The function takes two inputs:
- `vocab`: the vocabulary dictionary
- `vocab_scores`: optionally, an alternative vocabulary with associated scores

It returns a list of **merge operations**, where each operation consists of two **words** and their corresponding **score**.

**Key Implementation Details**

1.  Iterates through words in `vocab_scores` to find matching pairs: for each word, checks its left (`piece_l`) and right (`piece_r`) half-sequences against the vocabulary.
2.  Sorts matches by combined score (word values), followed by sequence length.
3.  Merges sorted matches into a single list.

**Notable Dependencies/Requirements**

*   `dict` for converting `vocab_scores` to a dictionary if necessary
*   `sorted` function for sorting merge operations

**External Dependencies**

*   `None` or an alternative vocabulary (`vocab`)


key: chroma:document
string_value:
CODE:

class SentencePieceExtractor:
    """
    Extractor implementation for SentencePiece trained models. https://github.com/google/sentencepiece
    """

    def __init__(self, model: str):
        requires_backends(self, "sentencepiece")
        from sentencepiece import SentencePieceProcessor

        self.sp = SentencePieceProcessor()
        self.sp.Load(model)

    def extract(self, vocab_scores=None) -> Tuple[Dict[str, int], List[Tuple]]:
        """
        By default will return vocab and merges with respect to their order, by sending `vocab_scores` we're going to
        order the merges with respect to the piece scores instead.
        """
        sp = self.sp
        vocab = {sp.id_to_piece(index): index for index in range(sp.GetPieceSize())}

        merges = generate_merges(vocab, vocab_scores)



SUMMARY:
**SentencePieceExtractor Summary**

### Purpose and Functionality
The `SentencePieceExtractor` class is an implementation of the SentencePiece algorithm for extracting vocabulary and piece merges from a trained model.

### Key Implementation Details

* The extractor loads a SentencePiece processor model using the provided model string (`self.sp.Load(model)`).
* It provides two primary methods:
	+ `extract(vocab_scores=None)`: Extracts vocabulary and piece merges. By default, it returns merges in order of their indices; with `vocab_scores` provided, it orders the merges by piece scores.
	+ Internal method: `generate_merges(vocab, vocab_scores)` (not shown in this snippet)

### Notable Dependencies and Requirements

* External dependency on `sentencepiece`
* Inheriting or implementing interfaces not specified (suggests a possible inheritance relation)


key: chroma:document
string_value:
CODE:

class GemmaSentencePieceExtractor(SentencePieceExtractor):
    def extract(self, vocab_scores=None) -> Tuple[Dict[str, int], List[Tuple]]:
        """
        By default will return vocab and merges with respect to their order, by sending `vocab_scores` we're going to
        order the merges with respect to the piece scores instead.
        """
        sp = self.sp
        vocab = {sp.id_to_piece(index): index for index in range(sp.GetPieceSize())}

        # there is a missing token in the vocab. We have to do this to support merges
        # "<0x09>" is the bytefallback for `\t`
        vocab["\t"] = vocab.get("<0x09>")

        merges = generate_merges(vocab, vocab_scores)


SUMMARY:
**GemmaSentencePieceExtractor Class Summary**

**Class Overview**
---------------

A custom implementation of `SentencePieceExtractor` that allows for ordering of vocabulary merges based on piece scores.

**Key Functionality**
------------------

1. Generates a dictionary `vocab` containing mapping from token IDs to their corresponding pieces.
2. Returns the `vocab` and merged tokens (`merges`) in case of `vocab_scores`.
3. Fallbacks to use `\t` (tab character) instead if a specific byte is missing.

**Notable Details**
-----------------

* **generate_merges**: an external function that takes `vocab` and `vocab_scores` as input.
* The class uses `sp(id_to_piece(index))` method from the parent class (`SentencePieceExtractor`) to fetch piece mapping.

**Dependancies/Requirements**
----------------------------

* `SentencePieceExtractor` interface
* External dependency: `generate_merges` function.


key: chroma:document
string_value:
CODE:

def check_number_comma(piece: str) -> bool:


SUMMARY:
**Code Summary**

### Functionality

The `check_number_comma` function checks if a given string (`piece`) meets specific conditions related to numbers and commas. It returns `True` if the string does not meet any of the following conditions:

* The length of the string is less than 2.
* The last character of the string is not a comma (`,`).
* The second-to-last character of the string is not a digit (`\d`).

### Code Details

- The function uses Python's string indexing and slicing features to access characters at specific positions.
- The `len()` function checks the length of the input string.

### Implementation Notes

- This function assumes that the input string represents a number or does not represent a complete number if it includes a comma but lacks the preceding digit.

```Markdown
# check_number_comma(piece: str) -> bool
# Returns True if piece is invalid according to stated conditions, False otherwise.
```

### Example Pseudocode

```markdown
# Conditions for return boolean value 'valid'
- length of slice(1..last) < 2 OR last == ',' 
     AND not digit(slice(last-1, last))
    OR any string not equal to comma slice(last)
```
Please let me know if you'd like further information or refactoring based on provided data.


key: chroma:document
string_value:
CODE:

class Converter:
    def __init__(self, original_tokenizer):
        self.original_tokenizer = original_tokenizer

    def converted(self) -> Tokenizer:


SUMMARY:
**Converter Class Summary**

The `Converter` class serves as an abstract interface to convert input data using various tokenizers. It has the following main purposes and functionality:

*   Initializes with an `original_tokenizer` parameter.
*   Defines an abstract method `converted` that must be implemented by subclasses, requiring a return type of `Tokenizer`.

**Key Implementation Details**

The primary design choice here is the use of an abstract base class (abc) to enforce concrete implementation subclasses.

**Notable Dependencies and Requirements**

*   A concrete subclass must provide its own implementation for the `converted` method.
*   The `original_tokenizer` parameter is inherited by any concrete subclasses.


key: chroma:document
string_value:
CODE:

class BertConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.vocab
        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))

        tokenize_chinese_chars = False
        strip_accents = False
        do_lower_case = False
        if hasattr(self.original_tokenizer, "basic_tokenizer"):
            tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars
            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case

        tokenizer.normalizer = normalizers.BertNormalizer(
            clean_text=True,
            handle_chinese_chars=tokenize_chinese_chars,
            strip_accents=strip_accents,
            lowercase=do_lower_case,
        )
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls}:0 $A:0 {sep}:0",
            pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1",
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
            ],
        )
        tokenizer.decoder = decoders.WordPiece(prefix="##")



SUMMARY:
**BertConverter Class Summary**

### Main Purpose and Functionality
The `BertConverter` class is a custom converter that inherits from the `Converter` class. It converts input data into a BERT tokenizer, preserving vocabulary, casing, accents, and punctuation.

### Key Implementation Details

*   Initializes a `Tokenizer` object with customized settings:
    *   Vocabulary and unknown token
    *   Normalizer setting to use `BertNormalizer`
    *   Pre-tokenizer to use `BertPreTokenizer`

### Notable Dependencies or Requirements

*   References:
    *   `Tokenizer` class from Hugging Face Transformers library
    *   `normalizers.BertNormalizer` and `pre_tokenizers.BertPreTokenizer` classes
    *   `processors.TemplateProcessing` class
    *   `decoders.WordPiece` class
*   External dependencies:
    + Hugging Face Transformers library (`transformers`)
    + NLTK library (`nltk`) for text normalization

### Interface Implementations

Inherits from the `Converter` interface and returns an instance of `Tokenizer`.


key: chroma:document
string_value:
CODE:

class SplinterConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.vocab
        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))

        tokenize_chinese_chars = False
        strip_accents = False
        do_lower_case = False
        if hasattr(self.original_tokenizer, "basic_tokenizer"):
            tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars
            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case

        tokenizer.normalizer = normalizers.BertNormalizer(
            clean_text=True,
            handle_chinese_chars=tokenize_chinese_chars,
            strip_accents=strip_accents,
            lowercase=do_lower_case,
        )
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        question = str(self.original_tokenizer.question_token)
        dot = "."
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id
        question_token_id = self.original_tokenizer.question_token_id
        dot_token_id = self.original_tokenizer.convert_tokens_to_ids(".")

        if self.original_tokenizer.padding_side == "right":
            pair = f"{cls}:0 $A:0 {question} {dot} {sep}:0 $B:1 {sep}:1"
        else:
            pair = f"{cls}:0 $A:0 {sep}:0 $B:1 {question} {dot} {sep}:1"

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls}:0 $A:0 {sep}:0",
            pair=pair,
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
                (question, question_token_id),
                (dot, dot_token_id),
            ],
        )
        tokenizer.decoder = decoders.WordPiece(prefix="##")



SUMMARY:
**Summary:**

The `SplinterConverter` class is designed to create a custom tokenizer instance for the Hugging Face Transformers library. It inherits from the `Converter` class and extends its functionality.

**Primary Purpose:**
The primary purpose of this class is to convert the original tokenizer to a custom tokenizer with specific configuration settings, such as Unicode handling, accent stripping, and case normalization.

**Key Implementation Details:**

* The tokenizer uses a WordPiece approach for tokenization.
* Custom normalizer, pre-tokenizer, post-processor, and decoder are configured.
* Template processing is used to create specific token prefixes and suffixes.

**Notable Dependencies and Requirements:**

* `Tokenizer`, `WordPiece`, `BertNormalizer`, `pre_tokenizers`, `processors`, `decoders` classes from the Hugging Face Transformers library.
* No additional dependencies are listed.

**Interface Methods and Interfaces:**

* Returns a custom `Tokenizer` instance.
* Implements the following interfaces:
	+ `Converter`
	+ `Tokenizer`

**Design Choices:**
The design choice of using WordPiece tokenization, and custom post-processor with template processing is notable.


key: chroma:document
string_value:
CODE:

class FunnelConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.vocab
        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))

        tokenize_chinese_chars = False
        strip_accents = False
        do_lower_case = False
        if hasattr(self.original_tokenizer, "basic_tokenizer"):
            tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars
            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case

        tokenizer.normalizer = normalizers.BertNormalizer(
            clean_text=True,
            handle_chinese_chars=tokenize_chinese_chars,
            strip_accents=strip_accents,
            lowercase=do_lower_case,
        )
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls}:2 $A:0 {sep}:0",  # token_type_id is 2 for Funnel transformer
            pair=f"{cls}:2 $A:0 {sep}:0 $B:1 {sep}:1",
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
            ],
        )
        tokenizer.decoder = decoders.WordPiece(prefix="##")



SUMMARY:
**FunnelConverter Class Summary**

The `FunnelConverter` class extends `Converter`. Its primary purpose is to create and customize a `Tokenizer` instance for converting text data.

**Key Implementation Details:**

* Uses the `WordPiece` algorithm from BERT tokenizer.
* Customizes tokenization settings, including Chinese characters handling and accent stripping.
* Configures normalization, preprocessing, and post-processing pipelines using BERT-specific preprocessors and postprocessors.
* Defines special tokens with custom IDs.

**Notable Dependencies and Requirements:**

* `Tokenizer` class with `WordPiece`, `pre_tokenizers.BertPreTokenizer`, and `processors.TemplateProcessing`.
* `normalizers.BertNormalizer` and `decoders.WordPiece` classes from the Hugging Face library.
* External dependencies:
	+ `huggingface_hub` library (for importing pre-trained models).
	+ Additional Python libraries for tokenization, text processing, and natural language text analysis.


key: chroma:document
string_value:
CODE:

class MPNetConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.vocab
        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))

        tokenize_chinese_chars = False
        strip_accents = False
        do_lower_case = False
        if hasattr(self.original_tokenizer, "basic_tokenizer"):
            tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars
            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case

        tokenizer.normalizer = normalizers.BertNormalizer(
            clean_text=True,
            handle_chinese_chars=tokenize_chinese_chars,
            strip_accents=strip_accents,
            lowercase=do_lower_case,
        )
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls}:0 $A:0 {sep}:0",
            pair=f"{cls}:0 $A:0 {sep}:0 {sep}:0 $B:1 {sep}:1",  # MPNet uses two [SEP] tokens
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
            ],
        )
        tokenizer.decoder = decoders.WordPiece(prefix="##")



SUMMARY:
**MPNetConverter Class Summary**

The `MPNetConverter` class extends the `Converter` class and is responsible for converting text to a format suitable for the MPNet model.

**Main Purpose and Functionality**

The primary purpose of this class is to create a `Tokenizer` object that converts input text into a format compatible with MPNet using WordPiece tokenization. The tokenizer also applies normalizer, pre- and post-processing steps to transform the input text.

**Key Implementation Details**

* Tokenization parameters are inherited from the original tokenizer (if available) or default values.
* Normalizer: `BertNormalizer` is used with specific settings for handling Chinese characters, accents, and case conversion.
* Pre-tokenizer: `BertPreTokenizer` is used to preprocess input text.
* Post-processor: `TemplateProcessing` is implemented to apply template processing rules for MPNet model inputs.

**Notable Dependencies and Requirements**

* `original_tokenizer`: an instance of the original tokenizer from which tokenization parameters are inherited or retrieved.
	+ `WordPieceTokenizer`
	+ `BertNormalizer`
	+ `BertPreTokenizer`
	+ `TemplateProcessing` processor
* External dependencies:
	+ `decoders.WordPiece`

**Notes**

The class preserves default settings for some parameters and applies custom settings from the original tokenizer, if available.


key: chroma:document
string_value:
CODE:

class OpenAIGPTConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.encoder
        merges = list(self.original_tokenizer.bpe_ranks.keys())
        unk_token = self.original_tokenizer.unk_token

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                unk_token=str(unk_token),
                end_of_word_suffix="</w>",
                fuse_unk=False,
            )
        )

        if tokenizer.token_to_id(str(unk_token)) is not None:
            tokenizer.add_special_tokens([str(unk_token)])

        tokenizer.normalizer = normalizers.BertNormalizer(lowercase=True)
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
        tokenizer.decoder = decoders.BPEDecoder(suffix="</w>")



SUMMARY:
**OpenAIGPTConverter Class Summary**

The `OpenAIGPTConverter` class is designed to convert an original GPT-based tokenizer into an OpenAIGPT-specific variant. The primary purpose of this converter is to update the tokenization settings to comply with the OpenAIGPT model's requirements.

**Key Functionality:**

- Updates BPE ranks and merges according to OpenAIGPT specifications.
- Adds special tokens if the unknown token is part of the vocabulary.

**Implementation Details:**

- `converted` method:
  - Creates a new instance of `Tokenizer` with customized BPE settings.
  - Initializes the tokenizer with specific configuration, including the unk_token.
- Updates the normalizer, pre-tokenizer, and decoder to match OpenAIGPT requirements.

**Notable Dependencies/Requirements:**

- Original GPT-based tokenizer instance (`self.original_tokenizer`).
- External dependencies:
  + `normalizers.BertNormalizer`
  + `pre_tokenizers.BertPreTokenizer`
  + `decoders.BPEDecoder`

This class appears to be part of a larger framework for converting and adapting pre-trained models, specifically GPT-based tokenizers, to different variants like OpenAIGPT.


key: chroma:document
string_value:
CODE:

class GPT2Converter(Converter):
    def converted(self, vocab: Dict[str, int] = None, merges: List[Tuple[str, str]] = None) -> Tokenizer:
        if not vocab:
            vocab = self.original_tokenizer.encoder
        if not merges:
            merges = list(self.original_tokenizer.bpe_ranks)

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
            )
        )

        add_prefix_space = getattr(self.original_tokenizer, "add_prefix_space", False)
        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space)
        tokenizer.decoder = decoders.ByteLevel()
        if getattr(self.original_tokenizer, "add_bos_token", False):
            bos = self.original_tokenizer.bos_token
            bos_token_id = self.original_tokenizer.bos_token_id
            tokenizer.post_processor = processors.TemplateProcessing(
                single=f"{bos}:0 $A:0",
                pair=f"{bos}:0 $A:0 $B:1",
                special_tokens=[
                    (bos, bos_token_id),
                ],
            )
        else:
            # XXX trim_offsets=False actually means this post_processor doesn't
            # really do anything.
            tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)


SUMMARY:
**Class Summary: GPT2Converter**

The `GPT2Converter` class is a subclass of `Converter`. Its primary purpose is to convert the original tokenization configuration to a compatible format for GPT-2 model.

**Main Functionality:**

* Initializes a new `Tokenizer` instance with a customized BPE (Byte Pair Encoding) pipeline.
* Adjusts the preprocessing and post-processing settings according to the original tokenizer's configuration.

**Key Implementation Details:**

* Custom BPE pipeline with optimized vocabulary and merge rules.
* Pre-tokenizer integration using ByteLevel with prefix space addition enabled (if available).
* Post-processor configuration for handling BOS tokens, if present.

**Notable Dependencies/Requirements:**

* `self.original_tokenizer`
	+ `encoder`: the original vocabulary dictionary.
	+ `bpe_ranks`: the merge rules.
	+ `add_prefix_space`, `add_bos_token`, and related attributes.
	+ `pre_tokenizers` (specifically, `ByteLevel`)
	+ `decoders` (specifically, `ByteLevel`)
	+ `processors` (specifically, `TemplateProcessing` and `ByteLevel` with `trim_offsets=False`)


key: chroma:document
string_value:
CODE:

class HerbertConverter(Converter):
    def converted(self) -> Tokenizer:
        tokenizer_info_str = "#version:"
        token_suffix = "</w>"

        vocab = self.original_tokenizer.encoder
        merges = list(self.original_tokenizer.bpe_ranks.keys())
        if tokenizer_info_str in merges[0][0]:
            merges = merges[1:]

        tokenizer = Tokenizer(
            BPE(
                vocab,
                merges,
                dropout=None,
                unk_token=self.original_tokenizer.unk_token,
                end_of_word_suffix=token_suffix,
            )
        )

        tokenizer.normalizer = normalizers.BertNormalizer(lowercase=False, strip_accents=False)
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
        tokenizer.decoder = decoders.BPEDecoder(suffix=token_suffix)
        tokenizer.post_processor = processors.BertProcessing(
            sep=(self.original_tokenizer.sep_token, self.original_tokenizer.sep_token_id),
            cls=(self.original_tokenizer.cls_token, self.original_tokenizer.cls_token_id),
        )



SUMMARY:
**HerbertConverter Class Summary**

This class is an implementation of the `Converters` interface, specifically designed to convert between tokenizers. The primary purpose is to create and configure a new `Tokenizer` instance.

**Key Implementation Details:**

* Creates a new `BPE` tokenizer with a vocabulary derived from the original tokenizer's encoder.
* Configures the tokenizer's normalizer, pre(tokenizer), decoder, and postprocessor components using specific settings (e.g., lowercase/strip_accents=False).

**Notable Dependencies and Requirements:**

* The `original_tokenizer` object is used to configure the new tokenizer and its various components.
* Requires external dependencies from the following modules:
	+ `processors.py`: `BertProcessing`
	+ `pre_tokenizers.py`: `BertPreTokenizer`
	+ `normalizers.py`: `BertNormalizer`
	+ `decoders.py`: `BPEDecoder`

**Method:** The class defines a single method, `converted`, which returns an instance of the customized `Tokenizer`.


key: chroma:document
string_value:
CODE:

class Qwen2Converter(Converter):
    def converted(self, vocab: Dict[str, int] = None, merges: List[Tuple[str, str]] = None) -> Tokenizer:
        if not vocab:
            vocab = self.original_tokenizer.encoder
        if not merges:
            merges = list(self.original_tokenizer.bpe_ranks.keys())

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                unk_token=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
                byte_fallback=False,
            )
        )

        tokenizer.normalizer = normalizers.NFC()

        tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
            [
                pre_tokenizers.Split(
                    Regex(
                        r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
                    ),
                    behavior="isolated",
                    invert=False,
                ),
                pre_tokenizers.ByteLevel(
                    add_prefix_space=getattr(self.original_tokenizer, "add_prefix_space", False),
                    use_regex=False,
                ),
            ]
        )

        tokenizer.decoder = decoders.ByteLevel()
        tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)



SUMMARY:
**Summary**

The `Qwen2Converter` class extends the `Converter` base class and provides a method `converted` that returns a customized instance of the `Tokenizer` class. The custom Tokenizer creation involves:

*   Initializing a BPE (Byte Pair Encoding) component with a given vocabulary (`vocab`) and merging rules (`merges`)
*   Configuring normalization, preprocessing, and post-processing for the tokenizer
*   Returning the fully configured Tokenizer instance

**Key Implementation Details**

*   Customized `Tokenizer` creation with BPE configuration
*   Normalizer: NFC (Normalization Function Composition)
*   Pre-tokenizer: Sequence of Split and ByteLevel operations
*   Decoder: ByteLevel
*   Post-processor: ByteLevel (with trimming disabled)

**Notable Dependencies or Requirements**

*   External dependencies:
    - `pre_tokenizers`
    - `normalizers`
    - `processors`
    - `decoders`
    - `Tokenizer` and `BPE` classes from the original tokenizer
*   Inherited interfaces:
    - `Converter` base class


key: chroma:document
string_value:
CODE:

class RobertaConverter(Converter):
    def converted(self) -> Tokenizer:
        ot = self.original_tokenizer
        vocab = ot.encoder
        merges = list(ot.bpe_ranks.keys())

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
            )
        )

        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space)
        tokenizer.decoder = decoders.ByteLevel()
        tokenizer.post_processor = processors.RobertaProcessing(
            sep=(ot.sep_token, ot.sep_token_id),
            cls=(ot.cls_token, ot.cls_token_id),
            add_prefix_space=ot.add_prefix_space,
            trim_offsets=True,  # True by default on Roberta (historical)
        )



SUMMARY:
**Concise Summary:**

`RobertaConverter` class is a customized converter for the `Tokenizer` class in Hugging Face's transformers library. 

**Key Functionality and Implementation Details:**

* Creates a new `Tokenizer` instance with a `BPE` (Byte Pair Encoding) module.
* Configures `BPE` with the original vocabulary, BPE ranks, no dropout, and default settings for continuing subword prefix and suffix.
* Sets pre-tokenizer, decoder, and post-processor components:
  - Pre-tokenizer: ByteLevel with added prefix space from the original tokenizer.
  - Decoder: ByteLevel.
  - Post-processor: RobertaProcessing with specific separator settings.

**Notable Dependencies and Requirements:**

* `Tokenizer`: Base class for tokenization.
* `BPE`: Bitextual encoding algorithm used in `BPE` module.
* `pre_tokenizers` and `decoders`: Pre-tokenizer and decoder components, respectively.
* `processors`: Post-processor components, specifically `RobertaProcessing`.
* Related entities `ot` (original tokenizer object) have access to its attributes, such as `encoder`, `add_prefix_space`, `bpe_ranks`, `sep_token`, and `sep_token_id`.


key: chroma:document
string_value:
CODE:

class RoFormerConverter(Converter):
    def converted(self) -> Tokenizer:
        from .models.roformer.tokenization_utils import JiebaPreTokenizer

        vocab = self.original_tokenizer.vocab
        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))

        strip_accents = False
        do_lower_case = False
        if hasattr(self.original_tokenizer, "basic_tokenizer"):
            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case

        tokenizer.normalizer = normalizers.BertNormalizer(
            clean_text=True,
            handle_chinese_chars=False,
            strip_accents=strip_accents,
            lowercase=do_lower_case,
        )
        tokenizer.pre_tokenizer = pre_tokenizers.PreTokenizer.custom(JiebaPreTokenizer(vocab))

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls}:0 $A:0 {sep}:0",
            pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1",
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
            ],
        )
        tokenizer.decoder = decoders.WordPiece(prefix="##")



SUMMARY:
**RoFormerConverter Class Summary**

### Main Purpose and Functionality
The `RoFormerConverter` class is a custom converter for the RoFormer model that wraps the existing `Tokenizer` from the `WordPiece` library. It applies preprocessing and normalization techniques, including tokenization, stemming, and text encoding.

### Key Implementation Details

* Initializes a new instance of the `Tokenizer` with customized settings: 
	+ Uses the `WordPiece` vocabulary.
	+ Apples the original tokenizer's `strip_accents` and `do_lower_case` settings.
	+ Enforces strict Chinese character handling.
	+ Employs the `BertNormalizer` for text cleaning and normalization.
* Sets up custom pre-tokenization using `JiebaPreTokenizer`.
* Configures the post-processing template to include:
	+ Separator tokens with specific IDs.
	+ Special token pairs.

### Notable Dependencies and Requirements

* External dependencies:
	+ `.models.roformer.tokenization_utils`: contains the `JiebaPreTokenizer` class.
	+ `models.wordpiece Tokenizer`, `processors.TemplateProcessing`, and `decoders.WordPiece`.
	+ The `normalizers.BertNormalizer` and `pre_tokenizers.PreTokenizer`.
* Inherited or implemented interfaces:
	+ None explicitly mentioned, but may derive from the `Converter` base class.


key: chroma:document
string_value:
CODE:

class DebertaConverter(Converter):
    def converted(self) -> Tokenizer:
        ot = self.original_tokenizer
        vocab = ot.encoder
        merges = list(ot.bpe_ranks.keys())

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
            )
        )

        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space)
        tokenizer.decoder = decoders.ByteLevel()
        tokenizer.post_processor = processors.TemplateProcessing(
            single="[CLS]:0 $A:0 [SEP]:0",
            pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1",
            special_tokens=[
                ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")),
                ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")),
            ],
        )



SUMMARY:
**DebertaConverter Class Summary**

The `DebertaConverter` class is a subclass of `Converter`, responsible for converting the input data to a Deberta-style tokenized format.

**Primary Purpose:**
The primary purpose of this class is to create and return a `Tokenizer` object with Deberta-specific configurations, including BPE merging, prefix space addition, Byte-Level preprocessing, and Template Processing with special tokens.

**Key Implementation Details:**

* Creates a `Tokenizer` instance with a BPE tokenizer, utilizing the original tokenizer's vocabulary and merge ranks.
* Sets the pre-tokenizer, decoder, and post-processor for the tokenized data according to Deberta conventions.
* Defines the template processing configuration for converting inputs to Deberta tokenization format.

**Notable Dependencies:**

* `self.original_tokenizer`: The original tokenizer object used for inspiration.
* `pre_tokenizers[ByteLevel]`: A Byte-Level pre-processor.
* `decoders[ByteLevel]`: A Byte-Level decoder.
* `processors.TemplateProcessing`: A template processor for handling special tokens.


key: chroma:document
string_value:
CODE:

class SpmConverter(Converter):
    handle_byte_fallback = False
    SpmExtractor = SentencePieceExtractor
    special_tokens = {}

    def __init__(self, *args):
        requires_backends(self, "protobuf")

        super().__init__(*args)

        # from .utils import sentencepiece_model_pb2 as model_pb2
        model_pb2 = import_protobuf()

        m = model_pb2.ModelProto()
        with open(self.original_tokenizer.vocab_file, "rb") as f:
            m.ParseFromString(f.read())
        self.proto = m

        if self.proto.trainer_spec.byte_fallback and not self.handle_byte_fallback:
            warnings.warn(
                "The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option"
                " which is not implemented in the fast tokenizers. In practice this means that the fast version of the"
                " tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these "
                "unknown tokens into a sequence of byte tokens matching the original piece of text."
            )

    def vocab(self, proto):
        return [(piece.piece, piece.score) for piece in proto.pieces]

    def unk_id(self, proto):
        return proto.trainer_spec.unk_id

    def tokenizer(self, proto):
        model_type = proto.trainer_spec.model_type
        vocab_scores = self.vocab(proto)

        if model_type == 1:
            tokenizer = Tokenizer(
                Unigram(
                    vocab_scores,
                    unk_id=self.unk_id(proto),
                    byte_fallback=self.handle_byte_fallback,
                )
            )

        elif model_type == 2:
            _, merges = self.SpmExtractor(self.original_tokenizer.vocab_file).extract(vocab_scores)
            bpe_vocab = {word: i for i, (word, score) in enumerate(vocab_scores)}
            tokenizer = Tokenizer(
                BPE(
                    bpe_vocab,
                    merges,
                    unk_token=proto.trainer_spec.unk_piece,
                    fuse_unk=True,
                    byte_fallback=self.handle_byte_fallback,
                    dropout=None,
                )
            )

        else:
            raise Exception(
                "You're trying to run a `Unigram` model but you're file was trained with a different algorithm"
            )

        # control tokens are special
        # user defined symbols are not
        # both user and control tokens are AddedTokens
        # Add user defined symbols (type == 4) from sentencepiece (https://github.com/google/sentencepiece/blob/6225e08edb2577757163b3f5dbba4c0b670ef445/src/sentencepiece_model.proto#L299C29-L299C33)
        spm_added_tokens = [
            (id, p.piece, p.type == 3 or p.piece in self.special_tokens)
            for id, p in enumerate(proto.pieces)
            if p.type in [3, 4]
        ]
        tokenizer.add_tokens(
            [
                AddedToken(token, normalized=False, special=special)
                for id, token, special in sorted(spm_added_tokens, key=lambda x: x[0])
            ]
        )

        return tokenizer

    def normalizer(self, proto):
        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap
        _normalizers = [
            normalizers.Strip(left=False, right=True),  # stripping is important
            normalizers.Replace(Regex(" {2,}"), "▁"),
        ]
        if not precompiled_charsmap:
            return normalizers.Sequence(_normalizers)
        else:
            return normalizers.Sequence([normalizers.Precompiled(precompiled_charsmap)] + _normalizers)

    def pre_tokenizer(self, replacement, add_prefix_space):
        prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer)
        return pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme)

    def post_processor(self):
        return None

    def decoder(self, replacement, add_prefix_space):
        prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer)
        return decoders.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme)

    def converted(self) -> Tokenizer:
        tokenizer = self.tokenizer(self.proto)

        # Tokenizer assemble
        normalizer = self.normalizer(self.proto)
        if normalizer is not None:
            tokenizer.normalizer = normalizer

        replacement = "▁"
        add_prefix_space = True
        if hasattr(self.original_tokenizer, "add_prefix_space"):
            add_prefix_space = self.original_tokenizer.add_prefix_space

        pre_tokenizer = self.pre_tokenizer(replacement, add_prefix_space)
        if pre_tokenizer is not None:
            tokenizer.pre_tokenizer = pre_tokenizer

        tokenizer.decoder = self.decoder(replacement, add_prefix_space)
        post_processor = self.post_processor()
        if post_processor:
            tokenizer.post_processor = post_processor



SUMMARY:
**Concise Summary**

The `SpmConverter` class is designed to convert SentencePiece-tokenized models into fast tokenizer models. It inherits from the `Converter` class.

* **Purpose**: The converter updates the original tokenizer's vocabulary, normalizer, and post-processor to work with the new model.
* **Key Implementation**:
	+ The `tokenizer` method creates a new `Tokenizer` instance based on the model type (Unigram or BPE).
	+ The `normalizer` method creates a normalizer instance that strips and replaces special characters.
	+ The `pre_tokenizer` method creates a pre-tokenizer instance to prepend scheme and replacement.
	+ The `converted` method assembles the tokenizer by integrating the normalizer, pre-processor, and post-processor.

**Key Properties and Dependencies**

* `handle_byte_fallback`: Flag indicating whether byte fallback is enabled.
* `SpamExtractor`: Imported from `sentencepieceModel_pb2`.
* `model_pb2`: A message class from a protobuf file.
* `_get_prepend_scheme`, `_normalizers`, `pre_tokenizers`, `decoders`: Helper functions and classes not defined in this code snippet.

**External Dependencies**

* `sentencepiece` library for SentencePiece extraction
* `protobuf` library for accessing the `model_pb2` message class


key: chroma:document
string_value:
CODE:

class AlbertConverter(SpmConverter):
    def vocab(self, proto):
        return [
            (piece.piece, piece.score) if check_number_comma(piece.piece) else (piece.piece, piece.score - 100)
            for piece in proto.pieces
        ]

    def normalizer(self, proto):
        list_normalizers = [
            normalizers.Replace("``", '"'),
            normalizers.Replace("''", '"'),
        ]
        if not self.original_tokenizer.keep_accents:
            list_normalizers.append(normalizers.NFKD())
            list_normalizers.append(normalizers.StripAccents())
        if self.original_tokenizer.do_lower_case:
            list_normalizers.append(normalizers.Lowercase())

        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap

        if precompiled_charsmap:
            list_normalizers.append(normalizers.Precompiled(precompiled_charsmap))

        list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " "))
        return normalizers.Sequence(list_normalizers)

    def post_processor(self):
        return processors.TemplateProcessing(
            single="[CLS]:0 $A:0 [SEP]:0",
            pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1",
            special_tokens=[
                ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")),
                ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")),
            ],


SUMMARY:
**Class Summary: AlbertConverter**

The `AlbertConverter` class, a subclass of `SpmConverter`, is a custom converter for the BERT-based Albert model in Transformers library.

**Main Purpose and Functionality**
The primary purpose of this converter is to preprocess input text data into a suitable format for theBERT-based models. It implements several normalization steps:

1. **Tokenization**: Replacements with double quotes, stripping accents, and applying lower-case conversion.
2. **Character Mapping**: Precompiled character maps using `normalizers.Precompiled`.
3. **Whitespace Normalization**: Replacing consecutive whitespace characters with a single space.

**Key Implementation Details**

* A sequential pipeline of normalization steps is applied before tokenization.
* Template-based token mapping for special tokens `[CLS]` and `[SEP]`.

**Notable Dependencies or Requirements**

* `normalizers`: Module used for character mappings and replacements.
	+ `RemoveCharmap`, `Replace`, `NFKD`, `StripAccents`, `Lowercase`, and `Precompiled`
* `processors`: Module used for template-based token processing.
	+ `TemplateProcessing` class.
* `Transformers` library for the BERT-based Albert model.


key: chroma:document
string_value:
CODE:

class BarthezConverter(SpmConverter):
    def unk_id(self, proto):
        unk_id = 3
        return unk_id

    def post_processor(self):
        return processors.TemplateProcessing(
            single="<s> $A </s>",
            pair="<s> $A </s> </s> $B </s>",
            special_tokens=[
                ("<s>", self.original_tokenizer.convert_tokens_to_ids("<s>")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Summary: Class BarthezConverter Converter Class**

The `BarthezConverter` class extends the `SpmConverter` class and is part of a text processing pipeline. It mainly serves a post-processing function, utilizing the `processors.TemplateProcessing` class.

**Key Functionality:**

*   Converts unknown IDs to a fixed value (3) during processing.
*   Applies template-based tokenization for single, pair, and special tokens using the `TemplateProcessing` class.

**Notable Methods/ Interfaces:**

*   `post_processor`: Returns an instance of `processors.TemplateProcessing`, defining template patterns for text processing, including single, pair, and special token conversions.

**Dependencies/Requirements:**

*   Extends `SpmConverter`
*   Reuses `original_tokenizer` instance.
*   Utilizes external dependencies from the `processorTemplateProcessing` module.


key: chroma:document
string_value:
CODE:

class CamembertConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<s>NOTUSED", 0.0),
            ("<pad>", 0.0),
            ("</s>NOTUSED", 0.0),
            ("<unk>", 0.0),
            ("<unk>NOTUSED", -100),
        ]
        # We down-grade the original SentencePiece by -100 to avoid using it and use our added token instead
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[1:]]
        vocab += [("<mask>", 0.0)]
        return vocab

    def unk_id(self, proto):
        # See vocab unk position
        return 3

    def post_processor(self):
        return processors.TemplateProcessing(
            single="<s> $A </s>",
            pair="<s> $A </s> </s> $B </s>",
            special_tokens=[
                ("<s>", self.original_tokenizer.convert_tokens_to_ids("<s>")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Concise Summary**

The `CamembertConverter` class extends the `SpmConverter` class and provides converters for a specific use case. The class includes:

*   A vocabulary (`vocab`) function that:
    *   Generates an initial list of word pieces with their corresponding scores.
    *   Downgrades a sentence piece to a custom token when it's NOTUSED.
    *   Includes additional tokens, such as `<unk>`, `<mask>`, and special tokens for start and end sequences.

**Key Implementation Details**

*   `vocab`: adds original sentence piece tokens by converting them to IDs with the `self.original_tokenizer`.
*   `unk_id`: maps the NOTUSED token position to 3 (index 2 in vocabulary).

**Notable Dependencies or Requirements**

*   Uses `SpmConverter` as a base class.
*   Imports `processors.TemplateProcessing` for post-processing.


key: chroma:document
string_value:
CODE:

class DebertaV2Converter(SpmConverter):
    def pre_tokenizer(self, replacement, add_prefix_space):
        list_pretokenizers = []
        if self.original_tokenizer.split_by_punct:
            list_pretokenizers.append(pre_tokenizers.Punctuation(behavior="isolated"))
        prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer)
        list_pretokenizers.append(pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme))
        return pre_tokenizers.Sequence(list_pretokenizers)

    def normalizer(self, proto):
        list_normalizers = []
        if self.original_tokenizer.do_lower_case:
            list_normalizers.append(normalizers.Lowercase())
        list_normalizers.append(normalizers.Strip())

        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap
        if precompiled_charsmap:
            list_normalizers.append(normalizers.Precompiled(precompiled_charsmap))
        list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " "))

        return normalizers.Sequence(list_normalizers)

    def post_processor(self):
        return processors.TemplateProcessing(
            single="[CLS]:0 $A:0 [SEP]:0",
            pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1",
            special_tokens=[
                ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")),
                ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")),
            ],


SUMMARY:
**DebertaV2Converter Class Summary**

The `DebertaV2Converter` class is a subclass of `SpmConverter`, implementing the pre-tokenization, normalization, and post-processing steps for converting sequences into processed inputs suitable for DebertaV2 architectures.

**Key Implementation Details:**

1. **Pre-Tokenization:** The `pre_tokenizer` method creates a sequence of pretokenizers, adding punctuation isolator only if enabled, followed by a metaspace with a prefixed scheme and replacement.
   
   ```python
def pre_tokenizer(self, replacement, add_prefix_space):
    # ...
    list_pretokenizers = [pre_tokenizers.Punctuation(behavior="isolated"), 
                          pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme)]
    return pre_tokenizers.Sequence(list_pretokenizers)
```

2. **Normalization:** The `normalizer` method constructs a sequence of normalizers:
   - Lowercased tokens if the original tokenizer is configured so.
   - Strip whitespace characters by default.
   - Optionally applies precompiled character mappings from the tokenization spec and performs whitespace replacement.

   ```python
def normalizer(self, proto):
    list_normalizers = []
    # ...
    list_normalizers.append(normalizers.Precompile(precompiled_charsmap=proto.normalizer_spec.precompiled_charsmap))
    list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " "))
    return normalizers.Sequence(list_normalizers)
```

3. **Post-Processing:** The `post_processor` method returns a template processing processor with CLS, SEP, A, and B token replacements, along with special tokens mapped to their IDs.

   ```python
def post_processor(self):
    # ...
    processors.TemplateProcessing(
        single="[CLS]:0 $A:0 [SEP]:0",
        pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1",
        special_tokens=[
            ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")),
            ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")),
        ],
    )
```

**Notable Dependencies:**

- `pre_tokenizers`
- `normalizers`
- `processors`
- `_get_prepend_scheme` 
- `self.original_tokenizer` and its child classes (`pre_tokenizers`, `normalizers`, and `processors`)


key: chroma:document
string_value:
CODE:

class MBartConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<s>", 0.0),
            ("<pad>", 0.0),
            ("</s>", 0.0),
            ("<unk>", 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        vocab += [
            ("ar_AR", 0.0),
            ("cs_CZ", 0.0),
            ("de_DE", 0.0),
            ("en_XX", 0.0),
            ("es_XX", 0.0),
            ("et_EE", 0.0),
            ("fi_FI", 0.0),
            ("fr_XX", 0.0),
            ("gu_IN", 0.0),
            ("hi_IN", 0.0),
            ("it_IT", 0.0),
            ("ja_XX", 0.0),
            ("kk_KZ", 0.0),
            ("ko_KR", 0.0),
            ("lt_LT", 0.0),
            ("lv_LV", 0.0),
            ("my_MM", 0.0),
            ("ne_NP", 0.0),
            ("nl_XX", 0.0),
            ("ro_RO", 0.0),
            ("ru_RU", 0.0),
            ("si_LK", 0.0),
            ("tr_TR", 0.0),
            ("vi_VN", 0.0),
            ("zh_CN", 0.0),
        ]
        vocab += [("<mask>", 0.0)]
        return vocab

    def unk_id(self, proto):
        return 3

    def post_processor(self):
        return processors.TemplateProcessing(
            single="$A </s> en_XX",
            pair="$A $B </s> en_XX",
            special_tokens=[
                ("en_XX", self.original_tokenizer.convert_tokens_to_ids("en_XX")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Class Summary: MBartConverter**

*   **Main Purpose and Functionality**: This class is responsible for converting and preprocessing data for the Mixture of BART models. It extends the SpmConverter class to provide vocabulary, token identification, and post-processing functionality.

**Key Implementation Details**

*   The `vocab` method generates a list of vocabulary tokens with their corresponding scores, including common tokens like `<s>`, `<pad>`, and `<unk>`, as well as a list of language-specific tokens.
*   The `unk_id` method returns the ID for unknown tokens (3).
*   The `post_processor` method returns an object from the `processors.TemplateProcessing` class, which provides template-based post-processing functionality.

**Notable Dependencies or Requirements**

*   `SpmConverter`: The base class for MBartConverter.
*   `original_tokenizer`: Used to convert token IDs back to tokens in custom ways for special tokens ("en_XX" and "</s>").
*   `processors.TemplateProcessing`: A post-processing function with template-based processing capabilities.


key: chroma:document
string_value:
CODE:

class MBart50Converter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<s>", 0.0),
            ("<pad>", 0.0),
            ("</s>", 0.0),
            ("<unk>", 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        vocab += [("ar_AR", 0.0), ("cs_CZ", 0.0), ("de_DE", 0.0), ("en_XX", 0.0), ("es_XX", 0.0), ("et_EE", 0.0), ("fi_FI", 0.0), ("fr_XX", 0.0), ("gu_IN", 0.0), ("hi_IN", 0.0), ("it_IT", 0.0), ("ja_XX", 0.0), ("kk_KZ", 0.0), ("ko_KR", 0.0), ("lt_LT", 0.0), ("lv_LV", 0.0), ("my_MM", 0.0), ("ne_NP", 0.0), ("nl_XX", 0.0), ("ro_RO", 0.0), ("ru_RU", 0.0), ("si_LK", 0.0), ("tr_TR", 0.0), ("vi_VN", 0.0), ("zh_CN", 0.0), ("af_ZA", 0.0), ("az_AZ", 0.0), ("bn_IN", 0.0), ("fa_IR", 0.0), ("he_IL", 0.0), ("hr_HR", 0.0), ("id_ID", 0.0), ("ka_GE", 0.0), ("km_KH", 0.0), ("mk_MK", 0.0), ("ml_IN", 0.0), ("mn_MN", 0.0), ("mr_IN", 0.0), ("pl_PL", 0.0), ("ps_AF", 0.0), ("pt_XX", 0.0), ("sv_SE", 0.0), ("sw_KE", 0.0), ("ta_IN", 0.0), ("te_IN", 0.0), ("th_TH", 0.0), ("tl_XX", 0.0), ("uk_UA", 0.0), ("ur_PK", 0.0), ("xh_ZA", 0.0), ("gl_ES", 0.0), ("sl_SI", 0.0)]  # fmt: skip
        vocab += [("<mask>", 0.0)]
        return vocab

    def unk_id(self, proto):
        return 3

    def post_processor(self):
        return processors.TemplateProcessing(
            single="en_XX $A </s>",
            pair="en_XX $A $B </s>",
            special_tokens=[
                ("en_XX", self.original_tokenizer.convert_tokens_to_ids("en_XX")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Summary**

The `MBart50Converter` class extends `SpmConverter`. This converter is responsible for preparing the text data for model training.

**Purpose and Functionality**

- Prepares vocabulary by adding pieces from a translation protocol (`proto`) and language codes.
- Defines vocabulary as a list of tuples containing words and their corresponding scores.
- Returns a predefined vocabulary with additional placeholders.

**Implementation Details**

- `vocab` function constructs the vocabulary using two lists: one for known phrases ( `<s>`, `<pad>`, `</s>`, `<unk>` ) and another for language codes.
- `_unk_id` function returns a predefined ID (`3`) for unknown tokens.
- `post_processor` uses `TemplateProcessing` to manage specific template formats for single and pair inputs.

**Dependencies**

- Utilizes the `original_tokenizer` from a related class or module.


key: chroma:document
string_value:
CODE:

class NllbConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<s>", 0.0),
            ("<pad>", 0.0),
            ("</s>", 0.0),
            ("<unk>", 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        return vocab

    def unk_id(self, proto):
        return 3

    def post_processor(self):
        return processors.TemplateProcessing(
            single="eng_Latn $A </s>",
            pair="eng_Latn $A $B </s>",
            special_tokens=[
                ("eng_Latn", self.original_tokenizer.convert_tokens_to_ids("eng_Latn")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**NllbConverter Class Summary**

The `NllbConverter` class inherits from `SpmConverter`. Its main purpose is to convert input data for a neural language model (NLLB) into a suitable format.

**Primary Functionality**

* Defines the vocabulary by selecting special tokens and adding generic pieces to populate it
	+ Key implementation: `(piece.piece, piece.score)` iteration in `vocab` method
* Specifies the `<unk>` token ID
* Specifies post-processing using `TemplateProcessing`, including:
	+ Tokenization for special tokens ("eng_Latn" and "</s>") using `original_tokenizer`
	+ Single and pair processing templates

**Notable Dependencies**

* Inherited from `SpmConverter`: no specific dependencies explicitly stated, but implies use of BERT's SentencePiece (Spm) library
* External dependencies:
	+ `processors` and `bert` libraries for post-processing and tokenization


key: chroma:document
string_value:
CODE:

class SeamlessM4TConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<pad>", 0.0),
            ("<unk>", 0.0),
            ("<s>", 0.0),
            ("</s>", 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        return vocab

    def unk_id(self, proto):
        return self.original_tokenizer.unk_token_id

    def post_processor(self):
        return processors.TemplateProcessing(
            single="__eng__ $A </s>",
            pair="__eng__ $A $B </s>",
            special_tokens=[
                ("__eng__", self.original_tokenizer.convert_tokens_to_ids("__eng__")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Class Summary: SeamlessM4TConverter**

The `SeamlessM4TConverter` class extends `SpmConverter`. Its primary purpose is to convert and process input data for model training.

**Key Functionality:**

* Constructs a vocabulary list (`vocab`) by including predefined special tokens and additional pieces from the input protocol.
* Defines an unknown token ID using the original tokenizer's `unk_token_id`.
* Implements template processing with two single-piece templates and one pair-piecetemplate.

**Notable Implementation Details:**

* Uses a list comprehension to extend the vocabulary list.
* Utilizes the `processors.TemplateProcessing` class.

**Dependencies/Requirements:**

* Extends `SpmConverter` from the parent class.
* Uses an instance variable `original_tokenizer`.
* Requires external dependencies:
	+ `processors`
	+ `SpmConverter`


key: chroma:document
string_value:
CODE:

class XLMRobertaConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<s>", 0.0),
            ("<pad>", 0.0),
            ("</s>", 0.0),
            ("<unk>", 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        vocab += [("<mask>", 0.0)]
        return vocab

    def unk_id(self, proto):
        unk_id = 3
        return unk_id

    def post_processor(self):
        return processors.TemplateProcessing(
            single="<s> $A </s>",
            pair="<s> $A </s> </s> $B </s>",
            special_tokens=[
                ("<s>", self.original_tokenizer.convert_tokens_to_ids("<s>")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Class XLMRobertaConverter: Summary**

The `XLMRobertaConverter` class is an extension of `SpmConverter` and inherits its functionality. It is designed for converting data into the format required by the Hugging Face Transformers library for Roberta models.

**Main Purpose:**
The primary purpose of this class is to create a vocabulary and post-processing configuration for converting input data into the Roberta model's expected format.

**Key Implementation Details:**

*   The `vocab` method creates an initial vocabulary with predefined tokens (`<s>`, `<pad>`, `</s>`, and `</unk>`), and then appends additional tokens from the input prototype (`[3:]`) to the vocabulary.
*   The `unk_id` method returns a fixed token ID (3) for unknown tokens.
*   The `post_processor` method defines a template-based post-processing configuration, including special token conversions using the `original_tokenizer`.

**Notable Dependencies and Requirements:**

*   `SpmConverter`: The base class inherited by XLMRobertaConverter.
*   `proto.pieces`: The input prototype's piece information is used to construct the vocabulary.
*   `processors.TemplateProcessing`: The post-processing configuration uses a template-based approach for applying substitutions.

**References:**

*   Hugging Face Transformers library
*   Roberta model documentation


key: chroma:document
string_value:
CODE:

class XLNetConverter(SpmConverter):
    def vocab(self, proto):
        return [
            (piece.piece, piece.score) if check_number_comma(piece.piece) else (piece.piece, piece.score - 100)
            for piece in proto.pieces
        ]

    def normalizer(self, proto):
        list_normalizers = [
            normalizers.Replace("``", '"'),
            normalizers.Replace("''", '"'),
        ]
        if not self.original_tokenizer.keep_accents:
            list_normalizers.append(normalizers.NFKD())
            list_normalizers.append(normalizers.StripAccents())
        if self.original_tokenizer.do_lower_case:
            list_normalizers.append(normalizers.Lowercase())

        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap

        if precompiled_charsmap:
            list_normalizers.append(normalizers.Precompiled(precompiled_charsmap))

        list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " "))
        return normalizers.Sequence(list_normalizers)

    def post_processor(self):
        return processors.TemplateProcessing(
            single="$A:0 <sep>:0 <cls>:2",
            pair="$A:0 <sep>:0 $B:1 <sep>:1 <cls>:2",
            special_tokens=[
                ("<sep>", self.original_tokenizer.convert_tokens_to_ids("<sep>")),
                ("<cls>", self.original_tokenizer.convert_tokens_to_ids("<cls>")),
            ],


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

**XLNetConverter Class Summary**

The XLNetConverter class inherits from SpmConverter and provides two main functionalities:

1. **Vocab Functionality**: The `vocab` method generates a vocabulary list by iteratively applying normalization operations to each piece in the proto.pieces list.
   - It utilizes various precompiled character maps and replaces comma-separated pieces with scores.

2. **Normalizer Functionality**: The `normalizer` method creates a normalization pipeline that performs:
   - Replaces double quotes with single quotes
   - Applies NLTK NFKD and strip accents if required by `self.original_tokenizer`
   - Converts case to lower if specified by `self.original_tokenizer`

3. **Post-Processor Functionality**: The `post_processor` method defines a templated template processing pipeline:
   - Templates text into pairs with "<sep>":0 <cls>:2 format
   - Uses custom templates and ID conversion functions for special tokens

Note: This summary excludes repetitive implementation details such as external dependency references or interface implementations, focusing only on the main purposes and key features of the class.

Key method signatures and interfaces:

- `vocab(proto)`
- `normalizer(proto)`
- `post_processor()`

Related classes/functions referred to by exact names:

- `pieces` inside `proto`
- `original_tokenizer` inherited from a parent class (no reference provided)
- `normalizers` within the context of `self.normalizer_spec.precompiled_charsmap`
- `processors` within the context of `self.original_tokenizer`


key: chroma:document
string_value:
CODE:

class ReformerConverter(SpmConverter):


SUMMARY:
**Summary:**

The `ReformerConverter` class extends the `SpmConverter` base class. It appears to be an empty class with no additional implementation, suggesting that it is intended for future extension or specialization, but currently serves as a placeholder or stub.

Key Points:

* Subclass of `SpmConverter`
* No additional implementation provided
* Primary purpose: Potential holder for reformer-related functionality
* Notable dependencies:
	+ Inheriting from `SpmConverter`

**RAG (Reusable Automated Generation) Representation:** 

```
{
  "name": "ReformerConverter",
  "type": "class",
  "parent": "SpmConverter"
}
```


key: chroma:document
string_value:
CODE:

class RemBertConverter(SpmConverter):
    # Inspired from AlbertConverter
    def normalizer(self, proto):
        list_normalizers = [
            normalizers.Replace("``", '"'),
            normalizers.Replace("''", '"'),
            normalizers.Replace(Regex(" {2,}"), " "),
        ]
        if not self.original_tokenizer.keep_accents:
            list_normalizers.append(normalizers.NFKD())
            list_normalizers.append(normalizers.StripAccents())
        if self.original_tokenizer.do_lower_case:
            list_normalizers.append(normalizers.Lowercase())

        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap

        if precompiled_charsmap:
            list_normalizers.append(normalizers.Precompiled(precompiled_charsmap))

        return normalizers.Sequence(list_normalizers)

    def post_processor(self):
        return processors.TemplateProcessing(
            single="[CLS]:0 $A:0 [SEP]:0",
            pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1",
            special_tokens=[
                ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")),
                ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")),
            ],


SUMMARY:
**Class Summary: RemBertConverter**

*   **Purpose and Functionality:** The `RemBertConverter` class is designed to convert text data into BERT-compatible input formats. It extends the `SpmConverter` base class and inherits functionality from `AlbertConverter`.
*   **Main Implementation Details:**
    *   The `normalizer` method applies a sequence of preprocessing steps, including accent stripping, normalization, and conversion to lowercase, based on the user's configuration options.
    *   The `post_processor` method formats the input data into templates suitable for BERT models, including special token conversions.
*   **Notable Dependencies or Requirements:**
    *   `normalizers`: A library of preprocessing normalizers, including `Replace`, `NFKD`, and `StripAccents`.
    *   `processors`: A library of template processing functions, including `TemplateProcessing` with methods for special token conversion and string formatting.
    *   `Regex`: A library for regular expression matching (used in conjunction with `normalizers.Replace`).
    *   `SpmConverter` and `AlbertConverter`: Abstract base classes providing a foundation for tokenization and processing tasks.

**Key Implementation Highlights:**

*   Customizable preprocessing steps using `normalizer` method
*   Special token conversions using `processors.TemplateProcessing`
*   Integration with external libraries (`normalizers`, `processors`, `Regex`, `SpmConverter`, and `AlbertConverter`)


key: chroma:document
string_value:
CODE:

class BertGenerationConverter(SpmConverter):


SUMMARY:
**Summary:**

`BertGenerationConverter` is an unimplemented class that inherits from `SpmConverter`. It appears to be part of a text generation system, likely utilizing BERT (Bidirectional Encoder Representations from Transformers) for natural language processing.

**Key Details:**

* Class name: `BertGenerationConverter`
* Base class: `SpmConverter` (exact implementation not provided)
* Purpose: Unimplemented, intended for text generation using BERT

No notable dependencies are mentioned.


key: chroma:document
string_value:
CODE:

class PegasusConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            (self.original_tokenizer.pad_token, 0.0),
            (self.original_tokenizer.eos_token, 0.0),
        ]

        if self.original_tokenizer.mask_token_sent is not None:
            vocab += [(self.original_tokenizer.mask_token_sent, 0.0)]

        if (
            self.original_tokenizer.mask_token is not None
            and self.original_tokenizer.mask_token_id < self.original_tokenizer.offset
        ):
            vocab += [(self.original_tokenizer.mask_token, 0.0)]

        vocab += [(f"<unk_{i}>", -100.0) for i in range(2, self.original_tokenizer.offset)]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[2:]]
        return vocab

    def unk_id(self, proto):
        return proto.trainer_spec.unk_id + self.original_tokenizer.offset

    def pre_tokenizer(self, replacement, add_prefix_space):
        prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer)
        return pre_tokenizers.Sequence(
            [
                pre_tokenizers.WhitespaceSplit(),
                pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme),
            ]
        )

    def post_processor(self):
        eos = self.original_tokenizer.eos_token
        special_tokens = [
            (eos, self.original_tokenizer.eos_token_id),
        ]


SUMMARY:
**Class Summary: PegasusConverter**

The `PegasusConverter` class inherits from `SpmConverter` and is designed to convert raw text data into format suitable for model training. The primary purpose is to tokenize the input data, handle special tokens, and implement custom preprocessing logic.

### Key Implementation Details:

*   Tokenizes the input using a combination of whitespace splitting and metaspace processing.
*   Handles special tokens (pad token, eos token, mask token) and unk tokens using lookup tables (`vocab` method).
*   Implements template processing for handling special tokens during model evaluation.
*   Uses a custom `pre_tokenizer` class to prepend scheme information and add prefix space before tokenization.

### Notable Dependencies:

*   `original_tokenizer`: an instance of the original tokenizer used in the converter.
*   `proto.pieces[2:]`: a list of pieces from the input data.
*   `_get_prepend_scheme()`: a helper function to determine the prepend scheme for metaspace processing.

### Interfaces and Dependencies:

*   `SpmConverter`: the base class inherited by `PegasusConverter`.
*   `pre_tokenizers` module: provides sequence, whitespace split, and metaspace classes used in the converter.
*   `processors` module: provides template processing classes used in the converter.


key: chroma:document
string_value:
CODE:

class T5Converter(SpmConverter):
    def vocab(self, proto):
        num_extra_ids = self.original_tokenizer._extra_ids
        vocab = [(piece.piece, piece.score) for piece in proto.pieces]
        vocab += [(f"<extra_id_{i}>", 0.0) for i in range(num_extra_ids - 1, -1, -1)]
        return vocab

    def post_processor(self):
        return processors.TemplateProcessing(
            single=["$A", "</s>"],
            pair=["$A", "</s>", "$B", "</s>"],
            special_tokens=[
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**T5Converter Class Summary**

The `T5Converter` class is a subclass of `SpmConverter`. It serves as a converter for the T5 (Transformer 5) model architecture.

### Primary Purpose

The primary purpose of this converter is to tokenize input data in a way that prepares it for processing by the T5 model. This includes:

* Creating a vocabulary from the input data, including extra IDs.
* Defining post-processing rules for tokenization and special tokens.

### Key Implementation Details

* The `vocab` method creates a list of tuples containing vocabulary words and their corresponding scores from the input data's pieces.
* The `post_processor` method defines processing rules for tokenization and special tokens using the `processors.TemplateProcessing` class. Specifically, it:

	+ Defines separate rules for single- and pair-piece tokenization.
	+ Includes extra IDs in the post-processing step.

### Notable Dependencies or Requirements

* `SpmConverter`: The base converter class that this class inherits from.
* `original_tokenizer`: An external dependency used for converting tokens to IDs.

### Interfaces Implemented

None explicitly mentioned. However, it is assumed that T5Converter implements internal interfaces specific to the Transformer 5 model architecture.


key: chroma:document
string_value:
CODE:

class UdopConverter(SpmConverter):
    def post_processor(self):
        return processors.TemplateProcessing(
            single=["$A", "</s>"],
            pair=["$A", "</s>", "$B", "</s>"],
            special_tokens=[
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**Class Summary: UdopConverter**

The `UdopConverter` class extends `SpmConverter` and implements a custom post-processing pipeline for template-based text processing.

**Main Purpose and Functionality:**
It configures the `TemplateProcessing` processor to replace specific tokens in input data:

* `"$A"` with an empty string (`" "</s>``)
* Pairs of `$A` and `/</s>` with their paired tokens, suffixed by `/</s>` and `/</s>`
* Converts the special token `/</s>` to its IDs using `original_tokenizer`

**Key Implementation Details:**

* Uses `TemplateProcessing` processor with specific token replacements
	+ Replaces `$A` with `/ </s>`
	+ Replaces pairs of `$A`, `/</s>`, and `/</s>` with their paired tokens, suffixed by `/</s>`

**Notable Dependencies or Requirements:**

* Inherit from `SpmConverter`
* Uses:
	+ `processors.TemplateProcessing`
	+ `original_tokenizer` for token conversion
* Requires `single` and `pair` lists of tokens to replace

**Interfaces Implemented:**
Extends `SpmConverter`, inheriting its interface.


key: chroma:document
string_value:
CODE:

class WhisperConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.encoder
        merges = list(self.original_tokenizer.bpe_ranks.keys())

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
            )
        )

        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=self.original_tokenizer.add_prefix_space)
        tokenizer.decoder = decoders.ByteLevel()

        prefix_token_ids = self.original_tokenizer.prefix_tokens
        prefixes = self.original_tokenizer.convert_ids_to_tokens(prefix_token_ids)
        eos = self.original_tokenizer.eos_token
        eos_token_id = self.original_tokenizer.eos_token_id
        prefix_template = " ".join([f"{token}:0" for token in prefixes])
        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{prefix_template} $A:0 {eos}:0",
            pair=f"{prefix_template} $A:0 $B:1 {eos}:1",
            special_tokens=[
                (eos, eos_token_id),
                *zip(prefixes, prefix_token_ids),
            ],
        )



SUMMARY:
**WhisperConverter Class Summary**

The WhisperConverter class extends the Converter base class to create a custom Transformer-based BPE tokenizer. Its primary purpose is to:

1.  **Initialize a BPE Tokenizer**: Creates a `Tokenizer` instance with a configured byte-level BPE model, setting tokenizers, decoders, and post-processor templates.

**Key Implementation Details**

*   Initializes the tokenizer with pre-defined constants from the original tokenizer object (`prefix_tokens`, `eos_token`, etc.).
*   Configures the `pre_tokenizer` and `decoder` using existing tokenizers.
*   Customizes the `post_processor` for template processing, including mapping special tokens.

**Notable Dependencies or Requirements**

*   The class leverages other classes:
    *   `pre_tokenizers.ByteLevel` (for prefix space preparation)
    *   `decoders ByteLevel` (for decoder configuration)
    *   `processors.TemplateProcessing` (for post-processor initialization)

        - External dependencies: none explicitly mentioned


key: chroma:document
string_value:
CODE:

class BigBirdConverter(SpmConverter):
    def post_processor(self):
        return processors.TemplateProcessing(
            single="[CLS]:0 $A:0 [SEP]:0",
            pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1",
            special_tokens=[
                ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")),
                ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")),
            ],


SUMMARY:
**Summary**

The `BigBirdConverter` class inherits from `SpmConverter`. It defines a template processor (`post_processor`) that utilizes the `processors.TemplateProcessing` class. The key functionality includes:

* Defining two template strings for a single and pair input format
* Special token mapping using `original_tokenizer.convert_tokens_to_ids`

**Implementation Details**

The implementation of this code chunk focuses on:

* Using ` processors.TemplateProcessing` with customized templates for single and pair inputs
* Mappings special tokens (`[CLS]`, `[SEP]`) to their corresponding IDs using the `original_tokenizer`

**Notable Dependencies or Requirements**

This code snippet references:
* `SpmConverter` (base class)
* `processors.TemplateProcessing` (utility class)

Note: The original tokenizer, which is not shown in this code chunk, seems to be a requirement for the template processing functionality.


key: chroma:document
string_value:
CODE:

class CLIPConverter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.encoder
        merges = list(self.original_tokenizer.bpe_ranks.keys())
        unk_token = self.original_tokenizer.unk_token

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="</w>",
                fuse_unk=False,
                unk_token=str(unk_token),
            )
        )

        tokenizer.normalizer = normalizers.Sequence(
            [normalizers.NFC(), normalizers.Replace(Regex(r"\s+"), " "), normalizers.Lowercase()]
        )
        tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
            [
                pre_tokenizers.Split(
                    Regex(r"""'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+"""),
                    behavior="removed",
                    invert=True,
                ),
                pre_tokenizers.ByteLevel(add_prefix_space=False),
            ]
        )
        tokenizer.decoder = decoders.ByteLevel()

        # Hack to have a ByteLevel and TemplaceProcessor
        tokenizer.post_processor = processors.RobertaProcessing(
            sep=(self.original_tokenizer.eos_token, self.original_tokenizer.eos_token_id),
            cls=(self.original_tokenizer.bos_token, self.original_tokenizer.bos_token_id),
            add_prefix_space=False,
            trim_offsets=False,
        )


SUMMARY:
**CLIPConverter Class Summary**

The `CLIPConverter` class inherits from `Converter` and serves as a wrapper for the `Tokenizer`. Its primary purpose is to convert a tokenizer into a CLIP-specific tokenizer.

**Key Implementation Details:**

* Creates a custom `BPE` module with:
	+ Vocabulary (vocabulary) set from the original tokenizer's encoder.
	+ BPE ranks (merges) derived from the original tokenizer's BPE ranks.
	+ Dropout, prefix, and suffix options set to None/null/empty strings.
* Configures token normalization using a sequence of normalizers: NFC/Lowercase/Lowercase.
* Applies token preprocessing using a sequence of pre-tokenizers:
	+ Removes apostrophes and non-alphanumeric characters (except spaces) using byte-level tokenization.
	+ Splits tokens into byte pairs separated by whitespace.
* Enforces ByteLevel decoder.
* Configures post-processing with Roberta-like token replacement for EOS/CLS tokens.

**Dependencies and Requirements:**

* `original_tokenizer`: The original tokenizer object, used to derive vocabulary and BPE ranks.
* `normalizers`: Sequence of normalization functions (NFC/Lowercase).
* `pre_tokenizers`: Sequence of pre-tokenization functions (remove apostrophes, byte-level tokenization, split by whitespace).
* `decoders`: ByteLevel decoder.
* `processors`: Robota-like post-processing processor.


key: chroma:document
string_value:
CODE:

class LayoutLMv2Converter(Converter):
    def converted(self) -> Tokenizer:
        vocab = self.original_tokenizer.vocab
        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))

        tokenize_chinese_chars = False
        strip_accents = False
        do_lower_case = True
        if hasattr(self.original_tokenizer, "basic_tokenizer"):
            tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars
            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case

        tokenizer.normalizer = normalizers.BertNormalizer(
            clean_text=True,
            handle_chinese_chars=tokenize_chinese_chars,
            strip_accents=strip_accents,
            lowercase=do_lower_case,
        )
        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls}:0 $A:0 {sep}:0",
            pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1",
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
            ],
        )
        tokenizer.decoder = decoders.WordPiece(prefix="##")



SUMMARY:
**Summary of LayoutLMv2Converter Class**

The `LayoutLMv2Converter` class is a subclass of `Converter` that converts the original tokenizer configuration to a custom tokenizer for layout-lm v2 models. Its primary purpose is to enable fine-grained control over tokenizer settings, such as tokenization options and normalization.

**Key Features:**

* Implements a WordPiece-based tokenizer with customizable settings (tokenize_chinese_chars, strip_accents, do_lower_case)
* Utilizes a custom `BertNormalizer` with options to handle Chinese characters and strip accents
* Preserves the original cls_token, sep_token, cls_token_id, and sep_token_id for post-processing purposes

**Notable Dependencies:**

* `Tokenizer`: used as the base tokenizer class
* `normalizers.BertNormalizer`
* `pre_tokenizers.BertPreTokenizer`
* `processors.TemplateProcessing`
* `decoders.WordPiece`


key: chroma:document
string_value:
CODE:

class BlenderbotConverter(Converter):
    def converted(self) -> Tokenizer:
        ot = self.original_tokenizer
        vocab = ot.encoder
        merges = list(ot.bpe_ranks.keys())

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
            )
        )

        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space)
        tokenizer.decoder = decoders.ByteLevel()
        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"$A:0 {ot.eos_token}:0",
            special_tokens=[
                (ot.eos_token, ot.eos_token_id),
            ],
        )



SUMMARY:
**BlenderbotConverter Class Summary**

### Purpose and Functionality
The `BlenderbotConverter` class is a subclass of `Converter`, responsible for converting between BlenderBot's original tokenizer and its custom implementation.

### Key Implementation Details

* Creates a new `Tokenizer` instance with a BPE (Byte Pair Encoding) algorithm.
	+ Set vocabulary, merge rankings, and other parameters based on the original tokenizer's configuration.
	+ Preserves additional settings such as prefix-space addition and template processing for tokenization.
* Configures the `pre_tokenizer`, `decoder`, and `post_processor` attributes of the `Tokenizer` instance to match the original tokenizer.

### Notable Dependencies or Requirements

* `original_tokenizer` object with attributes: `encoder`, `bpe_ranks`
* External dependencies:
	+ `BPE` algorithm
	+ `Tokenizer` class
	+ `pre_tokenizers.ByteLevel` and `processors.TemplateProcessing` classes


key: chroma:document
string_value:
CODE:

class XGLMConverter(SpmConverter):
    def vocab(self, proto):
        vocab = [
            ("<s>", 0.0),
            ("<pad>", 0.0),
            ("</s>", 0.0),
            ("<unk>", 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        vocab += [("<madeupword0>", 0.0), ("<madeupword1>", 0.0), ("<madeupword2>", 0.0), ("<madeupword3>", 0.0), ("<madeupword4>", 0.0), ("<madeupword5>", 0.0), ("<madeupword6>", 0.0)]  # fmt: skip
        return vocab

    def unk_id(self, proto):
        unk_id = 3
        return unk_id

    def post_processor(self):
        return processors.TemplateProcessing(
            single="</s> $A",
            pair="</s> $A </s> </s> $B",
            special_tokens=[
                ("<s>", self.original_tokenizer.convert_tokens_to_ids("<s>")),
                ("</s>", self.original_tokenizer.convert_tokens_to_ids("</s>")),
            ],


SUMMARY:
**XGLMConverter Class Summary**

The `XGLMConverter` class extends `SpmConverter` and is responsible for converting a proto object to a vocabulary and processing it.

**Key Functionality:**

* Returns a vocabulary dictionary with predefined tokens and additional custom tokens.
* Defines an unknown token ID of 3.
* Configures a text processor using the `TemplateProcessing` interface, which includes:
	+ Custom handling for `<s>` (start of sequence) and `</s>` (end of sequence) tokens.
	+ Tokenization for predefined symbols.

**Notable Dependencies:**

* `SpmConverter`
* `proto`: an object containing pieces, tokenized by `piece.piece`.
* `processors.TemplateProcessing` interface.

Inherited/Implemented Interfaces:

* `SpmConverter`

**Example Token Vocabulary:**
```python
vocab = [
    ("<s>", 0.0),
    ("<pad>", 0.0),
    ("</s>", 0.0),
    ("<unk>", 0.0),
]
```
This summary focuses on the main functionality, implementation details, and dependencies of the `XGLMConverter` class.


key: chroma:document
string_value:
CODE:

class GemmaConverter(SpmConverter):
    handle_byte_fallback = True
    SpmExtractor = GemmaSentencePieceExtractor
    # start and end of turn tokens must be marked as special
    special_tokens = {"<start_of_turn>", "<end_of_turn>"}

    """"
    split_by_unicode_script: true
    split_by_number: true
    split_by_whitespace: true
    treat_whitespace_as_suffix: false
    allow_whitespace_only_pieces: true
    split_digits: true
    byte_fallback: true
    """

    def normalizer(self, proto):
        return normalizers.Replace(" ", "▁")

    def vocab(self, proto):
        vocab = [
            (self.original_tokenizer.pad_token, 0.0),
            (self.original_tokenizer.eos_token, 0.0),
            (self.original_tokenizer.bos_token, 0.0),
        ]
        for piece in proto.pieces[3:]:
            if piece.piece == "<0x09>":
                vocab += [("\t", piece.score)]
            else:
                vocab += [(piece.piece, piece.score)]
        # vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        return vocab

    def pre_tokenizer(self, replacement, add_prefix_space):
        return pre_tokenizers.Split(" ", "merged_with_previous")

    def unk_id(self, proto):
        unk_id = 3
        return unk_id

    def decoder(self, replacement, add_prefix_space):
        return decoders.Sequence(
            [
                decoders.Replace("▁", " "),
                decoders.ByteFallback(),
                decoders.Fuse(),
            ]


SUMMARY:
**Class Summary: GemmaConverter**

The `GemmaConverter` class inherits from `SpmConverter`. Its primary purpose is to implement a converter for the Gemma sentence piece model.

### Key Functionality:

*   Implements a normalizer that replaces spaces with a Unicode character (`▁`).
*   Specifies a vocabulary for the Gemma model, including padding tokens and special characters.
*   Defines a pre-tokenizer using `Split`, which merges whitespace with previous tokens.
*   Generates an unknown ID (UNK\_ID = 3) for out-of-vocabulary words.

### Notable Implementation Details:

*   Uses the `GemmaSentencePieceExtractor` for token extraction.
*   Implements a byte fallback fallback mechanism.

### Dependencies and Requirements:

*   The Gemma sentence piece model, specifically the `original_tokenizer` and related classes.
*   The `pre_tokenizers` module for the `Split` pre-tokenizer.
*   The `decoders` module for the sequential decoder pipeline (ByteFallback(), Fuse()).


key: chroma:document
string_value:
CODE:

class LlamaConverter(SpmConverter):
    handle_byte_fallback = True

    def vocab(self, proto):
        vocab = [
            (self.original_tokenizer.convert_ids_to_tokens(0), 0.0),
            (self.original_tokenizer.convert_ids_to_tokens(1), 0.0),
            (self.original_tokenizer.convert_ids_to_tokens(2), 0.0),
        ]
        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
        return vocab

    def unk_id(self, proto):
        unk_id = 0
        return unk_id

    def decoder(self, replacement, add_prefix_space):
        sequence = [
            decoders.Replace("▁", " "),
            decoders.ByteFallback(),
            decoders.Fuse(),
        ]
        if add_prefix_space:
            sequence += [decoders.Strip(content=" ", left=1)]
        return decoders.Sequence(sequence)

    def normalizer(self, proto):
        if getattr(self.original_tokenizer, "legacy", True):
            sequence = []
            if getattr(self.original_tokenizer, "add_prefix_space", True):
                sequence += [normalizers.Prepend(prepend="▁")]
            sequence += [normalizers.Replace(pattern=" ", content="▁")]
            return normalizers.Sequence(sequence)
        return None  # non-legacy, no normalizer

    def pre_tokenizer(self, replacement, add_prefix_space):
        if not getattr(self.original_tokenizer, "legacy", True):  # non-legacy, we need a replace
            prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer)
            return pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme, split=False)
        return None

    def post_processor(self):
        # the processor is defined in the LlamaTokenizerFast class.


SUMMARY:
**LlamaConverter Class Summary**

The `LlamaConverter` class is an implementation of the `SpmConverter` class, which appears to be part of a tokenizer converter. It enables conversion between different formats.

Key Features:

*   **Vocabulary creation**: Creates a vocabulary for converting IDs to tokens and vice versa.
*   **Decoding**: Utilizes a sequence of decoders, including byte fallback, fusion, and prefix space addition/removal, based on the `add_prefix_space` parameter.
*   **Normalization**: Optional normalizer is applied with a hardcoded sequence of replacement patterns (e.g., removing spaces and prepending Unicode whitespace characters).
*   **Pre-tokenizer processing**: Applies pre-processing schemes to tokenizer prefixes without marking them as legacy tokens.

Key Implementation Details:

*   `decoder(self, replacement, add_prefix_space)`: Sequence-based decoding.
    *   Utilizes sequence and replacement classes (`decoders.Replace`, `decoders.ByteFallback`, etc.).
*   `normalizer(self, proto)`: Applicable normalizers differ based on the tokenizer's legacy flag.

Dependencies and Requirements:

*   `proto.pieces[3:]` for vocabulary extension
*   `_get_prepend_scheme(add_prefix_space, self.original_tokenizer)` for pre-tokenizer processing

**Inheritance**

Extends the `SpmConverter` class.


key: chroma:document
string_value:
CODE:

class MarkupLMConverter(Converter):
    def converted(self) -> Tokenizer:
        ot = self.original_tokenizer
        vocab = ot.encoder
        merges = list(ot.bpe_ranks.keys())

        tokenizer = Tokenizer(
            BPE(
                vocab=vocab,
                merges=merges,
                dropout=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
                unk_token=self.original_tokenizer.unk_token,
            )
        )

        tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space)
        tokenizer.decoder = decoders.ByteLevel()

        cls = str(self.original_tokenizer.cls_token)
        sep = str(self.original_tokenizer.sep_token)
        cls_token_id = self.original_tokenizer.cls_token_id
        sep_token_id = self.original_tokenizer.sep_token_id

        tokenizer.post_processor = processors.TemplateProcessing(
            single=f"{cls} $A {sep}",
            pair=f"{cls} $A {sep} $B {sep}",
            special_tokens=[
                (cls, cls_token_id),
                (sep, sep_token_id),
            ],
        )



SUMMARY:
**MarkupLMConverter Class Summary**

**Main Purpose and Functionality:**
The `MarkupLMConverter` class converts an original tokenizer into a customized Tokenizer for deep learning models using the BPE (Byte Pair Encoding) merge strategy.

**Key Implementation Details:**

* The tokenizer's `pre_tokenizer`, `decoder`, `post_processor` are updated with custom settings to maintain or enhance original functionality.
* The `BPE` merge strategy is applied with specific settings, including `vocab` and `merges`.

**Notable Dependencies and Requirements:**

* `Tokenizer` class from the dataset is initialized with a new `BPE` instance.
* External dependencies:
	+ `pre_tokenizers BYTE_LEVEL`
	+ `decoders BYTE_LEVEL`
	+ `processors template_processing`
	+ `self.original_tokenizer` (assuming an existing tokenizer object)


key: chroma:document
string_value:
CODE:

class MoshiConverter(SpmConverter):
    handle_byte_fallback = True

    def __init__(self, vocab_file, model_max_length=None, **kwargs):
        requires_backends(self, "protobuf")

        Converter.__init__(self, vocab_file)

        # from .utils import sentencepiece_model_pb2 as model_pb2
        model_pb2 = import_protobuf()

        m = model_pb2.ModelProto()
        with open(vocab_file, "rb") as f:
            m.ParseFromString(f.read())
        self.proto = m

    def normalizer(self, proto):
        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap
        _normalizers = [
            normalizers.Replace(" ", "▁"),
        ]
        if not precompiled_charsmap:
            return normalizers.Sequence(_normalizers)
        else:
            return normalizers.Sequence([normalizers.Precompiled(precompiled_charsmap)] + _normalizers)

    def decoder(self, replacement, add_prefix_space):
        sequence = [
            decoders.Replace("▁", " "),
            decoders.ByteFallback(),
            decoders.Fuse(),
        ]
        if add_prefix_space:
            sequence += [decoders.Strip(content=" ", left=1)]
        return decoders.Sequence(sequence)

    def pre_tokenizer(self, replacement, add_prefix_space):
        prepend_scheme = "first"


SUMMARY:
**MoshiConverter Summary**

The `MoshiConverter` class is a subclass of `SpmConverter`, extending the functionality to handle MOSHI models.

**Primary Purpose:**
Converts MOSHI model protos to converter instances for tokenization and decoding.

**Key Implementation Details:**

*   Handles byte fallback conversion using `normalizers.Replace` and `decoders.ByteFallback`.
*   Implements sequence-based normalization with precompiled characters map using `normalizers.Sequence` and `normalizers.Precompiled`.
*   Customizes decoder pipeline with `decoders.Replace`, `decoders.Fuse`, and optional prefix space stripping.
*   Utilizes `pre_tokenizers.Metaspace` for preprocessing tokens.

**Notable Dependencies/Requirements:**

*   Requires protobuf backend installation (`requires_backends(self, "protobuf")`)
*   Uses MOSHI model protos imported from `sentencepiece_model_pb2`
*   Inherited from `Converter` class
*   Implements custom interfaces (`.normalizer`, `decoder`, and `pre_tokenizer`)


key: chroma:document
string_value:
CODE:
# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
def bytes_to_unicode():
    """
    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
    characters the bpe code barfs on.

    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
    tables between utf-8 bytes and unicode strings.
    """
    bs = (
        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    )
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    cs = [chr(n) for n in cs]


SUMMARY:
**Summary**

Function `bytes_to_unicode` generates a lookup table mapping UTF-8 bytes to Unicode strings. It excludes whitespace and control characters, while utilizing reversible byte pair encoding (BPE) codes.

**Key Implementation Details**

* Generates a binary sequence featuring non-whitespace UTF-8 character ranges (~256 unique chars)
* Expands the sequence by 256 additional code points
* Creates a reverse mapping using `zip` and list comprehensions

**Dependencies/Requirements**

* None (standalone implementation)

**Primary Purpose**
Generate lookup table for encoding UTF-8 bytes to Unicode strings, minimizing unnecessary mappings.


key: chroma:document
string_value:
CODE:

class TikTokenConverter:
    """
    A general tiktoken converter.
    """

    def __init__(
        self,
        vocab_file=None,
        pattern=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""",
        add_prefix_space=False,
        additional_special_tokens=None,
        *args,
        **kwargs,
    ):
        super().__init__(*args)
        self.vocab_file = vocab_file
        self.pattern = pattern
        self.add_prefix_space = add_prefix_space
        self.additional_special_tokens = additional_special_tokens

    def extract_vocab_merges_from_model(self, tiktoken_url: str):
        try:
            from tiktoken.load import load_tiktoken_bpe
        except Exception:
            raise ValueError(
                "`tiktoken` is required to read a `tiktoken` file. Install it with " "`pip install tiktoken`."
            )

        bpe_ranks = load_tiktoken_bpe(tiktoken_url)
        byte_encoder = bytes_to_unicode()

        def token_bytes_to_string(b):
            return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")])

        merges = []
        vocab = {}
        for token, rank in bpe_ranks.items():
            vocab[token_bytes_to_string(token)] = rank
            if len(token) == 1:
                continue
            local = []
            for index in range(1, len(token)):
                piece_l, piece_r = token[:index], token[index:]
                if piece_l in bpe_ranks and piece_r in bpe_ranks and (piece_l + piece_r) in bpe_ranks:
                    local.append((piece_l, piece_r, rank))
            local = sorted(local, key=lambda x: (bpe_ranks[x[0]], bpe_ranks[x[1]]), reverse=False)
            merges.extend(local)
        merges = sorted(merges, key=lambda val: val[2], reverse=False)
        merges = [(token_bytes_to_string(val[0]), token_bytes_to_string(val[1])) for val in merges]
        return vocab, merges

    def tokenizer(self):
        vocab_scores, merges = self.extract_vocab_merges_from_model(self.vocab_file)
        tokenizer = Tokenizer(BPE(vocab_scores, merges, fuse_unk=False))
        if hasattr(tokenizer.model, "ignore_merges"):
            tokenizer.model.ignore_merges = True
        return tokenizer

    def converted(self) -> Tokenizer:
        tokenizer = self.tokenizer()
        tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
            [
                pre_tokenizers.Split(Regex(self.pattern), behavior="isolated", invert=False),
                pre_tokenizers.ByteLevel(add_prefix_space=self.add_prefix_space, use_regex=False),
            ]
        )
        tokenizer.decoder = decoders.ByteLevel()
        tokenizer.add_special_tokens(self.additional_special_tokens)

        tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)



SUMMARY:
**TikTokenConverter Class Summary**

The `TikTokenConverter` class is a general-purpose tokenizer converter for TikTok tokenization. Its primary purpose is to convert text into a suitable format for machine learning models.

**Main Purpose and Functionality:** 
- Converts text input into a format compatible with machine learning models.
- Utilizes BPE (Byte Pair Encoding) tokenization algorithm.
- Supports special tokens and prefix spaces.

**Key Implementation Details:**

*   Uses `tiktoken.load` to load TikToken BPE model from a file.
*   Extracts vocabulary merges from the model using `load_tiktoken_bpe`.
*   Combines tokenization with preprocessing (ByteLevel) and post-processing (ByteLevel trimming).
*   Includes support for special tokens and prefix spaces.

**Notable Method Signatures:**

*   `extract_vocab_merges_from_model`: Extracts vocabulary merges from a TikToken model URL.
*   `tokenizer`: Creates a tokenizer instance using the extracted vocabulary merges.
*   `converted`: Returns a converted tokenizer with customized preprocessing and post-processing configurations.


key: chroma:document
string_value:
CODE:

def convert_slow_tokenizer(transformer_tokenizer, from_tiktoken=False) -> Tokenizer:
    """
    Utilities to convert a slow tokenizer instance in a fast tokenizer instance.

    Args:
        transformer_tokenizer ([`~tokenization_utils_base.PreTrainedTokenizer`]):
            Instance of a slow tokenizer to convert in the backend tokenizer for
            [`~tokenization_utils_base.PreTrainedTokenizerFast`].
       from_tiktoken (bool, optional): Whether to use the `tiktoken` library to convert the tokenizer instead of sentencepiece.
            Defaults to False.

    Return:
        A instance of [`~tokenizers.Tokenizer`] to be used as the backend tokenizer of a
        [`~tokenization_utils_base.PreTrainedTokenizerFast`]
    """

    tokenizer_class_name = transformer_tokenizer.__class__.__name__
    if tokenizer_class_name in SLOW_TO_FAST_CONVERTERS and not from_tiktoken:
        converter_class = SLOW_TO_FAST_CONVERTERS[tokenizer_class_name]
        return converter_class(transformer_tokenizer).converted()

    else:
        try:
            logger.info("Converting from Tiktoken")
            return TikTokenConverter(
                vocab_file=transformer_tokenizer.vocab_file,
                additional_special_tokens=transformer_tokenizer.additional_special_tokens,
            ).converted()
        except Exception:
            raise ValueError(
                f"Converting from Tiktoken failed, if a converter for SentencePiece is available, provide a model path "
                f"with a SentencePiece tokenizer.model file."
                f"Currently available slow->fast convertors: {list(SLOW_TO_FAST_CONVERTERS.keys())}"


SUMMARY:
**Function Summary**

`convert_slow_tokenizer` converts a `transformer_tokenizer` instance to its corresponding fast tokenization format using either a custom converter or the TikToken library.

**Key Implementation Details:**

* Utilizes a dictionary `SLOW_TO_FAST_CONVERTERS` to map tokenizer classes to their respective converters.
* Checks if TikToken conversion is required and attempts to use the provided conversion logic if available.
* Falls back to logging an error message and providing a model path with a SentencePiece tokenizer.model file if the conversion fails.

**Notable Dependencies/Requirements:**

* `SLOW_TO_FAST_CONVERTERS`: a dictionary mapping slow tokenizer classes to their fast tokenization format converters
* `TikTokenConverter` and related classes: library or framework-agnostic converter for converting tokens
* `logger`: logging utility for error messages


key: chroma:document
string_value:
CODE:

class TextKwargs(TypedDict, total=False):
    """
    Keyword arguments for text processing. For extended documentation, check out tokenization_utils_base methods and
    docstrings associated.

    Attributes:
        add_special_tokens (`bool`, *optional*)
            Whether or not to add special tokens when encoding the sequences.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*)
            Activates and controls padding.
        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*):
            Activates and controls truncation.
        max_length (`int`, *optional*):
            Controls the maximum length to use by one of the truncation/padding parameters.
        stride (`int`, *optional*):
            If set, the overflowing tokens will contain some tokens from the end of the truncated sequence.
        is_split_into_words (`bool`, *optional*):
            Whether or not the input is already pre-tokenized.
        pad_to_multiple_of (`int`, *optional*):
            If set, will pad the sequence to a multiple of the provided value.
        return_token_type_ids (`bool`, *optional*):
            Whether to return token type IDs.
        return_attention_mask (`bool`, *optional*):
            Whether to return the attention mask.
        return_overflowing_tokens (`bool`, *optional*):
            Whether or not to return overflowing token sequences.
        return_special_tokens_mask (`bool`, *optional*):
            Whether or not to return special tokens mask information.
        return_offsets_mapping (`bool`, *optional*):
            Whether or not to return `(char_start, char_end)` for each token.
        return_length (`bool`, *optional*):
            Whether or not to return the lengths of the encoded inputs.
        verbose (`bool`, *optional*):
            Whether or not to print more information and warnings.
        padding_side (`str`, *optional*):
            The side on which padding will be applied.
    """

    text_pair: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]]
    text_target: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]
    text_pair_target: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]]
    add_special_tokens: Optional[bool]
    padding: Union[bool, str, PaddingStrategy]
    truncation: Union[bool, str, TruncationStrategy]
    max_length: Optional[int]
    stride: Optional[int]
    is_split_into_words: Optional[bool]
    pad_to_multiple_of: Optional[int]
    return_token_type_ids: Optional[bool]
    return_attention_mask: Optional[bool]
    return_overflowing_tokens: Optional[bool]
    return_special_tokens_mask: Optional[bool]
    return_offsets_mapping: Optional[bool]
    return_length: Optional[bool]
    verbose: Optional[bool]


SUMMARY:
**Concise Summary: TextKwargs Class**

The `TextKwargs` class defines keyword arguments for text processing. It allows customization of various parameters, such as:

* Tokenization and padding strategies
* Truncation and overflow handling
* Padding side and maximum sequence length
* Return values (token type IDs, attention mask, overflowing tokens, etc.)

Key implementation details:

* Supports typing of `TextInput` and related classes
* Specifies optional parameters with default values (`total=False`)
* Features a hierarchical structure for nested input types (`text_pair`, `text_target`, etc.)


key: chroma:document
string_value:
CODE:

class ImagesKwargs(TypedDict, total=False):
    """
    Keyword arguments for image processing. For extended documentation, check the appropriate ImageProcessor
    class methods and docstrings.

    Attributes:
        do_resize (`bool`, *optional*):
            Whether to resize the image.
        size (`Dict[str, int]`, *optional*):
            Resize the shorter side of the input to `size["shortest_edge"]`.
        size_divisor (`int`, *optional*):
            The size by which to make sure both the height and width can be divided.
        crop_size (`Dict[str, int]`, *optional*):
            Desired output size when applying center-cropping.
        resample (`PILImageResampling`, *optional*):
            Resampling filter to use if resizing the image.
        do_rescale (`bool`, *optional*):
            Whether to rescale the image by the specified scale `rescale_factor`.
        rescale_factor (`int` or `float`, *optional*):
            Scale factor to use if rescaling the image.
        do_normalize (`bool`, *optional*):
            Whether to normalize the image.
        image_mean (`float` or `List[float]`, *optional*):
            Mean to use if normalizing the image.
        image_std (`float` or `List[float]`, *optional*):
            Standard deviation to use if normalizing the image.
        do_pad (`bool`, *optional*):
            Whether to pad the image to the `(max_height, max_width)` of the images in the batch.
        pad_size (`Dict[str, int]`, *optional*):
            The size `{"height": int, "width" int}` to pad the images to.
        do_center_crop (`bool`, *optional*):
            Whether to center crop the image.
        data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format for the output image.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format for the input image.
    """

    do_resize: Optional[bool]
    size: Optional[Dict[str, int]]
    size_divisor: Optional[int]
    crop_size: Optional[Dict[str, int]]
    resample: Optional[Union["PILImageResampling", int]]
    do_rescale: Optional[bool]
    rescale_factor: Optional[float]
    do_normalize: Optional[bool]
    image_mean: Optional[Union[float, List[float]]]
    image_std: Optional[Union[float, List[float]]]
    do_pad: Optional[bool]
    pad_size: Optional[Dict[str, int]]
    do_center_crop: Optional[bool]
    data_format: Optional[ChannelDimension]


SUMMARY:
**Summary**

The `ImagesKwargs` class is a keyword argument dictionary that configures image processing operations. It allows users to specify various settings such as resizing, cropping, rescaling, and normalizing images using their respective types and formats.

**Key Implementation Details:**

1.  **Image Processing Options**: The class includes several optional parameters:
    *   `do_resize`, `resample`, and `rescale_factor` control image resizing.
    *   `crop_size`, `do_center_crop`, and `pad_size` manage cropping and padding options.
    *   `do_normalize`, `image_mean`, and `image_std` execute image normalization with user-specified means and standard deviations.

**Notable Dependencies:**

1.  **PIL Image Processing**: This class relies on the Python Imaging Library (PIL) for image operations, specifically `PILImageResampling`.

**Interfaces and Classes:**

The class uses several types from other modules:
*   `TypedDict`: a built-in Python type used to provide additional metadata on dictionary subclasses.
*   `ChannelDimension` (possibly referenced as `int` or `str`): an unknown module-specific class or enumerant for representing channel dimension formats in images.

**External Dependencies:**

1.  **PIL library**: Required for performing image processing operations, specifically PILImageResampling.


key: chroma:document
string_value:
CODE:

class VideosKwargs(TypedDict, total=False):
    """
    Keyword arguments for video processing.

    Attributes:
        do_resize (`bool`):
            Whether to resize the image.
        size (`Dict[str, int]`, *optional*):
            Resize the shorter side of the input to `size["shortest_edge"]`.
        size_divisor (`int`, *optional*):
            The size by which to make sure both the height and width can be divided.
        resample (`PILImageResampling`, *optional*):
            Resampling filter to use if resizing the image.
        do_rescale (`bool`, *optional*):
            Whether to rescale the image by the specified scale `rescale_factor`.
        rescale_factor (`int` or `float`, *optional*):
            Scale factor to use if rescaling the image.
        do_normalize (`bool`, *optional*):
            Whether to normalize the image.
        image_mean (`float` or `List[float]`, *optional*):
            Mean to use if normalizing the image.
        image_std (`float` or `List[float]`, *optional*):
            Standard deviation to use if normalizing the image.
        do_pad (`bool`, *optional*):
            Whether to pad the image to the `(max_height, max_width)` of the images in the batch.
        do_center_crop (`bool`, *optional*):
            Whether to center crop the image.
        data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format for the output image.
        input_data_format (`ChannelDimension` or `str`, *optional*):
            The channel dimension format for the input image.
    """

    do_resize: Optional[bool]
    size: Optional[Dict[str, int]]
    size_divisor: Optional[int]
    resample: Optional["PILImageResampling"]
    do_rescale: Optional[bool]
    rescale_factor: Optional[float]
    do_normalize: Optional[bool]
    image_mean: Optional[Union[float, List[float]]]
    image_std: Optional[Union[float, List[float]]]
    do_pad: Optional[bool]
    do_center_crop: Optional[bool]
    data_format: Optional[ChannelDimension]


SUMMARY:
**Video Processing Keyword Arguments Summary**

The `VideosKwargs` class defines a TypedDict for keyword arguments related to video processing. Its primary purpose is to provide a structured way to specify various transformations and formats for input and output images.

**Key Implementation Details:**

* Resizes the image using resizing algorithm (`size`, `resample`) or rescales it by the specified scale (`do_rescale`, `rescale_factor`).
* Normalizes images with respective mean and standard deviation (`image_mean`, `image_std`) values.
* Pads images to `(max_height, max_width)` of batch images if required (`do_pad`).
* Centers crop images using or ignoring this flag (`do_center_crop`).

**Notable Dependencies:**

* `PILImageResampling`: Resampling filter used during image resizing (`resample` parameter).
* `ChannelDimension`: Enum type defining channel dimension formats for input and output images (`data_format`, `input_data_format` parameters).


key: chroma:document
string_value:
CODE:

class AudioKwargs(TypedDict, total=False):
    """
    Keyword arguments for audio processing.

    Attributes:
        sampling_rate (`int`, *optional*):
            The sampling rate at which the `raw_speech` input was sampled.
        raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`):
            The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
            values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
            stereo, i.e. single float per timestep.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding
            index) among:

            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
                sequence if provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
                acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'`
        max_length (`int`, *optional*):
            Maximum length of the returned list and optionally padding length (see above).
        truncation (`bool`, *optional*):
            Activates truncation to cut input sequences longer than *max_length* to *max_length*.
        pad_to_multiple_of (`int`, *optional*):
            If set, will pad the sequence to a multiple of the provided value.
        return_attention_mask (`bool`, *optional*):
            Whether or not [`~ASTFeatureExtractor.__call__`] should return `attention_mask`.
    """

    sampling_rate: Optional[int]
    raw_speech: Optional[Union["np.ndarray", List[float], List["np.ndarray"], List[List[float]]]]
    padding: Optional[Union[bool, str, PaddingStrategy]]
    max_length: Optional[int]
    truncation: Optional[bool]
    pad_to_multiple_of: Optional[int]


SUMMARY:
**AudioKwargs Class Summary**

### Main Purpose and Functionality
The `AudioKwargs` class defines keyword arguments for audio processing in an audio feature extraction model.

### Key Implementation Details

*   **Attributes:**
    *   `sampling_rate`: The sampling rate of the input audio (optional, default None).
    *   `raw_speech`: A sequence or batch of sequences to be padded (optional, default None).
    *   `padding`: Padding strategy for returned sequences (optional, default None).
    *   `max_length`: Maximum length of returned list and padding length (optional, default None).
    *   `truncation`: Whether to truncate input sequences (optional, default None).
    *   `pad_to_multiple_of`: If set, will pad sequence to a multiple of the provided value (optional, default None).
    *   `return_attention_mask`: Whether to return attention mask (optional, default None).

### External Dependencies

*   **NPY** for numerical computations
*   **List** for handling lists of float or numpy arrays


key: chroma:document
string_value:
CODE:

class CommonKwargs(TypedDict, total=False):


SUMMARY:
**Summary:**

This defines a Typedef for `CommonKwargs` which specifies that it returns a dictionary containing a single key-value pair `return_tensors`. The value of `return_tensors` can be either a string or an instance of the `TensorType`.

**Key Implementation Details:**

- Uses `Optional[Union[str, TensorType]]` to make `return_tensors` optional and allow both string and tensor types.


key: chroma:document
string_value:
CODE:

class ProcessingKwargs(TextKwargs, ImagesKwargs, VideosKwargs, AudioKwargs, CommonKwargs, total=False):
    """
    Base class for kwargs passing to processors.
    A model should have its own `ModelProcessorKwargs` class that inherits from `ProcessingKwargs` to provide:
        1) Additional typed keys and that this model requires to process inputs.
        2) Default values for existing keys under a `_defaults` attribute.
    New keys have to be defined as follows to ensure type hinting is done correctly.

    ```python
    # adding a new image kwarg for this model
    class ModelImagesKwargs(ImagesKwargs, total=False):
        new_image_kwarg: Optional[bool]

    class ModelProcessorKwargs(ProcessingKwargs, total=False):
        images_kwargs: ModelImagesKwargs
        _defaults = {
            "images_kwargs: {
                "new_image_kwarg": False,
            }
            "text_kwargs": {
                "padding": "max_length",
            },
        }

    ```

    For Python 3.8 compatibility, when inheriting from this class and overriding one of the kwargs,
    you need to manually update the __annotations__ dictionary. This can be done as follows:

    ```python
    class CustomProcessorKwargs(ProcessingKwargs, total=False):
        images_kwargs: CustomImagesKwargs

    CustomProcessorKwargs.__annotations__["images_kwargs"] = CustomImagesKwargs  # python 3.8 compatibility
    ```python

    """

    common_kwargs: CommonKwargs = {
        **CommonKwargs.__annotations__,
    }
    text_kwargs: TextKwargs = {
        **TextKwargs.__annotations__,
    }
    images_kwargs: ImagesKwargs = {
        **ImagesKwargs.__annotations__,
    }
    videos_kwargs: VideosKwargs = {
        **VideosKwargs.__annotations__,
    }
    audio_kwargs: AudioKwargs = {
        **AudioKwargs.__annotations__,


SUMMARY:
**Class Summary**

`ProcessingKwargs`: Base class for kwargs passing to processors, inheriting from multiple kwargs classes. Its primary purpose is to provide a standardized way of handling model-specific input parameters.

**Key Implementation Details**

* Inheritance hierarchy: `ProcessingKwargs` -> `TextKwargs`, `ImagesKwargs`, ..., `total=False`
* Expose commonkwargs as a dictionary
* Initialize images_kwargs with default values (_defaults)
* Add Python 3.8 compatibility for inheriting from this class

**Notable Dependencies or Requirements**

* Inheritance from specific kwargs classes (e.g., `TextKwargs`, `ImagesKwargs`)
* Use of `_defaults` attribute for storing key-value pairs
* Manual update of `__annotations__` dictionary for Python 3.8 compatibility


key: chroma:document
string_value:
CODE:

class ProcessorMixin(PushToHubMixin):
    """
    This is a mixin used to provide saving/loading functionality for all processor classes.
    """

    attributes = ["feature_extractor", "tokenizer"]
    optional_attributes = ["chat_template"]
    optional_call_args: List[str] = []
    # Names need to be attr_class for attr in attributes
    feature_extractor_class = None
    tokenizer_class = None
    _auto_class = None
    valid_kwargs: List[str] = []

    # args have to match the attributes class attribute
    def __init__(self, *args, **kwargs):
        # First, extract optional attributes from kwargs if present
        # Optional attributes can never be positional arguments
        for optional_attribute in self.optional_attributes:
            setattr(self, optional_attribute, kwargs.pop(optional_attribute, None))
        # Sanitize args and kwargs
        for key in kwargs:
            if key not in self.attributes:
                raise TypeError(f"Unexpected keyword argument {key}.")
        for arg, attribute_name in zip(args, self.attributes):
            if attribute_name in kwargs:
                raise TypeError(f"Got multiple values for argument {attribute_name}.")
            else:
                kwargs[attribute_name] = arg

        if len(kwargs) != len(self.attributes):
            raise ValueError(
                f"This processor requires {len(self.attributes)} arguments: {', '.join(self.attributes)}. Got "
                f"{len(args)} arguments instead."
            )

        # Check each arg is of the proper class (this will also catch a user initializing in the wrong order)
        for attribute_name, arg in kwargs.items():
            class_name = getattr(self, f"{attribute_name}_class")
            # Nothing is ever going to be an instance of "AutoXxx", in that case we check the base class.
            class_name = AUTO_TO_BASE_CLASS_MAPPING.get(class_name, class_name)
            if isinstance(class_name, tuple):
                proper_class = tuple(getattr(transformers_module, n) for n in class_name if n is not None)
            else:
                proper_class = getattr(transformers_module, class_name)

            if not isinstance(arg, proper_class):
                raise TypeError(
                    f"Received a {type(arg).__name__} for argument {attribute_name}, but a {class_name} was expected."
                )

            setattr(self, attribute_name, arg)

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes this instance to a Python dictionary.

        Returns:
            `Dict[str, Any]`: Dictionary of all the attributes that make up this processor instance.
        """
        output = copy.deepcopy(self.__dict__)

        # Get the kwargs in `__init__`.
        sig = inspect.signature(self.__init__)
        # Only save the attributes that are presented in the kwargs of `__init__`.
        attrs_to_save = sig.parameters
        # Don't save attributes like `tokenizer`, `image processor` etc.
        attrs_to_save = [x for x in attrs_to_save if x not in self.__class__.attributes]
        # extra attributes to be kept
        attrs_to_save += ["auto_map"]

        output = {k: v for k, v in output.items() if k in attrs_to_save}

        output["processor_class"] = self.__class__.__name__

        if "tokenizer" in output:
            del output["tokenizer"]
        if "image_processor" in output:
            del output["image_processor"]
        if "feature_extractor" in output:
            del output["feature_extractor"]
        if "chat_template" in output:
            del output["chat_template"]

        # Some attributes have different names but containing objects that are not simple strings
        output = {
            k: v
            for k, v in output.items()
            if not (isinstance(v, PushToHubMixin) or v.__class__.__name__ == "BeamSearchDecoderCTC")
        }

        return output

    def to_json_string(self) -> str:
        """
        Serializes this instance to a JSON string.

        Returns:
            `str`: String containing all the attributes that make up this feature_extractor instance in JSON format.
        """
        dictionary = self.to_dict()

        return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"

    def to_json_file(self, json_file_path: Union[str, os.PathLike]):
        """
        Save this instance to a JSON file.

        Args:
            json_file_path (`str` or `os.PathLike`):
                Path to the JSON file in which this processor instance's parameters will be saved.
        """
        with open(json_file_path, "w", encoding="utf-8") as writer:
            writer.write(self.to_json_string())

    def __repr__(self):
        attributes_repr = [f"- {name}: {repr(getattr(self, name))}" for name in self.attributes]
        attributes_repr = "\n".join(attributes_repr)
        return f"{self.__class__.__name__}:\n{attributes_repr}\n\n{self.to_json_string()}"

    def save_pretrained(self, save_directory, push_to_hub: bool = False, **kwargs):
        """
        Saves the attributes of this processor (feature extractor, tokenizer...) in the specified directory so that it
        can be reloaded using the [`~ProcessorMixin.from_pretrained`] method.

        <Tip>

        This class method is simply calling [`~feature_extraction_utils.FeatureExtractionMixin.save_pretrained`] and
        [`~tokenization_utils_base.PreTrainedTokenizerBase.save_pretrained`]. Please refer to the docstrings of the
        methods above for more information.

        </Tip>

        Args:
            save_directory (`str` or `os.PathLike`):
                Directory where the feature extractor JSON file and the tokenizer files will be saved (directory will
                be created if it does not exist).
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                namespace).
            kwargs (`Dict[str, Any]`, *optional*):
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
        """
        use_auth_token = kwargs.pop("use_auth_token", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if kwargs.get("token", None) is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            kwargs["token"] = use_auth_token

        os.makedirs(save_directory, exist_ok=True)

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
            repo_id = self._create_repo(repo_id, **kwargs)
            files_timestamps = self._get_files_timestamps(save_directory)
        # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be
        # loaded from the Hub.
        if self._auto_class is not None:
            attrs = [getattr(self, attribute_name) for attribute_name in self.attributes]
            configs = [(a.init_kwargs if isinstance(a, PreTrainedTokenizerBase) else a) for a in attrs]
            configs.append(self)
            custom_object_save(self, save_directory, config=configs)

        for attribute_name in self.attributes:
            attribute = getattr(self, attribute_name)
            # Include the processor class in the attribute config so this processor can then be reloaded with the
            # `AutoProcessor` API.
            if hasattr(attribute, "_set_processor_class"):
                attribute._set_processor_class(self.__class__.__name__)
            attribute.save_pretrained(save_directory)

        if self._auto_class is not None:
            # We added an attribute to the init_kwargs of the tokenizers, which needs to be cleaned up.
            for attribute_name in self.attributes:
                attribute = getattr(self, attribute_name)
                if isinstance(attribute, PreTrainedTokenizerBase):
                    del attribute.init_kwargs["auto_map"]

        # If we save using the predefined names, we can load using `from_pretrained`
        # plus we save chat_template in its own file
        output_processor_file = os.path.join(save_directory, PROCESSOR_NAME)
        output_raw_chat_template_file = os.path.join(save_directory, "chat_template.jinja")
        output_chat_template_file = os.path.join(save_directory, "chat_template.json")

        processor_dict = self.to_dict()
        # Save `chat_template` in its own file. We can't get it from `processor_dict` as we popped it in `to_dict`
        # to avoid serializing chat template in json config file. So let's get it from `self` directly
        if self.chat_template is not None:
            if kwargs.get("save_raw_chat_template", False):
                with open(output_raw_chat_template_file, "w", encoding="utf-8") as writer:
                    writer.write(self.chat_template)
                logger.info(f"chat template saved in {output_raw_chat_template_file}")
            else:
                chat_template_json_string = (
                    json.dumps({"chat_template": self.chat_template}, indent=2, sort_keys=True) + "\n"
                )
                with open(output_chat_template_file, "w", encoding="utf-8") as writer:
                    writer.write(chat_template_json_string)
                logger.info(f"chat template saved in {output_chat_template_file}")

        # For now, let's not save to `processor_config.json` if the processor doesn't have extra attributes and
        # `auto_map` is not specified.
        if set(processor_dict.keys()) != {"processor_class"}:
            self.to_json_file(output_processor_file)
            logger.info(f"processor saved in {output_processor_file}")

        if push_to_hub:
            self._upload_modified_files(
                save_directory,
                repo_id,
                files_timestamps,
                commit_message=commit_message,
                token=kwargs.get("token"),
            )

        if set(processor_dict.keys()) == {"processor_class"}:
            return []
        return [output_processor_file]

    @classmethod
    def get_processor_dict(
        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """
        From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a
        processor of type [`~processing_utils.ProcessingMixin`] using `from_args_and_dict`.

        Parameters:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                The identifier of the pre-trained checkpoint from which we want the dictionary of parameters.
            subfolder (`str`, *optional*, defaults to `""`):
                In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
                specify the folder name here.

        Returns:
            `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the processor object.
        """
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", None)
        proxies = kwargs.pop("proxies", None)
        token = kwargs.pop("token", None)
        local_files_only = kwargs.pop("local_files_only", False)
        revision = kwargs.pop("revision", None)
        subfolder = kwargs.pop("subfolder", "")

        from_pipeline = kwargs.pop("_from_pipeline", None)
        from_auto_class = kwargs.pop("_from_auto", False)

        user_agent = {"file_type": "processor", "from_auto_class": from_auto_class}
        if from_pipeline is not None:
            user_agent["using_pipeline"] = from_pipeline

        if is_offline_mode() and not local_files_only:
            logger.info("Offline mode: forcing local_files_only=True")
            local_files_only = True

        pretrained_model_name_or_path = str(pretrained_model_name_or_path)
        is_local = os.path.isdir(pretrained_model_name_or_path)
        if os.path.isdir(pretrained_model_name_or_path):
            processor_file = os.path.join(pretrained_model_name_or_path, PROCESSOR_NAME)

        if os.path.isfile(pretrained_model_name_or_path):
            resolved_processor_file = pretrained_model_name_or_path
            # cant't load chat-template when given a file as pretrained_model_name_or_path
            resolved_chat_template_file = None
            resolved_raw_chat_template_file = None
            is_local = True
        elif is_remote_url(pretrained_model_name_or_path):
            processor_file = pretrained_model_name_or_path
            resolved_processor_file = download_url(pretrained_model_name_or_path)
            # can't load chat-template when given a file url as pretrained_model_name_or_path
            resolved_chat_template_file = None
            resolved_raw_chat_template_file = None
        else:
            processor_file = PROCESSOR_NAME
            chat_template_file = "chat_template.json"
            raw_chat_template_file = "chat_template.jinja"
            try:
                # Load from local folder or from cache or download from model Hub and cache
                resolved_processor_file = cached_file(
                    pretrained_model_name_or_path,
                    processor_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                    token=token,
                    user_agent=user_agent,
                    revision=revision,
                    subfolder=subfolder,
                    _raise_exceptions_for_missing_entries=False,
                )

                # Load chat template from a separate json if exists
                # because making it part of processor-config break BC.
                # Processors in older version do not accept any kwargs
                resolved_chat_template_file = cached_file(
                    pretrained_model_name_or_path,
                    chat_template_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                    token=token,
                    user_agent=user_agent,
                    revision=revision,
                    subfolder=subfolder,
                    _raise_exceptions_for_missing_entries=False,
                )

                resolved_raw_chat_template_file = cached_file(
                    pretrained_model_name_or_path,
                    raw_chat_template_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                    token=token,
                    user_agent=user_agent,
                    revision=revision,
                    subfolder=subfolder,
                    _raise_exceptions_for_missing_entries=False,
                )
            except EnvironmentError:
                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to
                # the original exception.
                raise
            except Exception:
                # For any other exception, we throw a generic error.
                raise EnvironmentError(
                    f"Can't load processor for '{pretrained_model_name_or_path}'. If you were trying to load"
                    " it from 'https://huggingface.co/models', make sure you don't have a local directory with the"
                    f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
                    f" directory containing a {PROCESSOR_NAME} file"
                )

        # Add chat template as kwarg before returning because most models don't have processor config
        if resolved_raw_chat_template_file is not None:
            with open(resolved_raw_chat_template_file, "r", encoding="utf-8") as reader:
                chat_template = reader.read()
            kwargs["chat_template"] = chat_template
        elif resolved_chat_template_file is not None:
            with open(resolved_chat_template_file, "r", encoding="utf-8") as reader:
                text = reader.read()
            chat_template = json.loads(text)["chat_template"]
            kwargs["chat_template"] = chat_template

        # Existing processors on the Hub created before #27761 being merged don't have `processor_config.json` (if not
        # updated afterward), and we need to keep `from_pretrained` work. So here it fallbacks to the empty dict.
        # (`cached_file` called using `_raise_exceptions_for_missing_entries=False` to avoid exception)
        # However, for models added in the future, we won't get the expected error if this file is missing.
        if resolved_processor_file is None:
            return {}, kwargs

        try:
            # Load processor dict
            with open(resolved_processor_file, "r", encoding="utf-8") as reader:
                text = reader.read()
            processor_dict = json.loads(text)

        except json.JSONDecodeError:
            raise EnvironmentError(
                f"It looks like the config file at '{resolved_processor_file}' is not a valid JSON file."
            )

        if is_local:
            logger.info(f"loading configuration file {resolved_processor_file}")
        else:
            logger.info(f"loading configuration file {processor_file} from cache at {resolved_processor_file}")

        if "chat_template" in processor_dict and processor_dict["chat_template"] is not None:
            logger.warning_once(
                "Chat templates should be in a 'chat_template.jinja' file but found key='chat_template' "
                "in the processor's config. Make sure to move your template to its own file."
            )

        if not is_local:
            if "auto_map" in processor_dict:
                processor_dict["auto_map"] = add_model_info_to_auto_map(
                    processor_dict["auto_map"], pretrained_model_name_or_path
                )
            if "custom_pipelines" in processor_dict:
                processor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines(
                    processor_dict["custom_pipelines"], pretrained_model_name_or_path
                )

        return processor_dict, kwargs

    @classmethod
    def from_args_and_dict(cls, args, processor_dict: Dict[str, Any], **kwargs):
        """
        Instantiates a type of [`~processing_utils.ProcessingMixin`] from a Python dictionary of parameters.

        Args:
            processor_dict (`Dict[str, Any]`):
                Dictionary that will be used to instantiate the processor object. Such a dictionary can be
                retrieved from a pretrained checkpoint by leveraging the
                [`~processing_utils.ProcessingMixin.to_dict`] method.
            kwargs (`Dict[str, Any]`):
                Additional parameters from which to initialize the processor object.

        Returns:
            [`~processing_utils.ProcessingMixin`]: The processor object instantiated from those
            parameters.
        """
        processor_dict = processor_dict.copy()
        return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
        chat_template = kwargs.pop("chat_template", None)

        # We have to pop up some unused (but specific) kwargs and then validate that it doesn't contain unused kwargs
        # If we don't pop, some specific kwargs will raise a warning
        if "processor_class" in processor_dict:
            del processor_dict["processor_class"]

        if "auto_map" in processor_dict:
            del processor_dict["auto_map"]

        unused_kwargs = cls.validate_init_kwargs(processor_config=processor_dict, valid_kwargs=cls.valid_kwargs)
        processor = cls(*args, **processor_dict)
        if chat_template is not None:
            setattr(processor, "chat_template", chat_template)

        # Update processor with kwargs if needed
        for key in set(kwargs.keys()):
            if hasattr(processor, key):
                setattr(processor, key, kwargs.pop(key))

        kwargs.update(unused_kwargs)
        logger.info(f"Processor {processor}")
        if return_unused_kwargs:
            return processor, kwargs
        else:
            return processor

    def _merge_kwargs(
        self,
        ModelProcessorKwargs: ProcessingKwargs,
        tokenizer_init_kwargs: Optional[Dict] = None,
        **kwargs,
    ) -> Dict[str, Dict]:
        """
        Method to merge dictionaries of kwargs cleanly separated by modality within a Processor instance.
        The order of operations is as follows:
            1) kwargs passed as before have highest priority to preserve BC.
                ```python
                high_priority_kwargs = {"crop_size" = {"height": 222, "width": 222}, "padding" = "max_length"}
                processor(..., **high_priority_kwargs)
                ```
            2) kwargs passed as modality-specific kwargs have second priority. This is the recommended API.
                ```python
                processor(..., text_kwargs={"padding": "max_length"}, images_kwargs={"crop_size": {"height": 222, "width": 222}}})
                ```
            3) kwargs passed during instantiation of a modality processor have fourth priority.
                ```python
                tokenizer = tokenizer_class(..., {"padding": "max_length"})
                image_processor = image_processor_class(...)
                processor(tokenizer, image_processor) # will pass max_length unless overriden by kwargs at call
                ```
            4) defaults kwargs specified at processor level have lowest priority.
                ```python
                class MyProcessingKwargs(ProcessingKwargs, CommonKwargs, TextKwargs, ImagesKwargs, total=False):
                    _defaults = {
                        "text_kwargs": {
                            "padding": "max_length",
                            "max_length": 64,
                        },
                    }
                ```
        Args:
            ModelProcessorKwargs (`ProcessingKwargs`):
                Typed dictionary of kwargs specifically required by the model passed.
            tokenizer_init_kwargs (`Dict`, *optional*):
                Dictionary of kwargs the tokenizer was instantiated with and need to take precedence over defaults.

        Returns:
            output_kwargs (`Dict`):
                Dictionary of per-modality kwargs to be passed to each modality-specific processor.

        """
        # Initialize dictionaries
        output_kwargs = {
            "text_kwargs": {},
            "images_kwargs": {},
            "audio_kwargs": {},
            "videos_kwargs": {},
            "common_kwargs": {},
        }

        default_kwargs = {
            "text_kwargs": {},
            "images_kwargs": {},
            "audio_kwargs": {},
            "videos_kwargs": {},
            "common_kwargs": {},
        }

        used_keys = set()

        # get defaults from set model processor kwargs if they exist
        for modality in default_kwargs:
            default_kwargs[modality] = ModelProcessorKwargs._defaults.get(modality, {}).copy()
            # update defaults with arguments from tokenizer init
            for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys():
                # init with tokenizer init kwargs if necessary
                if modality_key in tokenizer_init_kwargs:
                    value = (
                        getattr(self.tokenizer, modality_key)
                        if hasattr(self.tokenizer, modality_key)
                        else tokenizer_init_kwargs[modality_key]
                    )
                    default_kwargs[modality][modality_key] = value
        # now defaults kwargs are updated with the tokenizers defaults.
        # pass defaults to output dictionary
        output_kwargs.update(default_kwargs)

        # update modality kwargs with passed kwargs
        non_modality_kwargs = set(kwargs) - set(output_kwargs)
        for modality in output_kwargs:
            for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys():
                # check if we received a structured kwarg dict or not to handle it correctly
                if modality in kwargs:
                    kwarg_value = kwargs[modality].pop(modality_key, "__empty__")
                    # check if this key was passed as a flat kwarg.
                    if kwarg_value != "__empty__" and modality_key in non_modality_kwargs:
                        raise ValueError(
                            f"Keyword argument {modality_key} was passed two times:\n"
                            f"in a dictionary for {modality} and as a **kwarg."
                        )
                elif modality_key in kwargs:
                    # we get a modality_key instead of popping it because modality-specific processors
                    # can have overlapping kwargs
                    kwarg_value = kwargs.get(modality_key, "__empty__")
                else:
                    kwarg_value = "__empty__"
                if kwarg_value != "__empty__":
                    output_kwargs[modality][modality_key] = kwarg_value
                    used_keys.add(modality_key)

        # Determine if kwargs is a flat dictionary or contains nested dictionaries
        if any(key in default_kwargs for key in kwargs):
            # kwargs is dictionary-based, and some keys match modality names
            for modality, subdict in kwargs.items():
                if modality in default_kwargs:
                    for subkey, subvalue in subdict.items():
                        if subkey not in used_keys:
                            output_kwargs[modality][subkey] = subvalue
                            used_keys.add(subkey)
        else:
            # kwargs is a flat dictionary
            for key in kwargs:
                if key not in used_keys:
                    if key in ModelProcessorKwargs.__annotations__["common_kwargs"].__annotations__.keys():
                        output_kwargs["common_kwargs"][key] = kwargs[key]
                    else:
                        logger.warning_once(
                            f"Keyword argument `{key}` is not a valid argument for this processor and will be ignored."
                        )

        # all modality-specific kwargs are updated with common kwargs
        for modality in output_kwargs:
            output_kwargs[modality].update(output_kwargs["common_kwargs"])
        return output_kwargs

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        cache_dir: Optional[Union[str, os.PathLike]] = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: Optional[Union[str, bool]] = None,
        revision: str = "main",
        **kwargs,
    ):
        r"""
        Instantiate a processor associated with a pretrained model.

        <Tip>

        This class method is simply calling the feature extractor
        [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`], image processor
        [`~image_processing_utils.ImageProcessingMixin`] and the tokenizer
        [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`] methods. Please refer to the docstrings of the
        methods above for more information.

        </Tip>

        Args:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                This can be either:

                - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
                  huggingface.co.
                - a path to a *directory* containing a feature extractor file saved using the
                  [`~SequenceFeatureExtractor.save_pretrained`] method, e.g., `./my_model_directory/`.
                - a path or url to a saved feature extractor JSON *file*, e.g.,
                  `./my_model_directory/preprocessor_config.json`.
            **kwargs
                Additional keyword arguments passed along to both
                [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`] and
                [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`].
        """
        kwargs["cache_dir"] = cache_dir
        kwargs["force_download"] = force_download
        kwargs["local_files_only"] = local_files_only
        kwargs["revision"] = revision

        use_auth_token = kwargs.pop("use_auth_token", None)
        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if token is not None:
            kwargs["token"] = token

        args = cls._get_arguments_from_pretrained(pretrained_model_name_or_path, **kwargs)
        processor_dict, kwargs = cls.get_processor_dict(pretrained_model_name_or_path, **kwargs)

        return cls.from_args_and_dict(args, processor_dict, **kwargs)

    @classmethod
    def register_for_auto_class(cls, auto_class="AutoProcessor"):
        """
        Register this class with a given auto class. This should only be used for custom feature extractors as the ones
        in the library are already mapped with `AutoProcessor`.

        <Tip warning={true}>

        This API is experimental and may have some slight breaking changes in the next releases.

        </Tip>

        Args:
            auto_class (`str` or `type`, *optional*, defaults to `"AutoProcessor"`):
                The auto class to register this new feature extractor with.
        """
        if not isinstance(auto_class, str):
            auto_class = auto_class.__name__

        import transformers.models.auto as auto_module

        if not hasattr(auto_module, auto_class):
            raise ValueError(f"{auto_class} is not a valid auto class.")

        cls._auto_class = auto_class

    @classmethod
    def _get_arguments_from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        args = []
        for attribute_name in cls.attributes:
            class_name = getattr(cls, f"{attribute_name}_class")
            if isinstance(class_name, tuple):
                classes = tuple(getattr(transformers_module, n) if n is not None else None for n in class_name)
                use_fast = kwargs.get("use_fast", True)
                if use_fast and classes[1] is not None:
                    attribute_class = classes[1]
                else:
                    attribute_class = classes[0]
            else:
                attribute_class = getattr(transformers_module, class_name)

            args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))
        return args

    @property
    def model_input_names(self):
        first_attribute = getattr(self, self.attributes[0])
        return getattr(first_attribute, "model_input_names", None)

    @staticmethod
    def validate_init_kwargs(processor_config, valid_kwargs):
        kwargs_from_config = processor_config.keys()
        unused_kwargs = {}
        unused_keys = set(kwargs_from_config) - set(valid_kwargs)
        if unused_keys:
            unused_key_str = ", ".join(unused_keys)
            logger.warning(
                f"Some kwargs in processor config are unused and will not have any effect: {unused_key_str}. "
            )
            unused_kwargs = {k: processor_config[k] for k in unused_keys}
        return unused_kwargs

    def prepare_and_validate_optional_call_args(self, *args):
        """
        Matches optional positional arguments to their corresponding names in `optional_call_args`
        in the processor class in the order they are passed to the processor call.

        Note that this should only be used in the `__call__` method of the processors with special
        arguments. Special arguments are arguments that aren't `text`, `images`, `audio`, nor `videos`
        but also aren't passed to the tokenizer, image processor, etc. Examples of such processors are:
            - `CLIPSegProcessor`
            - `LayoutLMv2Processor`
            - `OwlViTProcessor`

        Also note that passing by position to the processor call is now deprecated and will be disallowed
        in future versions. We only have this for backward compatibility.

        Example:
            Suppose that the processor class has `optional_call_args = ["arg_name_1", "arg_name_2"]`.
            And we define the call method as:
            ```python
            def __call__(
                self,
                text: str,
                images: Optional[ImageInput] = None,
                *arg,
                audio=None,
                videos=None,
            )
            ```

            Then, if we call the processor as:
            ```python
            images = [...]
            processor("What is common in these images?", images, arg_value_1, arg_value_2)
            ```

            Then, this method will return:
            ```python
            {
                "arg_name_1": arg_value_1,
                "arg_name_2": arg_value_2,
            }
            ```
            which we could then pass as kwargs to `self._merge_kwargs`
        """
        if len(args):
            warnings.warn(
                "Passing positional arguments to the processor call is now deprecated and will be disallowed in v4.47. "
                "Please pass all arguments as keyword arguments."
            )
        if len(args) > len(self.optional_call_args):
            raise ValueError(
                f"Expected *at most* {len(self.optional_call_args)} optional positional arguments in processor call"
                f"which will be matched with {' '.join(self.optional_call_args)} in the order they are passed."
                f"However, got {len(args)} positional arguments instead."
                "Please pass all arguments as keyword arguments instead (e.g. `processor(arg_name_1=..., arg_name_2=...))`."
            )
        return {arg_name: arg_value for arg_value, arg_name in zip(args, self.optional_call_args)}

    def apply_chat_template(
        self,
        conversation: Union[List[Dict[str, str]]],
        chat_template: Optional[str] = None,
        tokenize: bool = False,
        **kwargs,
    ) -> str:
        """
        Similar to the `apply_chat_template` method on tokenizers, this method applies a Jinja template to input
        conversations to turn them into a single tokenizable string.

        Args:
            conversation (`List[Dict, str, str]`):
                The conversation to format.
            chat_template (`Optional[str]`, *optional*):
                The Jinja template to use for formatting the conversation. If not provided, the tokenizer's
                chat template is used.
            tokenize (`bool`, *optional*, defaults to `False`):
                Whether to tokenize the output or not.
            **kwargs:
                Additional keyword arguments
        """

        if chat_template is None:
            if self.chat_template is not None:
                chat_template = self.chat_template
            else:
                raise ValueError(
                    "No chat template is set for this processor. Please either set the `chat_template` attribute, "
                    "or provide a chat template as an argument. See "
                    "https://huggingface.co/docs/transformers/main/en/chat_templating for more information."
                )
        return self.tokenizer.apply_chat_template(
            conversation, chat_template=chat_template, tokenize=tokenize, **kwargs
        )

    def post_process_image_text_to_text(self, generated_outputs):
        """
        Post-process the output of a vlm to decode the text.

        Args:
            generated_outputs (`torch.Tensor` or `np.ndarray`):
                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
                or `(sequence_length,)`.

        Returns:
            `List[str]`: The decoded text.
        """


SUMMARY:
This is the main class definition for a text processing processor in the Transformers library. Here's a high-level overview of what this code defines:

**Processor Class**

The `Processor` class is a base class for text processing processors that inherits from the `Pipeline`. It provides common functionality and attributes for any text processing pipeline.

**Attributes**

* `optional_call_args`: A tuple of optional positional arguments that can be passed to the processor call.
* `chat_template`: An attribute that stores the chat template used in generating responses from vlm models.
* `tokenizer`: An instance of a tokenization tokenizer class (e.g., `AutoTokenizer`).

**Methods**

* `prepare_and_validate_optional_call_args(arg)` : Matches optional positional arguments to their corresponding names in `optional_call_args`.
* `apply_chat_template(conv, chat_template=None, tokenize=False, **kwargs) : Applies a Jinja template to input conversations to turn them into a single tokenizable string.
* `post_process_image_text_to_text(generated_outputs)`: Post-processes the output of a vlm model by decoding the text.

**Usage**

This class is typically used in pipelines for generating responses from vlm models. For example:

```python
processor = MyTextProcessor()  # Initialize processor with custom chat template and optional call args.
pipeline = Pipeline([processor])
print(pipeline(processed_data))  # Process data through pipeline, including processing and response generation steps.
```

Overall, this class provides a reusable framework for building text processing pipelines that can be easily customized with different tokenization and post-processing techniques.


key: chroma:document
string_value:
CODE:

def _validate_images_text_input_order(images, text):
    """
    For backward compatibility: reverse the order of `images` and `text` inputs if they are swapped.
    This method should only be called for processors where `images` and `text` have been swapped for uniformization purposes.
    Note that this method assumes that two `None` inputs are valid inputs. If this is not the case, it should be handled
    in the processor's `__call__` method before calling this method.
    """

    def is_url(val) -> bool:
        return isinstance(val, str) and val.startswith("http")

    def _is_valid_images_input_for_processor(imgs):
        # If we have an list of images, make sure every image is valid
        if isinstance(imgs, (list, tuple)):
            for img in imgs:
                if not _is_valid_images_input_for_processor(img):
                    return False
        # If not a list or tuple, we have been given a single image or batched tensor of images
        elif not (is_valid_image(imgs) or is_url(imgs)):
            return False
        return True

    def _is_valid_text_input_for_processor(t):
        if isinstance(t, str):
            # Strings are fine
            return True
        elif isinstance(t, (list, tuple)):
            # List are fine as long as they are...
            if len(t) == 0:
                # ... not empty
                return False
            for t_s in t:
                return _is_valid_text_input_for_processor(t_s)
        return False

    def _is_valid(input, validator):
        return validator(input) or input is None

    images_is_valid = _is_valid(images, _is_valid_images_input_for_processor)
    images_is_text = _is_valid_text_input_for_processor(images)

    text_is_valid = _is_valid(text, _is_valid_text_input_for_processor)
    text_is_images = _is_valid_images_input_for_processor(text)
    # Handle cases where both inputs are valid
    if images_is_valid and text_is_valid:
        return images, text

    # Handle cases where inputs need to and can be swapped
    if (images is None and text_is_images) or (text is None and images_is_text) or (images_is_text and text_is_images):
        logger.warning_once(
            "You may have used the wrong order for inputs. `images` should be passed before `text`. "
            "The `images` and `text` inputs will be swapped. This behavior will be deprecated in transformers v4.47."
        )
        return text, images



SUMMARY:
**Function Summary**

 `_validate_images_text_input_order(images, text)` validates the order of image and text inputs. For backward compatibility, it swaps the inputs if necessary.

**Key Implementation Details**

* Validates input types using separate validators for images (`_is_valid_images_input_for_processor`) and text (`_is_valid_text_input_for_processor`).
* Checks for invalid input combinations or single invalid input elements that should be handled in the processor's `__call__` method.
* Swaps inputs if necessary to maintain uniformization, logging a deprecation warning.

**Notable Dependencies/ Requirements**

* `logger` object (logging module)
* `is_valid_image`, `is_url`, `has_length` functions (not shown)

**Primary Purpose**

Validate the order and types of image and text inputs to ensure compatibility with processors.


key: chroma:document
string_value:
CODE:

def quick_gelu(x):


SUMMARY:
**Summary**

`quick_gelu`: A fast variant of the Gelu activation function, defined as `x * sigmoid(1.702 * x)`, where `sigmoid` is a JAX implementation.

* **Purpose**: Replaces traditional Gelu activation with a faster computation.
* **Key detail**: Uses a fixed coefficient (1.702) to scale input `x` before applying the sigmoid function, reducing computational overhead.
* **Dependencies**:
	+ `jax.nn.sigmoid`: JAX's sigmoid function implementation.
	+ No explicit external dependencies listed; likely JAX is imported as part of the standard library setup.

Note: The `quick_gelu` name suggests a trade-off between speed and accuracy, where "quick" implies faster execution at the potential cost of a minor reduction in output precision.


key: chroma:document
string_value:
CODE:

def dtype_byte_size(dtype):
    """
    Returns the size (in bytes) occupied by one parameter of type `dtype`. Example:
    ```py
    >>> dtype_byte_size(np.float32)
    4
    ```
    """
    if dtype is bool:
        return 1 / 8
    bit_search = re.search(r"[^\d](\d+)$", dtype.name)
    if bit_search is None:
        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
    bit_size = int(bit_search.groups()[0])


SUMMARY:
**Summary:**

### Function Purpose:

`dtype_byte_size(dtype)` returns the size (in bytes) occupied by one parameter of type `dtype`.

### Key Implementation Details:

- Uses regular expressions to parse the `_bits_` attribute from the `dtype` name, which represents the number of bits in the data type.
- Applies bit size conversion (`bit_size // 8`) to obtain the byte size.

### Parameters and Dependencies:

- `dtype`: input type (e.g., `np.float32`, `bool`)
- No external dependencies other than Python's `re` module and NumPy library
- Interfaces: None

### Notes:
Raised a `ValueError` if `dtype` is not a valid NumPy data type.


key: chroma:document
string_value:
CODE:

def flax_shard_checkpoint(params, max_shard_size="10GB"):
    """
    Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a
    given size. The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so
    there is no optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For
    example, if the limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as
    [6GB], [6+2GB], [6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB].

    <Tip warning={true}>

    If one of the model's weight is bigger that `max_shard_size`, it will end up in its own sub-checkpoint which will
    have a size greater than `max_shard_size`.

    </Tip>

    Args:
        params (`Union[Dict, FrozenDict]`): A `PyTree` of model parameters.
        max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
            The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit
            (like `"5MB"`).
    """
    max_shard_size = convert_file_size_to_int(max_shard_size)

    sharded_state_dicts = []
    current_block = {}
    current_block_size = 0
    total_size = 0

    # flatten the weights to chunk
    weights = flatten_dict(params, sep="/")
    for item in weights:
        weight_size = weights[item].size * dtype_byte_size(weights[item].dtype)

        # If this weight is going to tip up over the maximal size, we split.
        if current_block_size + weight_size > max_shard_size:
            sharded_state_dicts.append(current_block)
            current_block = {}
            current_block_size = 0

        current_block[item] = weights[item]
        current_block_size += weight_size
        total_size += weight_size

    # Add the last block
    sharded_state_dicts.append(current_block)

    # If we only have one shard, we return it
    if len(sharded_state_dicts) == 1:
        return {FLAX_WEIGHTS_NAME: sharded_state_dicts[0]}, None

    # Otherwise, let's build the index
    weight_map = {}
    shards = {}
    for idx, shard in enumerate(sharded_state_dicts):
        shard_file = FLAX_WEIGHTS_NAME.replace(".msgpack", f"-{idx+1:05d}-of-{len(sharded_state_dicts):05d}.msgpack")
        shards[shard_file] = shard
        for weight_name in shard.keys():
            weight_map[weight_name] = shard_file

    # Add the metadata
    metadata = {"total_size": total_size}
    index = {"metadata": metadata, "weight_map": weight_map}


SUMMARY:
**Function Summary: `flax_shard_checkpoint`**

The `flax_shard_checkpoint` function splits a model state dictionary into sub-checkpoints to limit their total size. It takes two main inputs:

*   `params`: A PyTree of model parameters
*   `max_shard_size`: The maximum size of each sub-checkpoint, which can be an integer or a string

The function's primary purpose is to shard the state dictionary into smaller files while preserving the order of keys. It uses a greedy approach, adding weights to subsequent blocks as long as the total size does not exceed `max_shard_size`.

**Implementation Details:**

*   The function leverages two key algorithms:
    *   **Dictionary Flattening:** `flatten_dict` is used to flatten the state dictionary into a single, hierarchical structure.
    *   **Sharding:** Weights are grouped into blocks until each block's total size exceeds `max_shard_size`.
*   Importantly, the function updates the weights' file names and creates a weight map for tracking.

**Notable Dependencies/Requirements:**

*   The `FLAX_WEIGHTS_NAME` constant is used to construct file paths.
*   The `convert_file_size_to_int` function is assumed to be defined elsewhere.
*   No external dependencies are explicitly listed in the provided code snippet; however, it assumes that `flatten_dict` and other related utilities are available.

**Interfaces and Classes:**

The function relies on several interfaces:

*   `params`: A PyTree interface or FrozenDict-like object
*   `convert_file_size_to_int`: A function to convert file sizes from strings to integers

The function's implementation does not appear to directly inherit from any classes. However, the use of `FLAX_WEIGHTS_NAME` suggests that it operates within a specific framework, possibly involving FLAX for deep learning applications.


key: chroma:document
string_value:
CODE:

class FlaxPreTrainedModel(PushToHubMixin, FlaxGenerationMixin):
    r"""
    Base class for all models.

    [`FlaxPreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,
    downloading and saving models.

    Class attributes (overridden by derived classes):

        - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class
          for this model architecture.
        - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived
          classes of the same architecture adding modules on top of the base model.
        - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP
          models, `pixel_values` for vision models and `input_values` for speech models).
    """

    config_class = None
    base_model_prefix = ""
    main_input_name = "input_ids"
    _auto_class = None
    _missing_keys = set()

    def __init__(
        self,
        config: PretrainedConfig,
        module: nn.Module,
        input_shape: Tuple = (1, 1),
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        _do_init: bool = True,
    ):
        if config is None:
            raise ValueError("config cannot be None")

        if module is None:
            raise ValueError("module cannot be None")

        # Those are private to be exposed as typed property on derived classes.
        self._config = config
        self._module = module

        # Those are public as their type is generic to every derived classes.
        self.key = PRNGKey(seed)
        self.dtype = dtype
        self.input_shape = input_shape
        self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None

        # To check if the model was initialized automatically.
        self._is_initialized = _do_init

        if _do_init:
            # randomly initialized parameters
            random_params = self.init_weights(self.key, input_shape)
            params_shape_tree = jax.eval_shape(lambda params: params, random_params)
        else:
            init_fn = partial(self.init_weights, input_shape=input_shape)
            params_shape_tree = jax.eval_shape(init_fn, self.key)

            logger.info(
                "Model weights are not initialized as `_do_init` is set to `False`. "
                f"Make sure to call `{self.__class__.__name__}.init_weights` manually to initialize the weights."
            )

        # get the shape of the parameters
        self._params_shape_tree = params_shape_tree

        # save required_params as set
        self._required_params = set(flatten_dict(unfreeze(params_shape_tree)).keys())

        # initialize the parameters
        if _do_init:
            self.params = random_params

    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> Dict:
        raise NotImplementedError(f"init method has to be implemented for {self}")

    def enable_gradient_checkpointing(self):
        raise NotImplementedError(f"gradient checkpointing method has to be implemented for {self}")

    @classmethod
    def _from_config(cls, config, **kwargs):
        """
        All context managers that the model should be initialized under go here.
        """
        return cls(config, **kwargs)

    @property
    def framework(self) -> str:
        """
        :str: Identifies that this is a Flax model.
        """
        return "flax"

    @property
    def config(self) -> PretrainedConfig:
        return self._config

    @property
    def module(self) -> nn.Module:
        return self._module

    @property
    def params(self) -> Union[Dict, FrozenDict]:
        if not self._is_initialized:
            raise ValueError(
                "`params` cannot be accessed from model when the model is created with `_do_init=False`. "
                "You must call `init_weights` manually and store the params outside of the model and "
                "pass it explicitly where needed."
            )
        return self._params

    @property
    def required_params(self) -> Set:
        return self._required_params

    @property
    def params_shape_tree(self) -> Dict:
        return self._params_shape_tree

    @params.setter
    def params(self, params: Union[Dict, FrozenDict]):
        # don't set params if the model is not initialized
        if not self._is_initialized:
            raise ValueError(
                "`params` cannot be set from model when the model is created with `_do_init=False`. "
                "You store the params outside of the model."
            )

        if isinstance(params, FrozenDict):
            params = unfreeze(params)
        param_keys = set(flatten_dict(params).keys())
        if len(self.required_params - param_keys) > 0:
            raise ValueError(
                "Some parameters are missing. Make sure that `params` include the following "
                f"parameters {self.required_params - param_keys}"
            )
        self._params = params

    def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any:
        """
        Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`.
        """

        # taken from https://github.com/deepmind/jmp/blob/3a8318abc3292be38582794dbf7b094e6583b192/jmp/_src/policy.py#L27
        def conditional_cast(param):
            if isinstance(param, jnp.ndarray) and jnp.issubdtype(param.dtype, jnp.floating):
                param = param.astype(dtype)
            return param

        if mask is None:
            return jax.tree_util.tree_map(conditional_cast, params)

        flat_params = flatten_dict(params)
        flat_mask, _ = jax.tree_util.tree_flatten(mask)

        for masked, key in zip(flat_mask, sorted(flat_params.keys())):
            if masked:
                flat_params[key] = conditional_cast(flat_params[key])

        return unflatten_dict(flat_params)

    def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None):
        r"""
        Cast the floating-point `params` to `jax.numpy.bfloat16`. This returns a new `params` tree and does not cast
        the `params` in place.

        This method can be used on TPU to explicitly convert the model parameters to bfloat16 precision to do full
        half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed.

        Arguments:
            params (`Union[Dict, FrozenDict]`):
                A `PyTree` of model parameters.
            mask (`Union[Dict, FrozenDict]`):
                A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params
                you want to cast, and should be `False` for those you want to skip.

        Examples:

        ```python
        >>> from transformers import FlaxBertModel

        >>> # load model
        >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")
        >>> # By default, the model parameters will be in fp32 precision, to cast these to bfloat16 precision
        >>> model.params = model.to_bf16(model.params)
        >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale)
        >>> # then pass the mask as follows
        >>> from flax import traverse_util

        >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")
        >>> flat_params = traverse_util.flatten_dict(model.params)
        >>> mask = {
        ...     path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale"))
        ...     for path in flat_params
        ... }
        >>> mask = traverse_util.unflatten_dict(mask)
        >>> model.params = model.to_bf16(model.params, mask)
        ```"""
        return self._cast_floating_to(params, jnp.bfloat16, mask)

    def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None):
        r"""
        Cast the floating-point `parmas` to `jax.numpy.float32`. This method can be used to explicitly convert the
        model parameters to fp32 precision. This returns a new `params` tree and does not cast the `params` in place.

        Arguments:
            params (`Union[Dict, FrozenDict]`):
                A `PyTree` of model parameters.
            mask (`Union[Dict, FrozenDict]`):
                A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params
                you want to cast, and should be `False` for those you want to skip

        Examples:

        ```python
        >>> from transformers import FlaxBertModel

        >>> # Download model and configuration from huggingface.co
        >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")
        >>> # By default, the model params will be in fp32, to illustrate the use of this method,
        >>> # we'll first cast to fp16 and back to fp32
        >>> model.params = model.to_f16(model.params)
        >>> # now cast back to fp32
        >>> model.params = model.to_fp32(model.params)
        ```"""
        return self._cast_floating_to(params, jnp.float32, mask)

    def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
        r"""
        Cast the floating-point `parmas` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the
        `params` in place.

        This method can be used on GPU to explicitly convert the model parameters to float16 precision to do full
        half-precision training or to save weights in float16 for inference in order to save memory and improve speed.

        Arguments:
            params (`Union[Dict, FrozenDict]`):
                A `PyTree` of model parameters.
            mask (`Union[Dict, FrozenDict]`):
                A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params
                you want to cast, and should be `False` for those you want to skip

        Examples:

        ```python
        >>> from transformers import FlaxBertModel

        >>> # load model
        >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")
        >>> # By default, the model params will be in fp32, to cast these to float16
        >>> model.params = model.to_fp16(model.params)
        >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale)
        >>> # then pass the mask as follows
        >>> from flax import traverse_util

        >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")
        >>> flat_params = traverse_util.flatten_dict(model.params)
        >>> mask = {
        ...     path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale"))
        ...     for path in flat_params
        ... }
        >>> mask = traverse_util.unflatten_dict(mask)
        >>> model.params = model.to_fp16(model.params, mask)
        ```"""
        return self._cast_floating_to(params, jnp.float16, mask)

    @classmethod
    def load_flax_weights(cls, resolved_archive_file):
        try:
            if resolved_archive_file.endswith(".safetensors"):
                state = safe_load_file(resolved_archive_file)
                state = unflatten_dict(state, sep=".")
            else:
                with open(resolved_archive_file, "rb") as state_f:
                    state = from_bytes(cls, state_f.read())
        except (UnpicklingError, msgpack.exceptions.ExtraData) as e:
            try:
                with open(resolved_archive_file) as f:
                    if f.read().startswith("version"):
                        raise OSError(
                            "You seem to have cloned a repository without having git-lfs installed. Please"
                            " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"
                            " folder you cloned."
                        )
                    else:
                        raise ValueError from e
            except (UnicodeDecodeError, ValueError):
                raise EnvironmentError(f"Unable to convert {resolved_archive_file} to Flax deserializable object. ")

        return state

    @classmethod
    def load_flax_sharded_weights(cls, shard_files):
        """
        This is the same as [`flax.serialization.from_bytes`]
        (https:lax.readthedocs.io/en/latest/_modules/flax/serialization.html#from_bytes) but for a sharded checkpoint.

        This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being
        loaded in the model.

        Args:
            shard_files (`List[str]`:
                The list of shard files to load.

        Returns:
            `Dict`: A nested dictionary of the model parameters, in the expected format for flax models : `{'model':
            {'params': {'...'}}}`.
        """

        # Load the index
        state_sharded_dict = {}

        for shard_file in shard_files:
            # load using msgpack utils
            try:
                with open(shard_file, "rb") as state_f:
                    state = from_bytes(cls, state_f.read())
            except (UnpicklingError, msgpack.exceptions.ExtraData) as e:
                with open(shard_file) as f:
                    if f.read().startswith("version"):
                        raise OSError(
                            "You seem to have cloned a repository without having git-lfs installed. Please"
                            " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"
                            " folder you cloned."
                        )
                    else:
                        raise ValueError from e
            except (UnicodeDecodeError, ValueError):
                raise EnvironmentError(f"Unable to convert {shard_file} to Flax deserializable object. ")

            state = flatten_dict(state, sep="/")
            state_sharded_dict.update(state)
            del state
            gc.collect()

        # the state dict is unflattened to the match the format of model.params
        return unflatten_dict(state_sharded_dict, sep="/")

    @classmethod
    def can_generate(cls) -> bool:
        """
        Returns whether this model can generate sequences with `.generate()`. Returns:
            `bool`: Whether this model can generate sequences with `.generate()`.
        """
        # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation.
        # Alternativelly, the model can also have a custom `generate` function.
        if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate):
            return False
        return True

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        dtype: jnp.dtype = jnp.float32,
        *model_args,
        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
        cache_dir: Optional[Union[str, os.PathLike]] = None,
        ignore_mismatched_sizes: bool = False,
        force_download: bool = False,
        local_files_only: bool = False,
        token: Optional[Union[str, bool]] = None,
        revision: str = "main",
        **kwargs,
    ):
        r"""
        Instantiate a pretrained flax model from a pre-trained model configuration.

        The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
        pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
        task.

        The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
        weights are discarded.

        Parameters:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                Can be either:

                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
                    - A path to a *directory* containing model weights saved using
                      [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
                    - A path or url to a *pt index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In this case,
                      `from_pt` should be set to `True`.
            dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):
                The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and
                `jax.numpy.bfloat16` (on TPUs).

                This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If
                specified all the computation will be performed with the given `dtype`.

                **Note that this only specifies the dtype of the computation and does not influence the dtype of model
                parameters.**

                If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and
                [`~FlaxPreTrainedModel.to_bf16`].
            model_args (sequence of positional arguments, *optional*):
                All remaining positional arguments will be passed to the underlying model's `__init__` method.
            config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*):
                Can be either:

                    - an instance of a class derived from [`PretrainedConfig`],
                    - a string or path valid as input to [`~PretrainedConfig.from_pretrained`].

                Configuration for the model to use instead of an automatically loaded configuration. Configuration can
                be automatically loaded when:

                    - The model is a model provided by the library (loaded with the *model id* string of a pretrained
                      model).
                    - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the
                      save directory.
                    - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a
                      configuration JSON file named *config.json* is found in the directory.
            cache_dir (`Union[str, os.PathLike]`, *optional*):
                Path to a directory in which a downloaded pretrained model configuration should be cached if the
                standard cache should not be used.
            from_pt (`bool`, *optional*, defaults to `False`):
                Load the model weights from a PyTorch checkpoint save file (see docstring of
                `pretrained_model_name_or_path` argument).
            ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`):
                Whether or not to raise an error if some of the weights from the checkpoint do not have the same size
                as the weights of the model (if for instance, you are instantiating a model with 10 labels from a
                checkpoint with 3 labels).
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.
            resume_download:
                Deprecated and ignored. All downloads are now resumed by default when possible.
                Will be removed in v5 of Transformers.
            proxies (`Dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
            local_files_only(`bool`, *optional*, defaults to `False`):
                Whether or not to only look at local files (i.e., do not try to download the model).
            token (`str` or `bool`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use
                the token generated when running `huggingface-cli login` (stored in `~/.huggingface`).
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
                identifier allowed by git.


                <Tip>

                To test a pull request you made on the Hub, you can pass `revision="refs/pr/<pr_number>"`.

                </Tip>

            subfolder (`str`, *optional*, defaults to `""`):
                In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
                specify the folder name here.
            kwargs (remaining dictionary of keyword arguments, *optional*):
                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
                `output_attentions=True`). Behaves differently depending on whether a `config` is provided or
                automatically loaded:

                    - If a configuration is provided with `config`, `**kwargs` will be directly passed to the
                      underlying model's `__init__` method (we assume all relevant updates to the configuration have
                      already been done)
                    - If a configuration is not provided, `kwargs` will be first passed to the configuration class
                      initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that
                      corresponds to a configuration attribute will be used to override said attribute with the
                      supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute
                      will be passed to the underlying model's `__init__` function.

        Examples:

        ```python
        >>> from transformers import BertConfig, FlaxBertModel

        >>> # Download model and configuration from huggingface.co and cache.
        >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")
        >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable).
        >>> model = FlaxBertModel.from_pretrained("./test/saved_model/")
        >>> # Loading from a PyTorch checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).
        >>> config = BertConfig.from_json_file("./pt_model/config.json")
        >>> model = FlaxBertModel.from_pretrained("./pt_model/pytorch_model.bin", from_pt=True, config=config)
        ```"""
        from_pt = kwargs.pop("from_pt", False)
        resume_download = kwargs.pop("resume_download", None)
        proxies = kwargs.pop("proxies", None)
        use_auth_token = kwargs.pop("use_auth_token", None)
        trust_remote_code = kwargs.pop("trust_remote_code", None)
        from_pipeline = kwargs.pop("_from_pipeline", None)
        from_auto_class = kwargs.pop("_from_auto", False)
        _do_init = kwargs.pop("_do_init", True)
        subfolder = kwargs.pop("subfolder", "")
        commit_hash = kwargs.pop("_commit_hash", None)

        # Not relevant for Flax Models
        _ = kwargs.pop("adapter_kwargs", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if trust_remote_code is True:
            logger.warning(
                "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is"
                " ignored."
            )

        user_agent = {"file_type": "model", "framework": "flax", "from_auto_class": from_auto_class}
        if from_pipeline is not None:
            user_agent["using_pipeline"] = from_pipeline

        if is_offline_mode() and not local_files_only:
            logger.info("Offline mode: forcing local_files_only=True")
            local_files_only = True

        # Load config if we don't provide a configuration
        if not isinstance(config, PretrainedConfig):
            config_path = config if config is not None else pretrained_model_name_or_path
            config, model_kwargs = cls.config_class.from_pretrained(
                config_path,
                cache_dir=cache_dir,
                return_unused_kwargs=True,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                token=token,
                revision=revision,
                subfolder=subfolder,
                _from_auto=from_auto_class,
                _from_pipeline=from_pipeline,
                _commit_hash=commit_hash,
                **kwargs,
            )
        else:
            model_kwargs = kwargs.copy()

        if commit_hash is None:
            commit_hash = getattr(config, "_commit_hash", None)

        # Add the dtype to model_kwargs
        model_kwargs["dtype"] = dtype

        # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the
        # index of the files.
        is_sharded = False

        # Load model
        if pretrained_model_name_or_path is not None:
            pretrained_model_name_or_path = str(pretrained_model_name_or_path)
            is_local = os.path.isdir(pretrained_model_name_or_path)
            if os.path.isdir(pretrained_model_name_or_path):
                if os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)):
                    # Load from a Flax checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)):
                    # Load from a sharded Flax checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)
                    is_sharded = True
                elif is_safetensors_available() and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)
                ):
                    # Load from a safetensors checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)
                elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)):
                    # Load from a PyTorch checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)
                elif from_pt and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME)
                ):
                    # Load from a sharded pytorch checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME)
                    is_sharded = True
                # At this stage we don't have a weight file so we will raise an error.
                elif is_safetensors_available() and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)
                ):
                    # Load from a sharded safetensors checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)
                    is_sharded = True
                    raise NotImplementedError("Support for sharded checkpoints using safetensors is coming soon!")
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)):
                    raise EnvironmentError(
                        f"Error no file named {FLAX_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} "
                        "but there is a file for PyTorch weights. Use `from_pt=True` to load this model from those "
                        "weights."
                    )
                else:
                    raise EnvironmentError(
                        f"Error no file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory "
                        f"{pretrained_model_name_or_path}."
                    )
            elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
                archive_file = pretrained_model_name_or_path
                is_local = True
            elif is_remote_url(pretrained_model_name_or_path):
                filename = pretrained_model_name_or_path
                resolved_archive_file = download_url(pretrained_model_name_or_path)
            else:
                if from_pt:
                    filename = WEIGHTS_NAME
                else:
                    filename = FLAX_WEIGHTS_NAME

                try:
                    # Load from URL or cache if already cached
                    cached_file_kwargs = {
                        "cache_dir": cache_dir,
                        "force_download": force_download,
                        "proxies": proxies,
                        "resume_download": resume_download,
                        "local_files_only": local_files_only,
                        "token": token,
                        "user_agent": user_agent,
                        "revision": revision,
                        "subfolder": subfolder,
                        "_raise_exceptions_for_gated_repo": False,
                        "_raise_exceptions_for_missing_entries": False,
                        "_commit_hash": commit_hash,
                    }
                    resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)

                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                    if resolved_archive_file is None and filename == FLAX_WEIGHTS_NAME:
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME, **cached_file_kwargs
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True

                    # Maybe the checkpoint is pytorch sharded, we try to grab the pytorch index name in this case.
                    if resolved_archive_file is None and from_pt:
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **cached_file_kwargs
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True

                    # If we still haven't found anything, look for `safetensors`.
                    if resolved_archive_file is None:
                        # No support for sharded safetensors yet, so we'll raise an error if that's all we find.
                        filename = SAFE_WEIGHTS_NAME
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path, SAFE_WEIGHTS_NAME, **cached_file_kwargs
                        )

                    # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None
                    # result when internet is up, the repo and revision exist, but the file does not.
                    if resolved_archive_file is None:
                        # Otherwise, maybe there is a TF or Torch model file.  We try those to give a helpful error
                        # message.
                        has_file_kwargs = {
                            "revision": revision,
                            "proxies": proxies,
                            "token": token,
                            "cache_dir": cache_dir,
                            "local_files_only": local_files_only,
                        }
                        if has_file(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, **has_file_kwargs):
                            is_sharded = True
                            raise NotImplementedError(
                                "Support for sharded checkpoints using safetensors is coming soon!"
                            )
                        elif has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):
                            raise EnvironmentError(
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {FLAX_WEIGHTS_NAME} but there is a file for PyTorch weights. Use `from_pt=True` to"
                                " load this model from those weights."
                            )
                        elif has_file(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **has_file_kwargs):
                            raise EnvironmentError(
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {FLAX_WEIGHTS_INDEX_NAME} but there is a sharded file for PyTorch weights. Use"
                                " `from_pt=True` to load this model from those weights."
                            )
                        else:
                            raise EnvironmentError(
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}."
                            )
                except EnvironmentError:
                    # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
                    # to the original exception.
                    raise
                except Exception:
                    # For any other exception, we throw a generic error.
                    raise EnvironmentError(
                        f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
                        " from 'https://huggingface.co/models', make sure you don't have a local directory with the"
                        f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
                        f" directory containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}."
                    )

            if is_local:
                logger.info(f"loading weights file {archive_file}")
                resolved_archive_file = archive_file
                filename = resolved_archive_file.split(os.path.sep)[-1]
            else:
                logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}")
        else:
            resolved_archive_file = None

        # We'll need to download and cache each checkpoint shard if the checkpoint is sharded.
        if is_sharded:
            # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case.
            resolved_archive_file, _ = get_checkpoint_shard_files(
                pretrained_model_name_or_path,
                resolved_archive_file,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                resume_download=resume_download,
                local_files_only=local_files_only,
                token=token,
                user_agent=user_agent,
                revision=revision,
                subfolder=subfolder,
                _commit_hash=commit_hash,
            )

        safetensors_from_pt = False
        if filename == SAFE_WEIGHTS_NAME:
            with safe_open(resolved_archive_file, framework="flax") as f:
                safetensors_metadata = f.metadata()
            if safetensors_metadata is None or safetensors_metadata.get("format") not in ["pt", "tf", "flax"]:
                raise OSError(
                    f"The safetensors archive passed at {resolved_archive_file} does not contain the valid metadata."
                    " Make sure you save your model with the `save_pretrained` method."
                )
            safetensors_from_pt = safetensors_metadata.get("format") == "pt"

        # init random models
        model = cls(config, *model_args, _do_init=_do_init, **model_kwargs)

        if from_pt or safetensors_from_pt:
            state = load_pytorch_checkpoint_in_flax_state_dict(model, resolved_archive_file, is_sharded)
        else:
            if is_sharded:
                state = cls.load_flax_sharded_weights(resolved_archive_file)
            else:
                state = cls.load_flax_weights(resolved_archive_file)
            # make sure all arrays are stored as jnp.arrays
            # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4:
            # https://github.com/google/flax/issues/1261
            if _do_init:
                state = jax.tree_util.tree_map(jnp.array, state)
            else:
                # keep the params on CPU if we don't want to initialize
                state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.local_devices(backend="cpu")[0]), state)

        if "batch_stats" in state:  # if flax model contains batch norm layers
            # if model is base model only use model_prefix key
            if (
                cls.base_model_prefix not in dict(model.params_shape_tree["params"])
                and cls.base_model_prefix in state["params"]
            ):
                state["params"] = state["params"][cls.base_model_prefix]
                state["batch_stats"] = state["batch_stats"][cls.base_model_prefix]

            # if model is head model and we are loading weights from base model
            # we initialize new params dict with base_model_prefix
            if (
                cls.base_model_prefix in dict(model.params_shape_tree["params"])
                and cls.base_model_prefix not in state["params"]
            ):
                state = {
                    "params": {cls.base_model_prefix: state["params"]},
                    "batch_stats": {cls.base_model_prefix: state["batch_stats"]},
                }

        else:
            # if model is base model only use model_prefix key
            if cls.base_model_prefix not in dict(model.params_shape_tree) and cls.base_model_prefix in state:
                state = state[cls.base_model_prefix]

            # if model is head model and we are loading weights from base model
            # we initialize new params dict with base_model_prefix
            if cls.base_model_prefix in dict(model.params_shape_tree) and cls.base_model_prefix not in state:
                state = {cls.base_model_prefix: state}

        # flatten dicts
        state = flatten_dict(state)

        random_state = flatten_dict(unfreeze(model.params if _do_init else model.params_shape_tree))

        missing_keys = model.required_params - set(state.keys())
        unexpected_keys = set(state.keys()) - model.required_params

        # Disabling warning when porting pytorch weights to flax, flax does not uses num_batches_tracked
        for unexpected_key in unexpected_keys.copy():
            if "num_batches_tracked" in unexpected_key[-1]:
                unexpected_keys.remove(unexpected_key)

        if missing_keys and not _do_init:
            logger.warning(
                f"The checkpoint {pretrained_model_name_or_path} is missing required keys: {missing_keys}. "
                "Make sure to call model.init_weights to initialize the missing weights."
            )
            cls._missing_keys = missing_keys

        # Mistmatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not
        # matching the weights in the model.
        mismatched_keys = []
        for key in state.keys():
            if key in random_state and state[key].shape != random_state[key].shape:
                if ignore_mismatched_sizes:
                    mismatched_keys.append((key, state[key].shape, random_state[key].shape))
                    state[key] = random_state[key]
                else:
                    raise ValueError(
                        f"Trying to load the pretrained weight for {key} failed: checkpoint has shape "
                        f"{state[key].shape} which is incompatible with the model shape {random_state[key].shape}. "
                        "Using `ignore_mismatched_sizes=True` if you really want to load this checkpoint inside this "
                        "model."
                    )

        # add missing keys as random parameters if we are initializing
        if missing_keys and _do_init:
            for missing_key in missing_keys:
                state[missing_key] = random_state[missing_key]

        # remove unexpected keys to not be saved again
        for unexpected_key in unexpected_keys:
            del state[unexpected_key]

        if len(unexpected_keys) > 0:
            logger.warning(
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"
                " with another architecture (e.g. initializing a BertForSequenceClassification model from a"
                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
                " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
            )
        else:
            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")

        if len(missing_keys) > 0:
            logger.warning(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} 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."
            )
        elif len(mismatched_keys) == 0:
            logger.info(
                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
                f" was trained on, you can already use {model.__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 {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and 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."
            )

        # dictionary of key: dtypes for the model params
        param_dtypes = jax.tree_util.tree_map(lambda x: x.dtype, state)
        # extract keys of parameters not in jnp.float32
        fp16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.float16]
        bf16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.bfloat16]

        # raise a warning if any of the parameters are not in jnp.float32
        if len(fp16_params) > 0:
            logger.warning(
                f"Some of the weights of {model.__class__.__name__} were initialized in float16 precision from "
                f"the model checkpoint at {pretrained_model_name_or_path}:\n{fp16_params}\n"
                "You should probably UPCAST the model weights to float32 if this was not intended. "
                "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this."
            )

        if len(bf16_params) > 0:
            logger.warning(
                f"Some of the weights of {model.__class__.__name__} were initialized in bfloat16 precision from "
                f"the model checkpoint at {pretrained_model_name_or_path}:\n{bf16_params}\n"
                "You should probably UPCAST the model weights to float32 if this was not intended. "
                "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this."
            )

        # If it is a model with generation capabilities, attempt to load the generation config
        if model.can_generate():
            try:
                model.generation_config = GenerationConfig.from_pretrained(
                    pretrained_model_name_or_path,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    resume_download=resume_download,
                    proxies=proxies,
                    local_files_only=local_files_only,
                    token=token,
                    revision=revision,
                    subfolder=subfolder,
                    _from_auto=from_auto_class,
                    _from_pipeline=from_pipeline,
                    **kwargs,
                )
            except OSError:
                logger.info(
                    "Generation config file not found, using a generation config created from the model config."
                )
                pass

        if _do_init:
            # set correct parameters
            model.params = unflatten_dict(state)
            return model
        else:
            return model, unflatten_dict(state)

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        params=None,
        push_to_hub=False,
        max_shard_size="10GB",
        token: Optional[Union[str, bool]] = None,
        safe_serialization: bool = False,
        **kwargs,
    ):
        """
        Save a model and its configuration file to a directory, so that it can be re-loaded using the
        `[`~FlaxPreTrainedModel.from_pretrained`]` class method

        Arguments:
            save_directory (`str` or `os.PathLike`):
                Directory to which to save. Will be created if it doesn't exist.
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                namespace).
            max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).

                <Tip warning={true}>

                If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
                which will be bigger than `max_shard_size`.

                </Tip>

            token (`str` or `bool`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use
                the token generated when running `huggingface-cli login` (stored in `~/.huggingface`).
            kwargs (`Dict[str, Any]`, *optional*):
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
            safe_serialization (`bool`, *optional*, defaults to `False`):
                Whether to save the model using `safetensors` or through msgpack.
        """
        use_auth_token = kwargs.pop("use_auth_token", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if token is not None:
            kwargs["token"] = token

        if os.path.isfile(save_directory):
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
            return

        os.makedirs(save_directory, exist_ok=True)

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
            repo_id = self._create_repo(repo_id, **kwargs)
            files_timestamps = self._get_files_timestamps(save_directory)

        # get abs dir
        save_directory = os.path.abspath(save_directory)
        # save config as well
        self.config.architectures = [self.__class__.__name__[4:]]

        # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
        # loaded from the Hub.
        if self._auto_class is not None:
            custom_object_save(self, save_directory, config=self.config)

        self.config.save_pretrained(save_directory)
        if self.can_generate():
            self.generation_config.save_pretrained(save_directory)

        # save model
        weights_name = SAFE_WEIGHTS_NAME if safe_serialization else FLAX_WEIGHTS_NAME
        output_model_file = os.path.join(save_directory, weights_name)

        shards, index = flax_shard_checkpoint(params if params is not None else self.params, max_shard_size)
        # Clean the folder from a previous save
        for filename in os.listdir(save_directory):
            full_filename = os.path.join(save_directory, filename)
            weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "")
            if (
                filename.startswith(weights_no_suffix)
                and os.path.isfile(full_filename)
                and filename not in shards.keys()
            ):
                os.remove(full_filename)

        if index is None:
            if safe_serialization:
                params = params if params is not None else self.params
                flat_dict = flatten_dict(params, sep=".")
                safe_save_file(flat_dict, output_model_file, metadata={"format": "flax"})
            else:
                with open(output_model_file, "wb") as f:
                    params = params if params is not None else self.params
                    model_bytes = to_bytes(params)
                    f.write(model_bytes)

        else:
            save_index_file = os.path.join(save_directory, FLAX_WEIGHTS_INDEX_NAME)
            # Save the index as well
            with open(save_index_file, "w", encoding="utf-8") as f:
                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
                f.write(content)
            logger.info(
                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
                f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
                f"index located at {save_index_file}."
            )
            for shard_file, shard in shards.items():
                # the shard item are unflattened, to save them we need to flatten them again
                with open(os.path.join(save_directory, shard_file), mode="wb") as f:
                    params = unflatten_dict(shard, sep="/")
                    shard_bytes = to_bytes(params)
                    f.write(shard_bytes)

        logger.info(f"Model weights saved in {output_model_file}")

        if push_to_hub:
            self._upload_modified_files(
                save_directory,
                repo_id,
                files_timestamps,
                commit_message=commit_message,
                token=token,
            )

    @classmethod
    def register_for_auto_class(cls, auto_class="FlaxAutoModel"):
        """
        Register this class with a given auto class. This should only be used for custom models as the ones in the
        library are already mapped with an auto class.

        <Tip warning={true}>

        This API is experimental and may have some slight breaking changes in the next releases.

        </Tip>

        Args:
            auto_class (`str` or `type`, *optional*, defaults to `"FlaxAutoModel"`):
                The auto class to register this new model with.
        """
        if not isinstance(auto_class, str):
            auto_class = auto_class.__name__

        import transformers.models.auto as auto_module

        if not hasattr(auto_module, auto_class):
            raise ValueError(f"{auto_class} is not a valid auto class.")



SUMMARY:
This code provides the implementation of the Flax library for PyTorch Transformers models. Here is a concise summary:

**Main Classes and Functions**

1. `FlaxAutoModel`: The base class for all custom Flax models. It inherits from `flax modules.LazyModule`.
2. `register_for_auto_class`: A static method that registers a custom auto class with the Flax library.
3. `to_bytes` and `from_bytes`: Functions to convert model weights to bytes and vice versa.

**Model Registration**

The code defines several methods for registering models:

1. `update_repo_id`: Updates the repository ID associated with a given save directory.
2. `_create_repo`: Creates a new repository with the provided ID.
3. `_get_files_timestamps`: Retrieves the timestamp of each file in the specified directory.

**Saving and Loading Models**

The code provides methods to save and load models:

1. `save_pretrained` (from `self.config`): Saves the config and model weights to disk.
2. `save_model_directory` (optional argument): Saves the model weights to a custom directory.
3. `load_model_from_dir`: Loads the model from the specified directory.

**Additional Features**

The code defines several additional features:

1. Customization using Flax's `auto_class` system: Allows users to define their own auto classes for custom models.
2. Support for `auto_class` parameter in `FlaxAutoModel**: Enables the use of an external auto class with a given model.

Overall, this code provides a comprehensive implementation of the Flax library for PyTorch Transformers models, allowing users to define and save their own custom models.


key: chroma:document
string_value:
CODE:

def overwrite_call_docstring(model_class, docstring):
    # copy __call__ function to be sure docstring is changed only for this function
    model_class.__call__ = copy_func(model_class.__call__)
    # delete existing docstring
    model_class.__call__.__doc__ = None
    # set correct docstring


SUMMARY:
**Function Summary**

`overwrite_call_docstring(model_class, docstring)`: Updates the docstring of a given class's `__call__` method by preserving its original implementation and adding new start docstrings.

**Key Details:**

* Updates `model_class.__call__` by:
	+ Copying its original functionality (`copy_func`) to ensure docstring change is isolated.
	+ Deleting existing docstring (`model_class.__call__.__doc__ = None`).
	+ Setting the corrected docstring using `add_start_docstrings_to_model_forward`.

**Notable Dependencies:**

* `copy_func` function (not shown in this excerpt)
* `add_start_docstrings_to_model_forward` function (not shown in this excerpt)

This function modifies the `__call__` method of a given class to replace its original docstring with new start docstrings while preserving its functionality.


key: chroma:document
string_value:
CODE:

def append_call_sample_docstring(
    model_class, checkpoint, output_type, config_class, mask=None, revision=None, real_checkpoint=None
):
    model_class.__call__ = copy_func(model_class.__call__)
    model_class.__call__ = add_code_sample_docstrings(
        checkpoint=checkpoint,
        output_type=output_type,
        config_class=config_class,
        model_cls=model_class.__name__,
        revision=revision,
        real_checkpoint=real_checkpoint,


SUMMARY:
**Function Summary**

`append_call_sample_docstring(model_class, checkpoint, output_type, config_class, mask=None, revision=None, real_checkpoint=None)`

**Purpose and Functionality**: Updates the `__call__` method of a given model class by injecting custom docstrings for code sample generation.

**Key Implementation Details**:

*   Creates a copy of the original `__call__` method using `copy_func`.
*   Applies an external function `add_code_sample_docstrings` to modify the cloned method, passing in various parameters.
*   Returns the modified `model_class.__call__`.

**Notable Dependencies and Requirements**:

*   - `config_class`: instance of a configuration class (class reference)
*   - `mask`, `revision`: optional input variables without default values
*   - `real_checkpoint`: instance of a checkpoint object referencing `real_checkpoint`


key: chroma:document
string_value:
CODE:

def append_replace_return_docstrings(model_class, output_type, config_class):
    model_class.__call__ = copy_func(model_class.__call__)
    model_class.__call__ = replace_return_docstrings(
        output_type=output_type,
        config_class=config_class,


SUMMARY:
**Function Summary:**

* `append_replace_return_docstrings`
	+ Modifies the `__call__` method of a model class by:
		- Prepending a copy of its original method using `copy_func`.
		- Post-appending a new returned docstring replacement function generated by `replace_return_docstrings`.

**Key Implementation Details:**

* Appends and replaces the original method with a new one.
* Uses `output_type` and `config_class` to customize the generation of the new returned docstring.

**Notable Dependencies/Requirements:**

* None explicitly mentioned, but assumes availability of:
	+ `copy_func(model_class.__call__)`
	+ `replace_return_docstrings(output_type=None, config_class=None)`
	+ Support for modifying a function's `__call__` attribute.


key: chroma:document
string_value:
CODE:

class BatchFeature(UserDict):
    r"""
    Holds the output of the [`~SequenceFeatureExtractor.pad`] and feature extractor specific `__call__` methods.

    This class is derived from a python dictionary and can be used as a dictionary.

    Args:
        data (`dict`, *optional*):
            Dictionary of lists/arrays/tensors returned by the __call__/pad methods ('input_values', 'attention_mask',
            etc.).
        tensor_type (`Union[None, str, TensorType]`, *optional*):
            You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at
            initialization.
    """

    def __init__(self, data: Optional[Dict[str, Any]] = None, tensor_type: Union[None, str, TensorType] = None):
        super().__init__(data)
        self.convert_to_tensors(tensor_type=tensor_type)

    def __getitem__(self, item: str) -> Union[Any]:
        """
        If the key is a string, returns the value of the dict associated to `key` ('input_values', 'attention_mask',
        etc.).
        """
        if isinstance(item, str):
            return self.data[item]
        else:
            raise KeyError("Indexing with integers is not available when using Python based feature extractors")

    def __getattr__(self, item: str):
        try:
            return self.data[item]
        except KeyError:
            raise AttributeError

    def __getstate__(self):
        return {"data": self.data}

    def __setstate__(self, state):
        if "data" in state:
            self.data = state["data"]

    # Copied from transformers.tokenization_utils_base.BatchEncoding.keys
    def keys(self):
        return self.data.keys()

    # Copied from transformers.tokenization_utils_base.BatchEncoding.values
    def values(self):
        return self.data.values()

    # Copied from transformers.tokenization_utils_base.BatchEncoding.items
    def items(self):
        return self.data.items()

    def _get_is_as_tensor_fns(self, tensor_type: Optional[Union[str, TensorType]] = None):
        if tensor_type is None:
            return None, None

        # Convert to TensorType
        if not isinstance(tensor_type, TensorType):
            tensor_type = TensorType(tensor_type)

        # Get a function reference for the correct framework
        if tensor_type == TensorType.TENSORFLOW:
            if not is_tf_available():
                raise ImportError(
                    "Unable to convert output to TensorFlow tensors format, TensorFlow is not installed."
                )
            import tensorflow as tf

            as_tensor = tf.constant
            is_tensor = tf.is_tensor
        elif tensor_type == TensorType.PYTORCH:
            if not is_torch_available():
                raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.")
            import torch  # noqa

            def as_tensor(value):
                if isinstance(value, (list, tuple)) and len(value) > 0:
                    if isinstance(value[0], np.ndarray):
                        value = np.array(value)
                    elif (
                        isinstance(value[0], (list, tuple))
                        and len(value[0]) > 0
                        and isinstance(value[0][0], np.ndarray)
                    ):
                        value = np.array(value)
                if isinstance(value, np.ndarray):
                    return torch.from_numpy(value)
                else:
                    return torch.tensor(value)

            is_tensor = torch.is_tensor
        elif tensor_type == TensorType.JAX:
            if not is_flax_available():
                raise ImportError("Unable to convert output to JAX tensors format, JAX is not installed.")
            import jax.numpy as jnp  # noqa: F811

            as_tensor = jnp.array
            is_tensor = is_jax_tensor
        else:

            def as_tensor(value, dtype=None):
                if isinstance(value, (list, tuple)) and isinstance(value[0], (list, tuple, np.ndarray)):
                    value_lens = [len(val) for val in value]
                    if len(set(value_lens)) > 1 and dtype is None:
                        # we have a ragged list so handle explicitly
                        value = as_tensor([np.asarray(val) for val in value], dtype=object)
                return np.asarray(value, dtype=dtype)

            is_tensor = is_numpy_array
        return is_tensor, as_tensor

    def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):
        """
        Convert the inner content to tensors.

        Args:
            tensor_type (`str` or [`~utils.TensorType`], *optional*):
                The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If
                `None`, no modification is done.
        """
        if tensor_type is None:
            return self

        is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)

        # Do the tensor conversion in batch
        for key, value in self.items():
            try:
                if not is_tensor(value):
                    tensor = as_tensor(value)

                    self[key] = tensor
            except:  # noqa E722
                if key == "overflowing_values":
                    raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")
                raise ValueError(
                    "Unable to create tensor, you should probably activate padding "
                    "with 'padding=True' to have batched tensors with the same length."
                )

        return self

    def to(self, *args, **kwargs) -> "BatchFeature":
        """
        Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
        different `dtypes` and sending the `BatchFeature` to a different `device`.

        Args:
            args (`Tuple`):
                Will be passed to the `to(...)` function of the tensors.
            kwargs (`Dict`, *optional*):
                Will be passed to the `to(...)` function of the tensors.
                To enable asynchronous data transfer, set the `non_blocking` flag in `kwargs` (defaults to `False`).

        Returns:
            [`BatchFeature`]: The same instance after modification.
        """
        requires_backends(self, ["torch"])
        import torch  # noqa

        new_data = {}
        device = kwargs.get("device")
        non_blocking = kwargs.get("non_blocking", False)
        # Check if the args are a device or a dtype
        if device is None and len(args) > 0:
            # device should be always the first argument
            arg = args[0]
            if is_torch_dtype(arg):
                # The first argument is a dtype
                pass
            elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
                device = arg
            else:
                # it's something else
                raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
        # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
        for k, v in self.items():
            # check if v is a floating point
            if isinstance(v, torch.Tensor) and torch.is_floating_point(v):
                # cast and send to device
                new_data[k] = v.to(*args, **kwargs)
            elif isinstance(v, torch.Tensor) and device is not None:
                new_data[k] = v.to(device=device, non_blocking=non_blocking)
            else:
                new_data[k] = v
        self.data = new_data


SUMMARY:
**BatchFeature Class Summary**

The `BatchFeature` class extends `UserDict` to provide a dictionary-like interface for storing and manipulating output data from sequence feature extractors. It supports conversion to tensors, indexing, and sententious tensor manipulation.

**Main Purpose and Functionality:**

*   Represents a batch of output data from sequence feature extractors.
*   Provides a dictionary-like interface with methods like `__getitem__`, `__getattr__`, and `items`.
*   Supports conversion of inner content to tensors based on the specified `tensor_type`.

**Key Implementation Details:**

*   The `_get_is_as_tensor_fns` method returns functions for converting data to tensors based on the specified `tensor_type`.
*   The `convert_to_tensors` method converts the inner content to tensors using the `_get_is_as_tensor_fns` function.
*   The `to` method sends values to a specific device, supporting different dtypes and devices.

**Notable Dependencies or Requirements:**

*   `pytorch` is required for tensor conversion and sententious manipulation.
*   `tensorflow` is required for TensorFlow tensors format conversion (optional).
*   `jax.numpy` is required for JAX tensors format conversion (optional).


key: chroma:document
string_value:
CODE:

class FeatureExtractionMixin(PushToHubMixin):
    """
    This is a feature extraction mixin used to provide saving/loading functionality for sequential and image feature
    extractors.
    """

    _auto_class = None

    def __init__(self, **kwargs):
        """Set elements of `kwargs` as attributes."""
        # Pop "processor_class" as it should be saved as private attribute
        self._processor_class = kwargs.pop("processor_class", None)
        # Additional attributes without default values
        for key, value in kwargs.items():
            try:
                setattr(self, key, value)
            except AttributeError as err:
                logger.error(f"Can't set {key} with value {value} for {self}")
                raise err

    def _set_processor_class(self, processor_class: str):
        """Sets processor class as an attribute."""
        self._processor_class = processor_class

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        cache_dir: Optional[Union[str, os.PathLike]] = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: Optional[Union[str, bool]] = None,
        revision: str = "main",
        **kwargs,
    ):
        r"""
        Instantiate a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a feature extractor, *e.g.* a
        derived class of [`SequenceFeatureExtractor`].

        Args:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                This can be either:

                - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
                  huggingface.co.
                - a path to a *directory* containing a feature extractor file saved using the
                  [`~feature_extraction_utils.FeatureExtractionMixin.save_pretrained`] method, e.g.,
                  `./my_model_directory/`.
                - a path or url to a saved feature extractor JSON *file*, e.g.,
                  `./my_model_directory/preprocessor_config.json`.
            cache_dir (`str` or `os.PathLike`, *optional*):
                Path to a directory in which a downloaded pretrained model feature extractor should be cached if the
                standard cache should not be used.
            force_download (`bool`, *optional*, defaults to `False`):
                Whether or not to force to (re-)download the feature extractor files and override the cached versions
                if they exist.
            resume_download:
                Deprecated and ignored. All downloads are now resumed by default when possible.
                Will be removed in v5 of Transformers.
            proxies (`Dict[str, str]`, *optional*):
                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
                'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
            token (`str` or `bool`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use
                the token generated when running `huggingface-cli login` (stored in `~/.huggingface`).
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
                identifier allowed by git.


                <Tip>

                To test a pull request you made on the Hub, you can pass `revision="refs/pr/<pr_number>"`.

                </Tip>

            return_unused_kwargs (`bool`, *optional*, defaults to `False`):
                If `False`, then this function returns just the final feature extractor object. If `True`, then this
                functions returns a `Tuple(feature_extractor, unused_kwargs)` where *unused_kwargs* is a dictionary
                consisting of the key/value pairs whose keys are not feature extractor attributes: i.e., the part of
                `kwargs` which has not been used to update `feature_extractor` and is otherwise ignored.
            kwargs (`Dict[str, Any]`, *optional*):
                The values in kwargs of any keys which are feature extractor attributes will be used to override the
                loaded values. Behavior concerning key/value pairs whose keys are *not* feature extractor attributes is
                controlled by the `return_unused_kwargs` keyword parameter.

        Returns:
            A feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`].

        Examples:

        ```python
        # We can't instantiate directly the base class *FeatureExtractionMixin* nor *SequenceFeatureExtractor* so let's show the examples on a
        # derived class: *Wav2Vec2FeatureExtractor*
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
            "facebook/wav2vec2-base-960h"
        )  # Download feature_extraction_config from huggingface.co and cache.
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
            "./test/saved_model/"
        )  # E.g. feature_extractor (or model) was saved using *save_pretrained('./test/saved_model/')*
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("./test/saved_model/preprocessor_config.json")
        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
            "facebook/wav2vec2-base-960h", return_attention_mask=False, foo=False
        )
        assert feature_extractor.return_attention_mask is False
        feature_extractor, unused_kwargs = Wav2Vec2FeatureExtractor.from_pretrained(
            "facebook/wav2vec2-base-960h", return_attention_mask=False, foo=False, return_unused_kwargs=True
        )
        assert feature_extractor.return_attention_mask is False
        assert unused_kwargs == {"foo": False}
        ```"""
        kwargs["cache_dir"] = cache_dir
        kwargs["force_download"] = force_download
        kwargs["local_files_only"] = local_files_only
        kwargs["revision"] = revision

        use_auth_token = kwargs.pop("use_auth_token", None)
        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if token is not None:
            kwargs["token"] = token

        feature_extractor_dict, kwargs = cls.get_feature_extractor_dict(pretrained_model_name_or_path, **kwargs)

        return cls.from_dict(feature_extractor_dict, **kwargs)

    def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
        """
        Save a feature_extractor object to the directory `save_directory`, so that it can be re-loaded using the
        [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`] class method.

        Args:
            save_directory (`str` or `os.PathLike`):
                Directory where the feature extractor JSON file will be saved (will be created if it does not exist).
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                namespace).
            kwargs (`Dict[str, Any]`, *optional*):
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
        """
        use_auth_token = kwargs.pop("use_auth_token", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if kwargs.get("token", None) is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            kwargs["token"] = use_auth_token

        if os.path.isfile(save_directory):
            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")

        os.makedirs(save_directory, exist_ok=True)

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
            repo_id = self._create_repo(repo_id, **kwargs)
            files_timestamps = self._get_files_timestamps(save_directory)

        # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be
        # loaded from the Hub.
        if self._auto_class is not None:
            custom_object_save(self, save_directory, config=self)

        # If we save using the predefined names, we can load using `from_pretrained`
        output_feature_extractor_file = os.path.join(save_directory, FEATURE_EXTRACTOR_NAME)

        self.to_json_file(output_feature_extractor_file)
        logger.info(f"Feature extractor saved in {output_feature_extractor_file}")

        if push_to_hub:
            self._upload_modified_files(
                save_directory,
                repo_id,
                files_timestamps,
                commit_message=commit_message,
                token=kwargs.get("token"),
            )

        return [output_feature_extractor_file]

    @classmethod
    def get_feature_extractor_dict(
        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
        """
        From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a
        feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`] using `from_dict`.

        Parameters:
            pretrained_model_name_or_path (`str` or `os.PathLike`):
                The identifier of the pre-trained checkpoint from which we want the dictionary of parameters.

        Returns:
            `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the feature extractor object.
        """
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", None)
        proxies = kwargs.pop("proxies", None)
        subfolder = kwargs.pop("subfolder", None)
        token = kwargs.pop("token", None)
        use_auth_token = kwargs.pop("use_auth_token", None)
        local_files_only = kwargs.pop("local_files_only", False)
        revision = kwargs.pop("revision", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        from_pipeline = kwargs.pop("_from_pipeline", None)
        from_auto_class = kwargs.pop("_from_auto", False)

        user_agent = {"file_type": "feature extractor", "from_auto_class": from_auto_class}
        if from_pipeline is not None:
            user_agent["using_pipeline"] = from_pipeline

        if is_offline_mode() and not local_files_only:
            logger.info("Offline mode: forcing local_files_only=True")
            local_files_only = True

        pretrained_model_name_or_path = str(pretrained_model_name_or_path)
        is_local = os.path.isdir(pretrained_model_name_or_path)
        if os.path.isdir(pretrained_model_name_or_path):
            feature_extractor_file = os.path.join(pretrained_model_name_or_path, FEATURE_EXTRACTOR_NAME)
        if os.path.isfile(pretrained_model_name_or_path):
            resolved_feature_extractor_file = pretrained_model_name_or_path
            is_local = True
        elif is_remote_url(pretrained_model_name_or_path):
            feature_extractor_file = pretrained_model_name_or_path
            resolved_feature_extractor_file = download_url(pretrained_model_name_or_path)
        else:
            feature_extractor_file = FEATURE_EXTRACTOR_NAME
            try:
                # Load from local folder or from cache or download from model Hub and cache
                resolved_feature_extractor_file = cached_file(
                    pretrained_model_name_or_path,
                    feature_extractor_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                    subfolder=subfolder,
                    token=token,
                    user_agent=user_agent,
                    revision=revision,
                )
            except EnvironmentError:
                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to
                # the original exception.
                raise
            except Exception:
                # For any other exception, we throw a generic error.
                raise EnvironmentError(
                    f"Can't load feature extractor for '{pretrained_model_name_or_path}'. If you were trying to load"
                    " it from 'https://huggingface.co/models', make sure you don't have a local directory with the"
                    f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
                    f" directory containing a {FEATURE_EXTRACTOR_NAME} file"
                )

        try:
            # Load feature_extractor dict
            with open(resolved_feature_extractor_file, "r", encoding="utf-8") as reader:
                text = reader.read()
            feature_extractor_dict = json.loads(text)

        except json.JSONDecodeError:
            raise EnvironmentError(
                f"It looks like the config file at '{resolved_feature_extractor_file}' is not a valid JSON file."
            )

        if is_local:
            logger.info(f"loading configuration file {resolved_feature_extractor_file}")
        else:
            logger.info(
                f"loading configuration file {feature_extractor_file} from cache at {resolved_feature_extractor_file}"
            )

        if not is_local:
            if "auto_map" in feature_extractor_dict:
                feature_extractor_dict["auto_map"] = add_model_info_to_auto_map(
                    feature_extractor_dict["auto_map"], pretrained_model_name_or_path
                )
            if "custom_pipelines" in feature_extractor_dict:
                feature_extractor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines(
                    feature_extractor_dict["custom_pipelines"], pretrained_model_name_or_path
                )

        return feature_extractor_dict, kwargs

    @classmethod
    def from_dict(cls, feature_extractor_dict: Dict[str, Any], **kwargs) -> PreTrainedFeatureExtractor:
        """
        Instantiates a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a Python dictionary of
        parameters.

        Args:
            feature_extractor_dict (`Dict[str, Any]`):
                Dictionary that will be used to instantiate the feature extractor object. Such a dictionary can be
                retrieved from a pretrained checkpoint by leveraging the
                [`~feature_extraction_utils.FeatureExtractionMixin.to_dict`] method.
            kwargs (`Dict[str, Any]`):
                Additional parameters from which to initialize the feature extractor object.

        Returns:
            [`~feature_extraction_utils.FeatureExtractionMixin`]: The feature extractor object instantiated from those
            parameters.
        """
        return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)

        # Update feature_extractor with kwargs if needed
        to_remove = []
        for key, value in kwargs.items():
            if key in feature_extractor_dict:
                feature_extractor_dict[key] = value
                to_remove.append(key)
        for key in to_remove:
            kwargs.pop(key, None)

        feature_extractor = cls(**feature_extractor_dict)

        logger.info(f"Feature extractor {feature_extractor}")
        if return_unused_kwargs:
            return feature_extractor, kwargs
        else:
            return feature_extractor

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes this instance to a Python dictionary. Returns:
            `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
        """
        output = copy.deepcopy(self.__dict__)
        output["feature_extractor_type"] = self.__class__.__name__
        if "mel_filters" in output:
            del output["mel_filters"]
        if "window" in output:
            del output["window"]
        return output

    @classmethod
    def from_json_file(cls, json_file: Union[str, os.PathLike]) -> PreTrainedFeatureExtractor:
        """
        Instantiates a feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`] from the path to
        a JSON file of parameters.

        Args:
            json_file (`str` or `os.PathLike`):
                Path to the JSON file containing the parameters.

        Returns:
            A feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`]: The feature_extractor
            object instantiated from that JSON file.
        """
        with open(json_file, "r", encoding="utf-8") as reader:
            text = reader.read()
        feature_extractor_dict = json.loads(text)
        return cls(**feature_extractor_dict)

    def to_json_string(self) -> str:
        """
        Serializes this instance to a JSON string.

        Returns:
            `str`: String containing all the attributes that make up this feature_extractor instance in JSON format.
        """
        dictionary = self.to_dict()

        for key, value in dictionary.items():
            if isinstance(value, np.ndarray):
                dictionary[key] = value.tolist()

        # make sure private name "_processor_class" is correctly
        # saved as "processor_class"
        _processor_class = dictionary.pop("_processor_class", None)
        if _processor_class is not None:
            dictionary["processor_class"] = _processor_class

        return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"

    def to_json_file(self, json_file_path: Union[str, os.PathLike]):
        """
        Save this instance to a JSON file.

        Args:
            json_file_path (`str` or `os.PathLike`):
                Path to the JSON file in which this feature_extractor instance's parameters will be saved.
        """
        with open(json_file_path, "w", encoding="utf-8") as writer:
            writer.write(self.to_json_string())

    def __repr__(self):
        return f"{self.__class__.__name__} {self.to_json_string()}"

    @classmethod
    def register_for_auto_class(cls, auto_class="AutoFeatureExtractor"):
        """
        Register this class with a given auto class. This should only be used for custom feature extractors as the ones
        in the library are already mapped with `AutoFeatureExtractor`.

        <Tip warning={true}>

        This API is experimental and may have some slight breaking changes in the next releases.

        </Tip>

        Args:
            auto_class (`str` or `type`, *optional*, defaults to `"AutoFeatureExtractor"`):
                The auto class to register this new feature extractor with.
        """
        if not isinstance(auto_class, str):
            auto_class = auto_class.__name__

        import transformers.models.auto as auto_module

        if not hasattr(auto_module, auto_class):
            raise ValueError(f"{auto_class} is not a valid auto class.")



SUMMARY:
The code provided appears to be part of the PyTorch Transformers library for feature extraction. It defines several classes and functions related to creating, serializing, and registering feature extractors.

**Key Classes:**

1. `PreTrainedFeatureExtractor`: The main class that handles all the functionality related to feature extraction. It provides methods for loading pre-trained models, creating instances from a Python dictionary or JSON file, converting objects to dictionaries, and more.
2. `AutoFeatureExtractor`: An abstract base class defined in the library, which serves as an interface for custom feature extractors.

**Key Functions:**

1. `from_dict()`: Creates an instance of `PreTrainedFeatureExtractor` from a Python dictionary.
2. `to_dict()` and `to_json_string()`: Serializes an object to a dictionary or JSON string respectively.
3. `register_for_auto_class()`: Registers this feature extractor class with the specified auto class.

**Additional Notes:**

* The code uses type hints, which indicates that it is designed to be compatible with Python 3.5 or later versions.
* The library depends on other libraries such as PyTorch and json.
* It's worth noting that the code might not cover all edge cases and potential issues related to data structures and class hierarchies.

Overall, the code appears to be robust and well-structured, but it may require additional testing and validation to ensure its correctness for specific use cases. 

Here is an example of how you could create an instance from a dictionary using this library:

```python
from transformers import AutoFeatureExtractor

# define the dictionary containing pre-trained model features
feature_extractor_dict = {
    "class": "AudioFeatureExtractor",
    # other key-value pairs would be added as necessary...
}

feature_extractor = AutoFeatureExtractor.from_dict(feature_extractor_dict)
```

And here's an example of how you could register this feature extractor with a custom auto class:

```python
from transformers.models.auto import AutoFeatureExtractor

# define the custom auto class and its features
class MyAutoFeatureExtractor(AutoFeatureExtractor):
    pass  # implement relevant methods for your use case...

AutoFeatureExtractor.register_for_auto_class(auto_class="MyAutoFeatureExtractor")

# now you can create an instance with the 'new' name:
feature_extractor = AutoFeatureExtractor.from_dict(feature_extractor_dict)
```


key: chroma:document
string_value:
CODE:

class SequenceFeatureExtractor(FeatureExtractionMixin):
    """
    This is a general feature extraction class for speech recognition.

    Args:
        feature_size (`int`):
            The feature dimension of the extracted features.
        sampling_rate (`int`):
            The sampling rate at which the audio files should be digitalized expressed in hertz (Hz).
        padding_value (`float`):
            The value that is used to fill the padding values / vectors.
    """

    def __init__(self, feature_size: int, sampling_rate: int, padding_value: float, **kwargs):
        self.feature_size = feature_size
        self.sampling_rate = sampling_rate
        self.padding_value = padding_value

        self.padding_side = kwargs.pop("padding_side", "right")
        self.return_attention_mask = kwargs.pop("return_attention_mask", True)

        super().__init__(**kwargs)

    def pad(
        self,
        processed_features: Union[
            BatchFeature,
            List[BatchFeature],
            Dict[str, BatchFeature],
            Dict[str, List[BatchFeature]],
            List[Dict[str, BatchFeature]],
        ],
        padding: Union[bool, str, PaddingStrategy] = True,
        max_length: Optional[int] = None,
        truncation: bool = False,
        pad_to_multiple_of: Optional[int] = None,
        return_attention_mask: Optional[bool] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
    ) -> BatchFeature:
        """
        Pad input values / input vectors or a batch of input values / input vectors up to predefined length or to the
        max sequence length in the batch.

        Padding side (left/right) padding values are defined at the feature extractor level (with `self.padding_side`,
        `self.padding_value`)

        <Tip>

        If the `processed_features` passed are dictionary of numpy arrays, PyTorch tensors or TensorFlow tensors, the
        result will use the same type unless you provide a different tensor type with `return_tensors`. In the case of
        PyTorch tensors, you will lose the specific device of your tensors however.

        </Tip>

        Args:
            processed_features ([`BatchFeature`], list of [`BatchFeature`], `Dict[str, List[float]]`, `Dict[str, List[List[float]]` or `List[Dict[str, List[float]]]`):
                Processed inputs. Can represent one input ([`BatchFeature`] or `Dict[str, List[float]]`) or a batch of
                input values / vectors (list of [`BatchFeature`], *Dict[str, List[List[float]]]* or *List[Dict[str,
                List[float]]]*) so you can use this method during preprocessing as well as in a PyTorch Dataloader
                collate function.

                Instead of `List[float]` you can have tensors (numpy arrays, PyTorch tensors or TensorFlow tensors),
                see the note above for the return type.
            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
                Select a strategy to pad the returned sequences (according to the model's padding side and padding
                index) among:

                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
                  sequence if provided).
                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
                  acceptable input length for the model if that argument is not provided.
                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
                  lengths).
            max_length (`int`, *optional*):
                Maximum length of the returned list and optionally padding length (see above).
            truncation (`bool`):
                Activates truncation to cut input sequences longer than `max_length` to `max_length`.
            pad_to_multiple_of (`int`, *optional*):
                If set will pad the sequence to a multiple of the provided value.

                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
                `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128.
            return_attention_mask (`bool`, *optional*):
                Whether to return the attention mask. If left to the default, will return the attention mask according
                to the specific feature_extractor's default.

                [What are attention masks?](../glossary#attention-mask)
            return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors instead of list of python integers. Acceptable values are:

                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return Numpy `np.ndarray` objects.
        """
        # If we have a list of dicts, let's convert it in a dict of lists
        # We do this to allow using this method as a collate_fn function in PyTorch Dataloader
        if isinstance(processed_features, (list, tuple)) and isinstance(processed_features[0], (dict, BatchFeature)):
            processed_features = {
                key: [example[key] for example in processed_features] for key in processed_features[0].keys()
            }

        # The model's main input name, usually `input_values`, has be passed for padding
        if self.model_input_names[0] not in processed_features:
            raise ValueError(
                "You should supply an instance of `transformers.BatchFeature` or list of `transformers.BatchFeature`"
                f" to this method that includes {self.model_input_names[0]}, but you provided"
                f" {list(processed_features.keys())}"
            )

        required_input = processed_features[self.model_input_names[0]]
        return_attention_mask = (
            return_attention_mask if return_attention_mask is not None else self.return_attention_mask
        )

        if len(required_input) == 0:
            if return_attention_mask:
                processed_features["attention_mask"] = []
            return processed_features

        # If we have PyTorch/TF tensors or lists as inputs, we cast them as Numpy arrays
        # and rebuild them afterwards if no return_tensors is specified
        # Note that we lose the specific device the tensor may be on for PyTorch

        first_element = required_input[0]
        if isinstance(first_element, (list, tuple)):
            # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element.
            index = 0
            while len(required_input[index]) == 0:
                index += 1
            if index < len(required_input):
                first_element = required_input[index][0]

        if return_tensors is None:
            if is_tf_tensor(first_element):
                return_tensors = "tf"
            elif is_torch_tensor(first_element):
                return_tensors = "pt"
            elif isinstance(first_element, (int, float, list, tuple, np.ndarray)):
                return_tensors = "np"
            else:
                raise ValueError(
                    f"type of {first_element} unknown: {type(first_element)}. "
                    "Should be one of a python, numpy, pytorch or tensorflow object."
                )

        for key, value in processed_features.items():
            if isinstance(value[0], (int, float)):
                processed_features[key] = to_numpy(value)
            else:
                processed_features[key] = [to_numpy(v) for v in value]

        # Convert padding_strategy in PaddingStrategy
        padding_strategy = self._get_padding_strategies(padding=padding, max_length=max_length)

        required_input = processed_features[self.model_input_names[0]]

        batch_size = len(required_input)
        if not all(len(v) == batch_size for v in processed_features.values()):
            raise ValueError("Some items in the output dictionary have a different batch size than others.")

        truncated_inputs = []
        for i in range(batch_size):
            inputs = {k: v[i] for k, v in processed_features.items()}
            # truncation
            inputs_slice = self._truncate(
                inputs,
                max_length=max_length,
                pad_to_multiple_of=pad_to_multiple_of,
                truncation=truncation,
            )
            truncated_inputs.append(inputs_slice)

        if padding_strategy == PaddingStrategy.LONGEST:
            # make sure that `max_length` cannot be longer than the longest truncated length
            max_length = max(len(input_slice[self.model_input_names[0]]) for input_slice in truncated_inputs)
            padding_strategy = PaddingStrategy.MAX_LENGTH

        batch_outputs = {}
        for i in range(batch_size):
            # padding
            outputs = self._pad(
                truncated_inputs[i],
                max_length=max_length,
                padding_strategy=padding_strategy,
                pad_to_multiple_of=pad_to_multiple_of,
                return_attention_mask=return_attention_mask,
            )

            for key, value in outputs.items():
                if key not in batch_outputs:
                    batch_outputs[key] = []
                if value.dtype is np.dtype(np.float64):
                    value = value.astype(np.float32)
                batch_outputs[key].append(value)

        return BatchFeature(batch_outputs, tensor_type=return_tensors)

    def _pad(
        self,
        processed_features: Union[Dict[str, np.ndarray], BatchFeature],
        max_length: Optional[int] = None,
        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
        pad_to_multiple_of: Optional[int] = None,
        return_attention_mask: Optional[bool] = None,
    ) -> dict:
        """
        Pad inputs (on left/right and up to predefined length or max length in the batch)

        Args:
            processed_features (`Union[Dict[str, np.ndarray], BatchFeature]`):
                Dictionary of input values (`np.ndarray[float]`) / input vectors (`List[np.ndarray[float]]`) or batch
                of inputs values (`List[np.ndarray[int]]`) / input vectors (`List[np.ndarray[int]]`)
            max_length (`int`, *optional*):
                Maximum length of the returned list and optionally padding length (see below)
            padding_strategy (`PaddingStrategy`, *optional*, default to `PaddingStrategy.DO_NOT_PAD`):
                PaddingStrategy to use for padding.

                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
                - PaddingStrategy.DO_NOT_PAD: Do not pad
                The feature_extractor padding sides are defined in self.padding_side:

                    - 'left': pads on the left of the sequences
                    - 'right': pads on the right of the sequences
            pad_to_multiple_of (`int`, *optional*):
                Integer if set will pad the sequence to a multiple of the provided value. This is especially useful to
                enable the use of Tensor Core on NVIDIA hardware with compute capability `>= 7.5` (Volta), or on TPUs
                which benefit from having sequence lengths be a multiple of 128.
            return_attention_mask (`bool`, *optional*):
                Set to False to avoid returning attention mask (default: set to model specifics)
        """
        required_input = processed_features[self.model_input_names[0]]

        if padding_strategy == PaddingStrategy.LONGEST:
            max_length = len(required_input)

        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
            max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) < max_length

        if return_attention_mask and "attention_mask" not in processed_features:
            processed_features["attention_mask"] = np.ones(len(required_input), dtype=np.int32)

        if needs_to_be_padded:
            difference = max_length - len(required_input)
            if self.padding_side == "right":
                if return_attention_mask:
                    processed_features["attention_mask"] = np.pad(
                        processed_features["attention_mask"], (0, difference)
                    )
                padding_shape = ((0, difference), (0, 0)) if self.feature_size > 1 else (0, difference)
                processed_features[self.model_input_names[0]] = np.pad(
                    required_input, padding_shape, "constant", constant_values=self.padding_value
                )
            elif self.padding_side == "left":
                if return_attention_mask:
                    processed_features["attention_mask"] = np.pad(
                        processed_features["attention_mask"], (difference, 0)
                    )
                padding_shape = ((difference, 0), (0, 0)) if self.feature_size > 1 else (difference, 0)
                processed_features[self.model_input_names[0]] = np.pad(
                    required_input, padding_shape, "constant", constant_values=self.padding_value
                )
            else:
                raise ValueError("Invalid padding strategy:" + str(self.padding_side))

        return processed_features

    def _truncate(
        self,
        processed_features: Union[Dict[str, np.ndarray], BatchFeature],
        max_length: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = None,
        truncation: Optional[bool] = None,
    ):
        """
        Truncate inputs to predefined length or max length in the batch

        Args:
            processed_features(`Union[Dict[str, np.ndarray], BatchFeature]`):
                Dictionary of input values (`np.ndarray[float]`) / input vectors (`List[np.ndarray[float]]`) or batch
                of inputs values (`List[np.ndarray[int]]`) / input vectors (`List[np.ndarray[int]]`)
            max_length (`int`, *optional*):
                maximum length of the returned list and optionally padding length (see below)
            pad_to_multiple_of (`int`, *optional*) :
                Integer if set will pad the sequence to a multiple of the provided value. This is especially useful to
                enable the use of Tensor Core on NVIDIA hardware with compute capability `>= 7.5` (Volta), or on TPUs
                which benefit from having sequence lengths be a multiple of 128.
            truncation (`bool`, *optional*):
                Activates truncation to cut input sequences longer than `max_length` to `max_length`.
        """
        if not truncation:
            return processed_features
        elif truncation and max_length is None:
            raise ValueError("When setting ``truncation=True``, make sure that ``max_length`` is defined.")

        required_input = processed_features[self.model_input_names[0]]

        # find `max_length` that fits `pad_to_multiple_of`
        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
            max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of

        needs_to_be_truncated = len(required_input) > max_length

        if needs_to_be_truncated:
            processed_features[self.model_input_names[0]] = processed_features[self.model_input_names[0]][:max_length]
            if "attention_mask" in processed_features:
                processed_features["attention_mask"] = processed_features["attention_mask"][:max_length]

        return processed_features

    def _get_padding_strategies(self, padding=False, max_length=None):
        """
        Find the correct padding strategy
        """

        # Get padding strategy
        if padding is not False:
            if padding is True:
                padding_strategy = PaddingStrategy.LONGEST  # Default to pad to the longest sequence in the batch
            elif not isinstance(padding, PaddingStrategy):
                padding_strategy = PaddingStrategy(padding)
            elif isinstance(padding, PaddingStrategy):
                padding_strategy = padding
        else:
            padding_strategy = PaddingStrategy.DO_NOT_PAD

        # Set max length if needed
        if max_length is None:
            if padding_strategy == PaddingStrategy.MAX_LENGTH:
                raise ValueError(
                    f"When setting ``padding={PaddingStrategy.MAX_LENGTH}``, make sure that max_length is defined"
                )

        # Test if we have a padding value
        if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.padding_value is None):
            raise ValueError(
                "Asking to pad but the feature_extractor does not have a padding value. Please select a value to use"
                " as `padding_value`. For example: `feature_extractor.padding_value = 0.0`."
            )



SUMMARY:
Here is a concise summary:

The provided code appears to be part of a class definition for a feature extractor in the context of natural language processing (NLP) or text analysis. The class seems to be designed to handle tasks such as padding, truncation, and masking of sequences.

Key features include:
- **Padding**: allows for padding the input sequences with a specified value (default is 0.0)
- **Truncation**: can truncate input sequences to a maximum length
- **Masking attention mask**: generates an attention mask if it's not already present

The class has methods to perform these tasks, including `_truncate`, `_pad`, and `_get_padding_strategies`. It also supports different padding strategies such as default padding to the longest sequence (`PaddingStrategy.MAX_LENGTH`), padding all sequences equally using `PaddingStrategy.LONGEST`, or no padding at all (`PaddingStrategy.DO_NOT_PAD`). 

Overall, this feature extractor seems designed to be flexible and perform standard operations in NLP processing.


key: chroma:document
string_value:
CODE:

def _compute_default_rope_parameters(
    config: Optional[PretrainedConfig] = None,
    device: Optional["torch.device"] = None,
    seq_len: Optional[int] = None,
    **rope_kwargs,
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies according to the original RoPE implementation
    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length. Unused for this type of RoPE.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
    """
    if config is not None and len(rope_kwargs) > 0:
        raise ValueError(
            "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
            f"`_compute_default_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
        )
    if len(rope_kwargs) > 0:
        base = rope_kwargs["base"]
        dim = rope_kwargs["dim"]
    elif config is not None:
        base = config.rope_theta
        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        dim = int(head_dim * partial_rotary_factor)

    attention_factor = 1.0  # Unused in this type of RoPE

    # Compute the inverse frequencies
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))


SUMMARY:
**Summary:**

The `_compute_default_rope_parameters` function computes the inverse frequencies for RoPE embeddings and returns them along with a post-processing scaling factor. The primary purpose is to calculate the RoPE parameters required for token embedding.

**Functionality:**

- Calculates the `inv_freq` tensor using the base value and dims obtained from either user-provided `rope_kwargs` or the provided `config`.
- Returns a tuple containing the calculated `inv_freq` and a post-processing scaling factor (`attention_factor`).

**Key Implementation Details:**

- Uses `torch.arange` to generate an array of indices for computing inverse frequencies.
- Applies `partial_rotary_factor` to calculate dim based on `config` if necessary.

**Notable Dependencies/Requirements:**

- torch.device (`device`) for computation.
- torch Tensor (`inv_freq`) for storing computed values.
- transformers.PretrainedConfig (`config`) providing base parameter.
- optional rope_kwargs dictionaries compatibility. 

**External Dependencies:**


key: chroma:document
string_value:
CODE:

def _compute_linear_scaling_rope_parameters(
    config: Optional[PretrainedConfig] = None,
    device: Optional["torch.device"] = None,
    seq_len: Optional[int] = None,
    **rope_kwargs,
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev
    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length. Unused for this type of RoPE.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
    """
    if config is not None and len(rope_kwargs) > 0:
        raise ValueError(
            "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
            f"`_compute_linear_scaling_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
        )
    if len(rope_kwargs) > 0:
        factor = rope_kwargs["factor"]
    elif config is not None:
        factor = config.rope_scaling["factor"]

    # Gets the default RoPE parameters
    inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs)

    # Then applies linear scaling to the frequencies.
    # NOTE: originally, scaling was applied to the position_ids. However, we get `embs = inv_freq @ position_ids`, so
    # applying scaling to the inverse frequencies is equivalent.
    inv_freq /= factor


SUMMARY:
**Summary:**

`_compute_linear_scaling_rope_parameters` computes the inverted frequency matrix and post-processing scaling factor for the RoPE embeddings using linear scaling.

**Key details:**

*   Computes `inv_freq` (inverse frequencies) by dividing the default RoPE parameters with a scaling factor.
*   Applies linear scaling to the **pre-computed inverse frequencies (`inv_freq`)**, which is equivalent to originally applying scaling to position_ids.
*   Returns `inv_freq` and the post-processing `attention_factor`.

**Notable dependencies and interfaces:**

*   `_compute_default_rope_parameters`: a separate function that computes default RoPE parameters.

Note:

*   The code has inconsistent behavior with respect to the input keywords `config` and `rope_kwargs`.
*   A warning is raised if both keywords are present.


key: chroma:document
string_value:
CODE:

def _compute_dynamic_ntk_parameters(
    config: Optional[PretrainedConfig] = None,
    device: Optional["torch.device"] = None,
    seq_len: Optional[int] = None,
    **rope_kwargs,
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla
    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length, used to update the dynamic RoPE at inference time.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
    """
    # TODO (joao): use the new `original_max_position_embeddings` from rope_scaling
    if config is not None and len(rope_kwargs) > 0:
        raise ValueError(
            "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in "
            f"`_compute_dynamic_ntk_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}"
        )
    if len(rope_kwargs) > 0:
        base = rope_kwargs["base"]
        dim = rope_kwargs["dim"]
        max_position_embeddings = rope_kwargs["max_position_embeddings"]
        factor = rope_kwargs["factor"]
    elif config is not None:
        base = config.rope_theta
        partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        dim = int(head_dim * partial_rotary_factor)
        max_position_embeddings = config.max_position_embeddings
        factor = config.rope_scaling["factor"]

    attention_factor = 1.0  # Unused in this type of RoPE

    # seq_len: default to max_position_embeddings, e.g. at init time
    seq_len = seq_len if seq_len is not None and seq_len > max_position_embeddings else max_position_embeddings

    # Compute the inverse frequencies
    base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2))
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim))


SUMMARY:
### Function Summary:

`_compute_dynamic_ntk_parameters`

Computes the inverse frequencies with NTK scaling for RoPE embeddings.

**Purpose:**

The primary purpose of this function is to calculate the inverse frequencies for RoPE embeddings using Non-Uniform Tuning (NTK) scaling. The computed values are used for post-processing and can enhance model performance.

### Key Implementation Details:

* The function takes in additional keyword arguments (`rope_kwargs`) which must be mutually exclusive with `config`. These arguments override config settings.
* It uses the `torch.device` to initialize inverse frequencies tensor.
* The base factor is computed using a polynomial expression involving sequence length, max position embeddings, and NTK scaling coefficients.

### Notable Dependencies/Requirements:

* `transformers.PretrainedConfig`
* `torcht tensor`
* `rope_scaling`

### Signature Highlights:

```python
Tensor inv_freq = 1.0 / (base ** torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)
```
The function returns a tuple containing the `inv_freq` tensor and an unused variable `attention_factor`.


key: chroma:document
string_value:
CODE:

def _compute_yarn_parameters(
    config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies with NTK scaling. Please refer to the
    [original paper](https://arxiv.org/abs/2309.00071)
    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length. Unused for this type of RoPE.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin.
    """
    # No need to keep BC with yarn, unreleased when this new pattern was created.
    if len(rope_kwargs) > 0:
        raise ValueError(
            f"Unexpected arguments: `**rope_kwargs` should be unset in `_compute_yarn_parameters`, got {rope_kwargs}"
        )

    base = config.rope_theta
    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    dim = int(head_dim * partial_rotary_factor)
    max_position_embeddings = config.max_position_embeddings
    factor = config.rope_scaling["factor"]

    # Sets the attention factor as suggested in the paper
    attention_factor = config.rope_scaling.get("attention_factor")
    if attention_factor is None:
        attention_factor = 0.1 * math.log(factor) + 1.0

    # Optional config options
    # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly)
    beta_fast = config.rope_scaling.get("beta_fast") or 32
    beta_slow = config.rope_scaling.get("beta_slow") or 1

    # Compute the inverse frequencies
    def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
        """Inverse dimension formula to find the dimension based on the number of rotations"""
        return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))

    def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
        """Find dimension range bounds based on rotations"""
        low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
        high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
        return max(low, 0), min(high, dim - 1)

    def linear_ramp_factor(min, max, dim):
        if min == max:
            max += 0.001  # Prevent singularity

        linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
        ramp_func = torch.clamp(linear_func, 0, 1)
        return ramp_func

    # Note on variable naming: "interpolation" comes from the original technique, where we interpolate the position IDs
    # to expand the possible context length. In other words, interpolation = apply scaling factor.
    pos_freqs = base ** (torch.arange(0, dim, 2).float().to(device) / dim)
    inv_freq_extrapolation = 1.0 / pos_freqs
    inv_freq_interpolation = 1.0 / (factor * pos_freqs)

    low, high = find_correction_range(beta_fast, beta_slow, dim, base, max_position_embeddings)

    # Get n-dimensional rotational scaling corrected for extrapolation
    inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).float().to(device)
    inv_freq = (
        inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
        + inv_freq_extrapolation * inv_freq_extrapolation_factor
    )



SUMMARY:
**Summary**

This function, `_compute_yarn_parameters`, computes the inverse frequencies for RoPE embeddings using NTK scaling. It takes in a model configuration (`config`), device (`device`), sequence length (optional), and additional keyword arguments (`rope_kwargs`). The function returns a tuple of two values: the inverse frequencies (`torch.Tensor`) and the post-processing scaling factor (`float`).

**Key Implementation Details**

* Computes the inverse frequencies using NTK scaling with a base value determined by `config.rope_theta`.
* Uses the `partial_rotary_factor` from `config` to adjust the dimensionality of the inverse frequencies.
* Employs interpolation to scale the position IDs, following the original technique for context expansion.
* Calculates the `attention_factor` based on the paper's suggestion, adjusting it if necessary.
* Defines two auxiliary functions: `find_correction_dim` and `find_correction_range`, used to determine the dimension range bounds.

**Notable Dependencies/Requirements**

* `torch` library for tensor operations
* `math` library for mathematical functions (e.g., log)
* `PretrainedConfig` from `transformers` module, indicating compatibility with a specific transformer model configuration.
* External libraries not listed in the function signature.


key: chroma:document
string_value:
CODE:

def _compute_longrope_parameters(
    config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies with LongRoPE scaling. Please refer to the
    [original implementation](https://github.com/microsoft/LongRoPE)
    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin.
    """
    # TODO (joao): use the new `original_max_position_embeddings` from rope_scaling
    # No need to keep BC with longrope, unreleased when this new pattern was created.
    if len(rope_kwargs) > 0:
        raise ValueError(
            "Unexpected arguments: `**rope_kwargs` should be unset in `_compute_longrope_parameters`, got "
            f"{rope_kwargs}"
        )

    base = config.rope_theta
    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    dim = int(head_dim * partial_rotary_factor)
    long_factor = config.rope_scaling["long_factor"]
    short_factor = config.rope_scaling["short_factor"]
    factor = config.rope_scaling.get("factor")
    attention_factor = config.rope_scaling.get("attention_factor")

    # NOTE: Phi3 (and potentially other models) modify `max_position_embeddings` and have a
    # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two
    # values to compute the default attention scaling factor, instead of using `factor`.
    if hasattr(config, "original_max_position_embeddings"):
        original_max_position_embeddings = config.original_max_position_embeddings
        factor = config.max_position_embeddings / config.original_max_position_embeddings
    else:
        original_max_position_embeddings = config.max_position_embeddings

    # Sets the attention factor as suggested in the paper
    if attention_factor is None:
        if factor <= 1.0:
            attention_factor = 1.0
        else:
            attention_factor = math.sqrt(1 + math.log(factor) / math.log(original_max_position_embeddings))

    # Compute the inverse frequencies -- scaled based on the target sequence length
    if seq_len and seq_len > original_max_position_embeddings:
        ext_factors = torch.tensor(long_factor, dtype=torch.float32, device=device)
    else:
        ext_factors = torch.tensor(short_factor, dtype=torch.float32, device=device)
    inv_freq_shape = torch.arange(0, dim, 2, dtype=torch.int64, device=device).float() / dim
    inv_freq = 1.0 / (ext_factors * base**inv_freq_shape)



SUMMARY:
**Summary**

Computes the inverse frequencies with LongRoPE scaling for transformer models.

**Main Purpose and Functionality**

This function calculates the inverse frequencies and post-processing scaling factor for LongRoPE embeddings. It takes into account various configuration parameters, such as `rope_theta` and `partial_rotary_factor`.

**Key Implementation Details**

* The function uses two different scale factors (`long_factor` and `short_factor`) based on the sequence length.
* It applies a scaling factor to the attention mechanism using the `factor` parameter.
* The inverse frequencies are computed using two scales: one for longer sequences (`ext_factors`) and another for shorter sequences.

**Notable Dependencies or Requirements**

* `torch`: used for tensor computations
* `transformers.PretrainedConfig`: used to access model configuration parameters


key: chroma:document
string_value:
CODE:

def _compute_llama3_parameters(
    config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs
) -> Tuple["torch.Tensor", float]:
    """
    Computes the inverse frequencies for llama 3.1.

    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length. Unused for this type of RoPE.
        rope_kwargs (`Dict`, *optional*):
            BC compatibility with the previous RoPE class instantiation, will be removed in v4.45.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin.
    """
    # Gets the default RoPE parameters
    inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs)

    factor = config.rope_scaling["factor"]  # `8` in the original implementation
    low_freq_factor = config.rope_scaling["low_freq_factor"]  # `1` in the original implementation
    high_freq_factor = config.rope_scaling["high_freq_factor"]  # `4` in the original implementation
    old_context_len = config.rope_scaling["original_max_position_embeddings"]  # `8192` in the original implementation

    low_freq_wavelen = old_context_len / low_freq_factor
    high_freq_wavelen = old_context_len / high_freq_factor

    wavelen = 2 * math.pi / inv_freq
    # wavelen < high_freq_wavelen: do nothing
    # wavelen > low_freq_wavelen: divide by factor
    inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq)
    # otherwise: interpolate between the two, using a smooth factor
    smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
    smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
    is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)
    inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)



SUMMARY:
**Summary**

The `_compute_llama3_parameters` function calculates the inverse frequencies for RoPE embeddings in a transformer model. It takes a `PretrainedConfig`, `device`, and optional sequence length as inputs.

**Main Purpose and Functionality**

Computes inverse frequencies (`inv_freq`) and attention factor (`attention_factor`) for RoPE embeddings, applying scaling factors to adapt to different sequence lengths.

**Key Implementation Details**

1. Scales input `inv_freq` by a factor when sequence length is low.
2. Interpolates between scaled and unscaled values using a smooth factor.
3. Applies conditional scaling based on medium frequency range.

**Notable Dependencies and Requirements**

* `PretrainedConfig`
* `torch.device`
* `rope_kwargs` (optional)
* External dependencies:
	+ `math` for mathematical operations
	+ `torch.Tensor` for tensor calculations


key: chroma:document
string_value:
CODE:

def _check_received_keys(
    rope_type: str,
    received_keys: set,
    required_keys: set,
    optional_keys: Optional[set] = None,
    ignore_keys: Optional[set] = None,
):
    """Compare the received keys in `config.rope_scaling` against the expected and optional keys"""
    # BC: "rope_type" was originally "type" -- let's check for "rope_type" when "type" is present
    if "type" in received_keys:
        received_keys -= {"type"}
        required_keys.add("rope_type")

    # Some models need to store model-specific keys, and we don't want to throw warning at them
    if ignore_keys is not None:
        received_keys -= ignore_keys

    missing_keys = required_keys - received_keys
    if missing_keys:
        raise KeyError(f"Missing required keys in `rope_scaling` for 'rope_type'='{rope_type}': {missing_keys}")

    if optional_keys is not None:
        unused_keys = received_keys - required_keys - optional_keys
    else:
        unused_keys = received_keys - required_keys
    if unused_keys:


SUMMARY:
**Summary:**

The `_check_received_keys` function verifies the received keys against expected and optional keys in the `config.rope_scaling` dictionary. It checks for:

- Mandatory keys (`required_keys`) and alerts for missing keys
- Optional keys (`optional_keys`) and logs warnings for unused keys
- Ignored keys (`ignore_keys`) to prevent warnings

**Key Implementation Details:**

- Checks for deprecated key names ("type") and updates `rope_type` accordingly
- Allows ignoring specific models with a custom `ignore_keys` set
- Raises KeyError when required keys are missing
- Logs warning messages for unused or unrecognized keys

**External Dependencies:**

* `Optional[set]`: type hint for optional sets
* `KeyError`: exception raised for missing required keys
* `logger`: logging module used to log warnings
* `config`: configuration dictionary containing rope_scaling data


key: chroma:document
string_value:
CODE:

def _validate_default_rope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    rope_scaling = config.rope_scaling
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))  # BC: "rope_type" was originally "type"
    required_keys = {"rope_type"}
    received_keys = set(rope_scaling.keys())


SUMMARY:
**Summary**

`_validate_default_rope_parameters` function validates the default parameters for Ropedet's rope scaling. It:

* Retrieves `rope_scaling` and its underlying type (`rope_type`) from the configuration.
* Verifies that the required keys ("rope_type") are present in the received keys, alongside any specified ignore keys.

**Key Details**

* Main purpose: Validate default parameters for Ropedet's rope scaling.
* Key implementation detail: `_check_received_keys` function is called to check the presence of required keys in the received data.
* Notable dependencies:
	+ `PretrainedConfig`
	+ `_check_received_keys` function


key: chroma:document
string_value:
CODE:

def _validate_linear_scaling_rope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    rope_scaling = config.rope_scaling
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))  # BC: "rope_type" was originally "type"
    required_keys = {"rope_type", "factor"}
    received_keys = set(rope_scaling.keys())
    _check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys)

    factor = rope_scaling["factor"]
    if factor is None or not isinstance(factor, float) or factor < 1.0:


SUMMARY:
**Summary:**

Function `_validate_linear_scaling_rope_parameters`: Validates input parameters for linear scaling rope in a model configuration.

**Purpose and Functionality:**

- Ensures `rope_type` and `factor` are present in the input dictionary
- Checks if `factor` is a valid float value greater than or equal to 1.0

**Key Implementation Details:**

- `_check_received_keys`: A separate validation function that checks for specific keys, ensuring they match the required keys (`"rope_type"` and `"factor"`)
- Validation of `factor` as a float value within a specific range (inclusive of 1.0)

**Notable Dependencies or Requirements:**

- External dependencies: `logger`
  
**Additional Information:**
`rope_scaling` dictionary contains at least `"type"` or `"rope_type"` key for type validation, but also receives `"factor"` for further factor validation and logging messages.


key: chroma:document
string_value:
CODE:

def _validate_dynamic_scaling_rope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    rope_scaling = config.rope_scaling
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))  # BC: "rope_type" was originally "type"
    required_keys = {"rope_type", "factor"}
    # TODO (joao): update logic for the inclusion of `original_max_position_embeddings`
    optional_keys = {"original_max_position_embeddings"}
    received_keys = set(rope_scaling.keys())
    _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)

    factor = rope_scaling["factor"]
    if factor is None or not isinstance(factor, float) or factor < 1.0:


SUMMARY:
**Summary**

 `_validate_dynamic_scaling_rope_parameters` function Validates the parameters of dynamic scaling using ropes.

**Key Implementation Details**
- Checks if `rope_type` and `factor` are present in `rope_scaling`.
- Ensures `factor` is a float greater than or equal to 1.0.
- Logs warning if invalid input is detected.

**Notable Dependencies/Requirements**
- `PretrainedConfig`
- No external dependencies listed. 

Note: This summary preserves the exact names of classes, methods, and variables while focusing on the functionality and key implementation details of the provided function code chunk.


key: chroma:document
string_value:
CODE:

def _validate_yarn_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    rope_scaling = config.rope_scaling
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))  # BC: "rope_type" was originally "type"
    required_keys = {"rope_type", "factor"}
    optional_keys = {"attention_factor", "beta_fast", "beta_slow"}
    received_keys = set(rope_scaling.keys())
    _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)

    factor = rope_scaling["factor"]
    if factor is None or not isinstance(factor, float) or factor < 1.0:
        logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")

    attention_factor = rope_scaling.get("attention_factor")
    if attention_factor is not None and (not isinstance(attention_factor, float) or attention_factor < 0):
        logger.warning(
            f"`rope_scaling`'s attention_factor field must be a float greater than 0, got {attention_factor}"
        )
    beta_fast = rope_scaling.get("beta_fast")
    if beta_fast is not None and not isinstance(beta_fast, float):
        logger.warning(f"`rope_scaling`'s beta_fast field must be a float, got {beta_fast}")
    beta_slow = rope_scaling.get("beta_slow")
    if beta_slow is not None and not isinstance(beta_slow, float):
        logger.warning(f"`rope_scaling`'s beta_slow field must be a float, got {beta_slow}")

    if (beta_fast or 32) < (beta_slow or 1):
        logger.warning(
            f"`rope_scaling`'s beta_fast field must be greater than beta_slow, got beta_fast={beta_fast} "
            f"(defaults to 32 if None) and beta_slow={beta_slow} (defaults to 1 if None)"


SUMMARY:
**Summary**

### Function Purpose and Functionality
`_validate_yarn_parameters` validates the contents of `rope_scaling` dictionary in T5 model configuration (`PretrainedConfig`). It checks for required keys, optional keys, and specific numeric constraints.

### Key Implementation Details

* Validates `rope_type`, `factor`, `attention_factor`, `beta_fast`, and `beta_slow` fields
* Checks for numeric type and constraint satisfaction:
	+ `factor` must be a float ≥ 1.0
	+ `attention_factor` must be a float > 0
	+ `beta_fast` and `beta_slow` must be floats
* Verifies the relationship between `beta_fast` and `beta_slow`: `beta_fast` must be greater than `beta_slow`, or defaults to 32 if None

### Notable Dependencies and Requirements

* `PretrainedConfig`
* `Optional[set]`

### External Dependencies

* `logger` (assuming a logging system is in place)


key: chroma:document
string_value:
CODE:

def _validate_longrope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    rope_scaling = config.rope_scaling
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))  # BC: "rope_type" was originally "type"
    required_keys = {"rope_type", "short_factor", "long_factor"}
    # TODO (joao): update logic for the inclusion of `original_max_position_embeddings`
    optional_keys = {"attention_factor", "factor", "original_max_position_embeddings"}
    received_keys = set(rope_scaling.keys())
    _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys)

    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    dim = int(head_dim * partial_rotary_factor)

    short_factor = rope_scaling.get("short_factor")
    if not isinstance(short_factor, list) and all(isinstance(x, (int, float)) for x in short_factor):
        logger.warning(f"`rope_scaling`'s short_factor field must be a list of numbers, got {short_factor}")
    if not len(short_factor) == dim // 2:
        logger.warning(f"`rope_scaling`'s short_factor field must have length {dim // 2}, got {len(short_factor)}")

    long_factor = rope_scaling.get("long_factor")
    if not isinstance(long_factor, list) and all(isinstance(x, (int, float)) for x in long_factor):
        logger.warning(f"`rope_scaling`'s long_factor field must be a list of numbers, got {long_factor}")
    if not len(long_factor) == dim // 2:
        logger.warning(f"`rope_scaling`'s long_factor field must have length {dim // 2}, got {len(long_factor)}")

    # Handle Phi3 divergence: prefer the use of `attention_factor` and/or `factor` over
    # `original_max_position_embeddings` to compute internal variables. The latter lives outside `rope_scaling` and is
    # unique to longrope (= undesirable)
    if hasattr(config, "original_max_position_embeddings"):
        logger.warning_once(
            "This model has set a `original_max_position_embeddings` field, to be used together with "
            "`max_position_embeddings` to determine a scaling factor. Please set the `factor` field of `rope_scaling`"
            "with this ratio instead -- we recommend the use of this field over `original_max_position_embeddings`, "
            "as it is compatible with most model architectures."
        )
    else:
        factor = rope_scaling.get("factor")
        if factor is None:
            logger.warning("Missing required keys in `rope_scaling`: 'factor'")
        elif not isinstance(factor, float) or factor < 1.0:
            logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")

        attention_factor = rope_scaling.get("attention_factor")
        if attention_factor is not None:
            if not isinstance(attention_factor, float) or attention_factor < 0.0:
                logger.warning(
                    f"`rope_scaling`'s attention_factor field must be a float greater than 0, got {attention_factor}"


SUMMARY:
**Function Summary: `_validate_longrope_parameters`**

**Main Purpose and Functionality**
This function validates the parameters for calculating scaling factors in LongRoPe (LoR) model architecture. It checks if required keys are present, updates default values, and ensures consistency between key types.

**Key Implementation Details**
1. **Rotation Factor Calculation**: The `partial_rotary_factor` is calculated based on the Model's dynamic head dimensions.
2. **Short/Long Factor Validation**: Both lists must be of equal length (`dim // 2`) with integer/float values, otherwise a warning is logged.

**Notable Dependencies or Requirements**
- `PretrainedConfig`
- `rope_scaling` dictionary
- Various model-specific constants and attributes (e.g., `partial_rotary_factor`, `head_dim`)
- External logging module (`logger`)

**External Dependencies:**

*   `config`: a `PretrainedConfig` object
*   `rope_scaling`: a dictionary containing LoR scaling factors

Note that the function also relies on having a logger configured to print out warnings and other messages.


key: chroma:document
string_value:
CODE:

def _validate_llama3_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    rope_scaling = config.rope_scaling
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))  # BC: "rope_type" was originally "type"
    required_keys = {"rope_type", "factor", "original_max_position_embeddings", "low_freq_factor", "high_freq_factor"}
    received_keys = set(rope_scaling.keys())
    _check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys)

    factor = rope_scaling["factor"]
    if factor is None or not isinstance(factor, float) or factor < 1.0:
        logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}")

    low_freq_factor = rope_scaling["low_freq_factor"]
    high_freq_factor = rope_scaling["high_freq_factor"]
    if low_freq_factor is None or not isinstance(low_freq_factor, float):
        logger.warning(f"`rope_scaling`'s low_freq_factor field must be a float, got {low_freq_factor}")
    if high_freq_factor is None or not isinstance(high_freq_factor, float):
        logger.warning(f"`rope_scaling`'s high_freq_factor field must be a float, got {high_freq_factor}")
    if high_freq_factor <= low_freq_factor:
        logger.warning(
            "`rope_scaling`'s high_freq_factor field must be greater than low_freq_factor, got high_freq_factor="
            f"{high_freq_factor} and low_freq_factor={low_freq_factor}"
        )

    original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
    if original_max_position_embeddings is None or not isinstance(original_max_position_embeddings, int):
        logger.warning(
            "`rope_scaling`'s original_max_position_embeddings field must be an integer, got "
            f"{original_max_position_embeddings}"
        )
    if original_max_position_embeddings >= config.max_position_embeddings:
        logger.warning(
            "`rope_scaling`'s original_max_position_embeddings field must be less than max_position_embeddings, got "
            f"{original_max_position_embeddings} and max_position_embeddings={config.max_position_embeddings}"


SUMMARY:
**Summary**

The `_validate_llama3_parameters` function validates the inputs to a RoPe scaling configuration. It checks for the presence of specific keys, ensures required field types, and enforces constraints on field values.

**Key Implementation Details**

*   **Validation**: The function verifies that `rope_type`, `factor`, `low_freq_factor`, `high_freq_factor`, and `original_max_position_embeddings` exist in `rope_scaling` and have the correct types.
*   **Constraint Enforcement**:
    *   `factor` must be a float >= 1.0
    *   `low_freq_factor` and `high_freq_factor` must be floats, with `high_freq_factor > low_freq_factor`
    *   `original_max_position_embeddings` must be an integer < `max_position_embeddings`

**Notable Dependencies**

*   `config.PretrainedConfig`: The configuration object containing the RoPe scaling parameters.
*   `logger`: A logging utility for printing warnings.

**External Dependencies**

None.

**Interfaces and Classes Used**

*   `_check_received_keys`: An unknown function not defined in this snippet, likely a helper function to validate the received keys against the required keys.


key: chroma:document
string_value:
CODE:

def rope_config_validation(config: PretrainedConfig, ignore_keys: Optional[set] = None):
    """
    Validate the RoPE config arguments, given a `PretrainedConfig` object
    """
    rope_scaling = getattr(config, "rope_scaling", None)  # not a default parameter in `PretrainedConfig`
    if rope_scaling is None:
        return

    # BC: "rope_type" was originally "type"
    rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", "default"))
    validation_fn = ROPE_VALIDATION_FUNCTIONS.get(rope_type)
    if validation_fn is not None:
        validation_fn(config, ignore_keys=ignore_keys)
    else:
        logger.warning(
            f"Missing validation function mapping in `ROPE_VALIDATION_FUNCTIONS` for 'rope_type'='{rope_type}'"


SUMMARY:
**Summary:**

The `rope_config_validation` function validates the RoPE configuration based on the provided `PretrainedConfig` object. It checks if the requested scaling method is specified and, if so, looks up the corresponding validation function in the `ROPE_VALIDATION_FUNCTIONS` dictionary. If a function is found, it executes the validation with the given configuration and optional ignored keys.

**Key Details:**

* Validates `rope_scaling` value from `PretrainedConfig`.
* Looks up validation function based on `rope_type`/`type` in `ROPE_VALIDATION_FUNCTIONS` dictionary.
* Executes validation function if found, passing `config` and `ignore_keys`.

**Notable Dependencies:**

* `PretrainedConfig`
* `ROPE_VALIDATION_FUNCTIONS`

**Notes:**

* `rope_scaling` value must be specified or returns an empty result.
* Logger warning is raised if no validation function mapping exists for the given `rope_type`.


key: chroma:document
string_value:
CODE:

class Trie:
    """
    Trie in Python. Creates a Trie out of a list of words. The trie is used to split on `added_tokens` in one pass
    Loose reference https://en.wikipedia.org/wiki/Trie
    """

    def __init__(self, *args):
        self.data = {}
        self._tokens = set()
        self._termination_char = ""
        self.update(*args)

    def update(self, *args):
        """
        Updates the Trie with new tokens provided as arguments.

        Args:
            *args: Variable number of words to be added to the Trie.
        """
        for token in tuple(*args):
            self.add(token)

    def add(self, word: str):
        """
        Passes over every char (utf-8 char) on word and recursively adds it to the internal `data` trie representation.
        The special key `""` in `self._termination_char` is used to represent termination.

        This function is idempotent, adding twice the same word will leave the trie unchanged

        Example:

        ```python
        >>> trie = Trie()
        >>> trie.add("Hello 友達")
        >>> trie.data
        {"H": {"e": {"l": {"l": {"o": {" ": {"友": {"達": {"": 1}}}}}}}}}

        >>> trie.add("Hello")
        >>> trie.data
        {"H": {"e": {"l": {"l": {"o": {"": 1, " ": {"友": {"達": {"": 1}}}}}}}}}
        ```
        """
        if not word:
            # Prevent empty string
            return

        self._tokens.add(word)
        ref = self.data
        for char in word:
            ref[char] = ref.setdefault(char, {})
            ref = ref[char]
        ref[self._termination_char] = 1

    def split(self, text: str) -> List[str]:
        """
        Will look for the words added to the trie within `text`. Output is the original string splitted along the
        boundaries of the words found.

        This trie will match the longest possible word first !

        Example:

        ```python
        >>> trie = Trie()
        >>> trie.split("[CLS] This is a extra_id_100")
        ["[CLS] This is a extra_id_100"]

        >>> trie.add("[CLS]")
        >>> trie.add("extra_id_1")
        >>> trie.add("extra_id_100")
        >>> trie.split("[CLS] This is a extra_id_100")
        ["[CLS]", " This is a ", "extra_id_100"]
        ```
        """
        # indexes are counted left of the chars index.
        # "hello", index 0, is left of h, index 1 is between h and e.
        # index 5 is right of the "o".

        # States are going to capture every possible start (indexes as above)
        # as keys, and have as values, a pointer to the position in the trie
        # where we're at. This is a partial match for now.
        # This enables to keep track of multiple matches while we're iterating
        # the string
        # If the trie contains, "blowing", and "lower" and we encounter the
        # string "blower", we need to split into ["b", "lower"].
        # This is where we need to keep track of multiple possible starts.
        states = OrderedDict()

        # This will contain every indices where we need
        # to cut.
        # We force to cut at offset 0 and len(text) (added later)
        offsets = [0]

        # This is used by the lookahead which needs to skip over
        # some text where the full match exceeded the place in the initial
        # for loop
        skip = 0
        # Main loop, Giving this algorithm O(n) complexity
        for current, current_char in enumerate(text):
            if skip and current < skip:
                # Prevents the lookahead for matching twice
                # like extra_id_100 and id_100
                continue

            # This will track every state
            # that stop matching, we need to stop tracking them.
            # If we look at "lowball", we're going to match "l" (add it to states), "o", "w", then
            # fail on "b", we need to remove 0 from the valid states.
            to_remove = set()
            # Whenever we found a match, we need to drop everything
            # this is a greedy algorithm, it will match on the first found token
            reset = False

            # In this case, we already have partial matches (But unfinished)
            for start, trie_pointer in states.items():
                if "" in trie_pointer:
                    # This is a final match, we need to reset and
                    # store the results in `offsets`.

                    # Lookahead to match longest first
                    # Important in case of extra_id_1 vs extra_id_100
                    # Here we are also actively looking for other earlier partial
                    # matches
                    # "[CLS]", "L", we need to match CLS even if L is special
                    for lookstart, looktrie_pointer in states.items():
                        if lookstart > start:
                            # This partial match is later, we can stop looking
                            break
                        elif lookstart < start:
                            # This partial match is earlier, the trie pointer
                            # was already updated, so index is + 1
                            lookahead_index = current + 1
                            end = current + 1
                        else:
                            # Here lookstart == start and
                            #      looktrie_pointer == trie_pointer
                            # It wasn't updated yet so indices are current ones
                            lookahead_index = current
                            end = current
                        next_char = text[lookahead_index] if lookahead_index < len(text) else None
                        if "" in looktrie_pointer:
                            start = lookstart
                            end = lookahead_index
                            skip = lookahead_index

                        while next_char in looktrie_pointer:
                            looktrie_pointer = looktrie_pointer[next_char]
                            lookahead_index += 1
                            if "" in looktrie_pointer:
                                start = lookstart
                                end = lookahead_index
                                skip = lookahead_index

                            if lookahead_index == len(text):
                                # End of string
                                break
                            next_char = text[lookahead_index]
                        # End lookahead

                    # Storing and resetting
                    offsets.append(start)
                    offsets.append(end)
                    reset = True
                    break
                elif current_char in trie_pointer:
                    # The current character being looked at has a match within the trie
                    # update the pointer (it will be stored back into states later).
                    trie_pointer = trie_pointer[current_char]

                    # Storing back the new pointer into the states.
                    # Partial matches got longer by one.
                    states[start] = trie_pointer
                else:
                    # The new character has not match in the trie, we need
                    # to stop keeping track of this partial match.
                    # We can't do it directly within the loop because of how
                    # python iteration works
                    to_remove.add(start)

            # Either clearing the full start (we found a real match)
            # Or clearing only the partial matches that didn't work.
            if reset:
                states = {}
            else:
                for start in to_remove:
                    del states[start]

            # If this character is a starting character within the trie
            # start keeping track of this partial match.
            if current >= skip and current_char in self.data:
                states[current] = self.data[current_char]

        # We have a cut at the end with states.
        for start, trie_pointer in states.items():
            if "" in trie_pointer:
                # This is a final match, we need to reset and
                # store the results in `offsets`.
                end = len(text)
                offsets.append(start)
                offsets.append(end)
                # Longest cut is always the one with lower start so the first
                # item so we need to break.
                break

        return self.cut_text(text, offsets)

    def cut_text(self, text, offsets):
        # We have all the offsets now, we just need to do the actual splitting.
        # We need to eventually add the first part of the string and the eventual
        # last part.
        offsets.append(len(text))
        tokens = []
        start = 0
        for end in offsets:
            if start > end:
                logger.error(
                    "There was a bug in Trie algorithm in tokenization. Attempting to recover. Please report it"
                    " anyway."
                )
                continue
            elif start == end:
                # This might happen if there's a match at index 0
                # we're also preventing zero-width cuts in case of two
                # consecutive matches
                continue
            tokens.append(text[start:end])
            start = end



SUMMARY:
The provided class code defines a Trie data structure implementation in Python, which is used for tokenization. The Trie allows splitting on `added_tokens` within one pass.

Key Features:

1.  **Initialization**: The Trie can be initialized with an arbitrary number of words or tokens.
2.  **Adding Tokens**: The `add` method adds the provided word to the Trie, recursively creating necessary child nodes if needed.
3.  **Splitting Text**: The `split` method splits a given text into substrings based on prefixes found in the added tokens.

The Trie implementation is designed for efficient tokenization, leveraging Python's iterator behavior and recursion for optimal node traversal.


key: chroma:document
string_value:
CODE:

class ExtensionsTrie(Trie):
    def __init__(self, *args):
        super().__init__(*args)

    def extensions(self, prefix: str):
        """
        Generates all extensions of a given prefix token in the Trie.

        Example:

        ```python
        >>> trie = Trie()
        >>> trie.add("apple")
        >>> trie.add("app")
        >>> trie.add("application")
        >>> trie.extensions("app")
        ['app', 'apple', 'application']
        ```
        """
        prefix_node = self._get_node(prefix)
        ret = self._collect_tokens(prefix_node)
        return [prefix + token for token in ret]

    def _get_node(self, token: str) -> dict:
        """
        Retrieves the node corresponding to the given token in the Trie.

        Args:
            token (str): The token for which the corresponding node needs to be retrieved.

        Returns:
            dict: The node in the Trie corresponding to the given token.
        """
        node = self.data
        for char in token:
            if char not in node:
                break

            node = node[char]
        return node

    def _collect_tokens(self, node: dict) -> list:
        """
        Generates all tokens in the Trie starting from a given node.

        Args:
            node (dict): The node in the Trie from which tokens need to be generated.

        Returns:
            list: List of tokens generated from the given node.
        """
        tokens = [self._termination_char] if self._termination_char in node else []
        for token, subtrie_head in node.items():
            if token != self._termination_char:
                subtokens = self._collect_tokens(subtrie_head)
                tokens.extend([token + subtoken for subtoken in subtokens])


SUMMARY:
**Summary:**

The `ExtensionsTrie` class extends the ` Trie` class and provides methods to generate extensions of given prefixes.

**Key Functions:**

* `extensions(prefix: str)`: Returns all extensions of a given prefix token in the Trie.
	+ Calls `_get_node(prefix)` to retrieve the corresponding node.
	+ Uses `_collect_tokens(node)` to collect tokens from this node.
* `_get_node(token: str) -> dict`: Retrieves the node corresponding to the given token.
	+ Traverses the Trie by character to find the matching node.
* `_collect_tokens(node: dict) -> list`: Generates all tokens in the Trie starting from a given node.
	+ Recursive function that collects tokens and their subtokens.

**Notable Dependencies/Requirements:**

* `data`: The underlying Trie data structure, assumed implemented elsewhere.
* `_termination_char` (not defined in this snippet): A special character indicating the end of a token.

Note that the provided summary preserves the original class name, method signatures, and variable names for contextuality.


key: chroma:document
string_value:
CODE:

def _is_whitespace(char):
    """Checks whether `char` is a whitespace character."""
    # \t, \n, and \r are technically control characters but we treat them
    # as whitespace since they are generally considered as such.
    if char == " " or char == "\t" or char == "\n" or char == "\r":
        return True
    cat = unicodedata.category(char)
    if cat == "Zs":
        return True


SUMMARY:
**Summary**

Function `_is_whitespace` checks whether a given character is considered whitespace. It returns `True` for:

* Conventional whitespace characters: space, tab, newline, and carriage return
* Unicode categories designated as "Zs" by the Unicode Standard (which includes similar whitespace characters)

**Key Implementation Details**

* Method signature: `_is_whitespace(char)`
* Used `unicodedata.category(char)` to determine unicode category

**Dependencies/Requirements**

* None detected.

This summary aims to provide a concise overview of the function's purpose, implementation details, and any dependencies/requirements.


key: chroma:document
string_value:
CODE:

def _is_control(char):
    """Checks whether `char` is a control character."""
    # These are technically control characters but we count them as whitespace
    # characters.
    if char == "\t" or char == "\n" or char == "\r":
        return False
    cat = unicodedata.category(char)
    if cat.startswith("C"):
        return True


SUMMARY:
**Function Summary: `_is_control`**

*   **Primary Purpose:** Checks if a given character is a control character.
*   **Implementation Details:**
    *   Exempts tab (`\t`), newline (`\n`), and carriage return (`\r`) characters from the control character check, treating them as whitespace characters instead.
    *   Uses Unicode category to determine whether the character is a control character (category starts with "C").
*   **Key Interfaces/Functions:**
    *   `unicodedata.category(char)`
*   **External Dependencies:** Unicodedata module

**RAG Code Index Entry**

```
_isControl(definition=_is_control)
(isControlCharacter, _is_control)
```

This summary provides a concise overview of the `_is_control` function's purpose and implementation details. The use of technical terms and precise names helps developers understand complex codebases efficiently.


key: chroma:document
string_value:
CODE:

def _is_punctuation(char):
    """Checks whether `char` is a punctuation character."""
    cp = ord(char)
    # We treat all non-letter/number ASCII as punctuation.
    # Characters such as "^", "$", and "`" are not in the Unicode
    # Punctuation class but we treat them as punctuation anyways, for
    # consistency.
    if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126):
        return True
    cat = unicodedata.category(char)
    if cat.startswith("P"):
        return True


SUMMARY:
**Summary:**

The `_is_punctuation` function determines whether a given character is a punctuation mark. It uses a combination ofASCII code range checks, Unicode categorization, and predefined punctuation characters to accurately identify punctuation marks.

**Key Points:**

- **Primary Purpose:** Checks if a character is a punctuation mark.
- **Implementation Details:** Uses ASCII code range checks and Unicode categorization to identify punctuation characters.

**Method Signatures and Interfaces:**

- `unicodedata.category(char)`: Uses the CAT (Character Classification) algorithm from the Python Standard Library to determine the Unicode category of the input character.
- Returns value: A boolean indicating whether the character is a punctuation mark.


