# embedding_metadata

key: chroma:document
string_value:
CODE:

class MambaCache:
    """
    Cache for mamba model which does not have attention mechanism and key value states.

    Arguments:
        config (`PretrainedConfig):
            The configuration file defining the shape-related attributes required to initialize the static cache.
        batch_size (`int`):
            The batch size with which the model will be used. Note that a new instance must be instantiated if a
            smaller batch size is used.
        dtype (`torch.dtype`, *optional*, defaults to `torch.float16`):
            The default `dtype` to use when initializing the layer.
        device (`torch.device` or `str`, *optional*):
            The device on which the cache should be initialized. Should be the same as the layer.

    Attributes:
        dtype: (`torch.dtype`):
            The default `dtype` used to initializing the cache.
        intermediate_size: (`int`):
            Model's intermediate_size taken from config.
        ssm_state_size: (`int`):
            Model's state_size taken from config.
        conv_kernel_size: (`int`):
            Model's convolution kernel size taken from config
        conv_states: (`torch.Tensor`):
            A tensor of shape `[layer_idx, batch_size, intermediate_size, conv_kernel_size]` that holds convolutional states.
        ssm_states: (`torch.Tensor`):
            A tensor of shape `[layer_idx, batch_size, intermediate_size, ssm_state_size]` that holds ssm states

    Example:

        ```python
        >>> from transformers import AutoTokenizer, MambaForCausalLM, MambaCache

        >>> model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")

        >>> inputs = tokenizer(text="My name is Mamba", return_tensors="pt")

        >>> # Prepare a cache class and pass it to model's forward
        >>> past_key_values = MambaCache(config=model.config, batch_size=1, device=model.device, dtype=model.dtype)
        >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
        >>> outputs.past_key_values
        MambaCache()
        ```
    """

    # TODO (joao): remove `=None` in non-optional arguments in v4.46. Remove from `OBJECTS_TO_IGNORE` as well.
    def __init__(
        self,
        config: PretrainedConfig,
        batch_size: int = None,
        dtype: torch.dtype = torch.float16,
        device: Optional[Union[torch.device, str]] = None,
        max_batch_size: Optional[int] = None,
    ):
        if batch_size is not None:
            logger.warning_once(
                f"The 'batch_size' argument of {self.__class__.__name__} is deprecated and will be removed in "
                "v4.49. Use the more precisely named 'max_batch_size' argument instead."
            )
        self.dtype = dtype
        self.max_batch_size = batch_size or max_batch_size
        self.intermediate_size = config.intermediate_size
        self.ssm_state_size = config.state_size
        self.conv_kernel_size = config.conv_kernel

        self.conv_states: torch.Tensor = torch.zeros(
            config.num_hidden_layers,
            self.max_batch_size,
            self.intermediate_size,
            self.conv_kernel_size,
            device=device,
            dtype=dtype,
        )
        self.ssm_states: torch.Tensor = torch.zeros(
            config.num_hidden_layers,
            self.max_batch_size,
            self.intermediate_size,
            self.ssm_state_size,
            device=device,
            dtype=dtype,
        )

        torch._dynamo.mark_static_address(self.conv_states)
        torch._dynamo.mark_static_address(self.ssm_states)

    def update_conv_state(
        self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor
    ) -> torch.Tensor:
        conv_state = self.conv_states[layer_idx]
        cache_position = cache_position.clamp(0, self.conv_kernel_size - 1)

        conv_state = conv_state.roll(shifts=-1, dims=-1)
        conv_state[:, :, cache_position] = new_conv_state.to(device=conv_state.device, dtype=conv_state.dtype)
        self.conv_states[layer_idx].zero_()
        self.conv_states[layer_idx] += conv_state
        return self.conv_states[layer_idx]

    def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor):
        self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states.device)
        return self.ssm_states[layer_idx]

    def reset(self):
        self.conv_states.zero_()
        self.ssm_states.zero_()

    @property
    def batch_size(self):
        logger.warning_once(
            f"The 'batch_size' attribute of {self.__class__.__name__} is deprecated and will be removed in "
            "v4.49. Use the more precisely named 'self.max_batch_size' attribute instead."
        )


SUMMARY:
**MambaCache Class Summary**

**Purpose and Functionality:**

The `MambaCache` class is a cache for mamba models, designed to store convolutional states and intermediate sizes. It is used in conjunction with the MambaForCausalLM model.

**Key Implementation Details:**

* Stores convolutional states (`conv_states`) of size `[layer_idx, batch_size, intermediate_size, conv_kernel_size]`.
* Stores ssm states (`ssm_states`) of size `[layer_idx, batch_size, intermediate_size, ssm_state_size]`.
* Provides methods to update conv and ssm states:
	+ `update_conv_state`: updates the convolutional state at a specific layer index.
	+ `update_ssm_state`: updates the ssm state at a specific layer index.
	+ `reset`: resets all cache values.

** dependencies:**

* `torch`: for tensor operations
* `PretrainedConfig`: an optional parameter for batch size
* `torch.dtype` and `torch.device`: used for data type and device mapping

**Interfaces/Methods:**

* `__init__`: initializes the cache with a configuration object.
* `update_conv_state`
* `update_ssm_state`
* `reset`
* `batch_size` (deprecated)


key: chroma:document
string_value:
CODE:

class OffloadedStaticCache(StaticCache):
    """
    Static cache class to be used with `torch.compile(model)` that offloads to the CPU or
    another device.

    Args:
        config (`PretrainedConfig):
            The configuration file defining the shape-related attributes required to initialize
            the static cache.
        max_batch_size (`int`):
            The maximum batch size with which the model will be used.
        max_cache_len (`int`):
            The maximum sequence length with which the model will be used.
        device (`Union[str, torch.device]`):
            The device on which the cache should be initialized. Should be the same as the
            layer device.
        dtype (`torch.dtype`, *optional*):
            The default `dtype` to use when initializing the cache.
        offload_device (`Union[str, torch.device]`, *optional*, defaults to `cpu`):
            The device to offload to. Defaults to CPU.
        layer_device_map (`Dict[int, Union[str, torch.device, int]]`, *optional*):
            Mapping between the layers and its device. This is required when you are manually initializing the cache and the model is splitted between differents gpus.
            You can know which layers mapped to which device by checking the associated device_map: `model.hf_device_map`.

    Attributes:
        key_cache (`List[torch.Tensor]`):
            Off-loaded key cache tensors. First one will be on device, where-as the others are
            off-loaded.
        value_cache (`List[torch.Tensor]`):
            Off-loaded value cache tensors. First one will be on device, where-as the others are
            off-loaded.
        max_batch_size (`int`):
            The maximum batch size with which this cache can be used.
        max_cache_len (`int`):
            The maximum sequence length with which this cache can be used.
        device (`torch.device`):
            The device on which the cache is used.
        offload_device (`torch.device`):
            The device used to offload to.
        dtype (`torch.dtype`):
            The `dtype` used to initializing the cache.

    Example:

        ```python
        >>> from transformers import AutoTokenizer, AutoModelForCausalLM, OffloadedStaticCache

        >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
        >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")

        >>> inputs = tokenizer(text="My name is GPT2", return_tensors="pt")

        >>> # Prepare a cache class and pass it to model's forward
        >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate
        >>> max_generated_length = inputs.input_ids.shape[1] + 10
        >>> past_key_values = OffloadedStaticCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype)
        >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True)
        >>> past_kv_length = outputs.past_key_values # access cache filled with key/values from generation
        ```
    """

    def __init__(
        self,
        config: PretrainedConfig,
        max_batch_size: int,
        max_cache_len: Optional[int],
        device: Union[str, torch.device],
        dtype: Optional[torch.dtype] = None,
        offload_device: Union[str, torch.device] = torch.device("cpu"),
        layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None,
    ) -> None:
        self.max_batch_size = max_batch_size
        self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len
        self.device = torch.device(device) if layer_device_map is None else layer_device_map[0]
        self.offload_device = torch.device(offload_device)
        self.dtype = dtype if dtype is not None else torch.float32

        # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads
        head_dim = config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads

        num_key_value_heads = (
            config.num_attention_heads
            if getattr(config, "num_key_value_heads", None) is None
            else config.num_key_value_heads
        )

        cache_shape = (max_batch_size, num_key_value_heads, self.max_cache_len, head_dim)

        # Create offloaded CPU tensors.
        self.key_cache: List[torch.Tensor] = []
        self.value_cache: List[torch.Tensor] = []

        for i in range(config.num_hidden_layers):
            # First layer is always on-device.
            device = self.device if i == 0 else self.offload_device

            key_cache, value_cache = self._create_key_value_cache_tensors(cache_shape, device)

            self.key_cache.append(key_cache)
            self.value_cache.append(value_cache)

        # Create device tensors.
        self._device_key_cache: List[torch.Tensor] = []
        self._device_value_cache: List[torch.Tensor] = []

        for i in range(2):
            key_cache, value_cache = self._create_key_value_cache_tensors(cache_shape, self.device)

            self._device_key_cache.append(key_cache)
            self._device_value_cache.append(value_cache)

        # For backwards compatibility.
        # TODO(gante): Remove this.
        self._seen_tokens = 0

        # Create new CUDA stream for parallel prefetching.
        self._prefetch_stream = torch.cuda.Stream() if self.device.type == "cuda" else None

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
        It is VERY important to index using a tensor, otherwise you introduce a copy to the device.

        Parameters:
            key_states (`torch.Tensor`):
                The new key states to cache.
            value_states (`torch.Tensor`):
                The new value states to cache.
            layer_idx (`int`):
                The index of the layer to cache the states for.
            cache_kwargs (`Dict[str, Any]`, *optional*):
                Additional arguments for the cache subclass. The `OffloadedStaticCache` needs the
                `cache_position` input to know how where to write in the cache.

        Return:
            A tuple containing the updated key and value states.
        """

        if layer_idx == 0:
            # Update seen tokens.
            # TODO(gante): Remove this.
            self._seen_tokens += key_states.shape[-2]

            # Always there.
            k_out = self.key_cache[0]
            v_out = self.value_cache[0]
        else:
            # Wait for prefetch stream.
            if self._prefetch_stream is not None:
                torch.cuda.default_stream(self.device).wait_stream(self._prefetch_stream)

            k_out = self._device_key_cache[layer_idx & 1]
            v_out = self._device_value_cache[layer_idx & 1]

        self._prefetch_layer(layer_idx + 1)

        cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None
        if cache_position is None:
            k_out.copy_(key_states)
            v_out.copy_(value_states)

            # Copy the values to the offloaded device as well.
            if layer_idx == 0:
                self.key_cache[layer_idx].copy_(key_states.to(self.offload_device))
                self.value_cache[layer_idx].copy_(value_states.to(self.offload_device))
        else:
            # Note: here we use `tensor.index_copy_(dim, index, tensor)` that is equivalent to
            # `tensor[:, :, index] = tensor`, but the first one is compile-friendly and it does
            # explicitly an in-place operation, that avoids copies and uses less memory.
            try:
                k_out.index_copy_(2, cache_position, key_states)
                v_out.index_copy_(2, cache_position, value_states)
            except NotImplementedError:
                # The operator 'aten::index_copy.out' is not currently implemented for the MPS
                # device.
                k_out[:, :, cache_position] = key_states
                v_out[:, :, cache_position] = value_states

            # Copy the values to the offloaded device as well.
            if layer_idx != 0:
                cache_position = cache_position.to(self.offload_device)
                key_states = key_states.to(self.offload_device)
                value_states = value_states.to(self.offload_device)

                try:
                    self.key_cache[layer_idx].index_copy_(2, cache_position, key_states)
                    self.value_cache[layer_idx].index_copy_(2, cache_position, value_states)
                except NotImplementedError:
                    # The operator 'aten::index_copy.out' is not currently implemented for the MPS
                    # device.
                    self.key_cache[layer_idx][:, :, cache_position] = key_states
                    self.value_cache[layer_idx][:, :, cache_position] = value_states

        return k_out, v_out

    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
        """Returns the sequence length of the cached states that were seen by the model."""

        # TODO(gante): Remove this.
        return self._seen_tokens

    def get_max_cache_shape(self) -> Optional[int]:
        """Returns the maximum sequence length of the cached states."""

        return self.max_cache_len

    def reset(self) -> None:
        """Resets the cache values while preserving the objects."""

        # For backwards compatibility.
        # TODO(gante): Remove this.
        self._seen_tokens = 0

        # Zero out cache.
        for layer_idx in range(len(self.key_cache)):
            # In-place ops prevent breaking the static address.
            self.key_cache[layer_idx].zero_()
            self.value_cache[layer_idx].zero_()

    @property
    def seen_tokens(self) -> int:
        # For backwards compatibility.
        # TODO(gante): Remove this.
        return self._seen_tokens

    def _create_key_value_cache_tensors(
        self, shape: Tuple[int, ...], device: torch.device
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Creates K/V cache tensors on a device. Pins memory for CPU tensors. Marks them as static
        addresses for non-CPU tensors.

        Args:
            shape (`Tuple[int, ...]`): Shape.
            device (`torch.device`): Device.

        Returns:
            Key and value cache tensors as a tuple.
        """

        is_cpu_device = device == torch.device("cpu")

        key_cache = torch.zeros(shape, dtype=self.dtype, device=device, pin_memory=is_cpu_device)
        value_cache = torch.zeros(shape, dtype=self.dtype, device=device, pin_memory=is_cpu_device)

        # Note: `mark_static_address` is used to tag the cache as a fixed data pointer,
        # preventing compiled graph breaks when updating the cache.
        torch._dynamo.mark_static_address(key_cache)
        torch._dynamo.mark_static_address(value_cache)

        return key_cache, value_cache

    def _prefetch_layer(self, layer_idx: int) -> None:
        """Prefetch a layer to the device. Needs to be called in order of layer indices."""

        # Don't fetch layers that do not exist.
        if layer_idx >= len(self.key_cache):
            return

        # Alternate between two on-device caches.
        if self._prefetch_stream is not None:
            with torch.cuda.stream(self._prefetch_stream):
                self._prefetch_layer_in_context(layer_idx)
        else:
            self._prefetch_layer_in_context(layer_idx)

    def _prefetch_layer_in_context(self, layer_idx: int) -> None:
        """Performs the actual copy of the layer to device cache."""

        self._device_key_cache[layer_idx & 1].copy_(self.key_cache[layer_idx], non_blocking=True)


SUMMARY:
This code is part of a TensorFlow model implementation, specifically dealing with caching and pre-fetching states to optimize model inference. Here's a concise summary:

**Purpose:** The module provides a way to cache and prefetch states for a deep learning model, aiming to improve the performance by reducing the computational overhead.

**Main components:**

1. **Cache**: Stores intermediate states of the model.
2. **Prefetching**: Transfers cached states to the device (e.g., GPU) in batches (layers) to optimize memory access and reduce latency.
3. **Device Cache**: Stores pre-fetched models on the device for fast access.

**Key features:**

1. **Static Address Marking**: Fixes data pointers to prevent recompilation issues during cache updates.
2. **Prefetching**: Transfers cached states to the device in batches using CUDA streams.
3. **Memory Management**: Optimizes memory usage by pinning CPU tensors and reducing the number of copies between devices.

**Benefits:**

1. Improved performance through reduced computational overhead and faster data transfer.
2. Enhanced model stability by minimizing the need for recompilation during cache updates.

Overall, this module is a key optimization component in deep learning models, helping to improve the efficiency and reliability of complex computations.


key: chroma:document
string_value:
CODE:

class RTDetrHungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network

    For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
    predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
    un-matched (and thus treated as non-objects).

    Args:
        config: RTDetrConfig
    """

    def __init__(self, config):
        super().__init__()
        requires_backends(self, ["scipy"])

        self.class_cost = config.matcher_class_cost
        self.bbox_cost = config.matcher_bbox_cost
        self.giou_cost = config.matcher_giou_cost

        self.use_focal_loss = config.use_focal_loss
        self.alpha = config.matcher_alpha
        self.gamma = config.matcher_gamma

        if self.class_cost == self.bbox_cost == self.giou_cost == 0:
            raise ValueError("All costs of the Matcher can't be 0")

    @torch.no_grad()
    def forward(self, outputs, targets):
        """Performs the matching

        Params:
            outputs: This is a dict that contains at least these entries:
                 "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
                 "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates

            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
                 "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
                           objects in the target) containing the class labels
                 "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates

        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions (in order)
                - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds:
                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        batch_size, num_queries = outputs["logits"].shape[:2]

        # We flatten to compute the cost matrices in a batch
        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]
        # Also concat the target labels and boxes
        target_ids = torch.cat([v["class_labels"] for v in targets])
        target_bbox = torch.cat([v["boxes"] for v in targets])
        # Compute the classification cost. Contrary to the loss, we don't use the NLL,
        # but approximate it in 1 - proba[target class].
        # The 1 is a constant that doesn't change the matching, it can be ommitted.
        if self.use_focal_loss:
            out_prob = F.sigmoid(outputs["logits"].flatten(0, 1))
            out_prob = out_prob[:, target_ids]
            neg_cost_class = (1 - self.alpha) * (out_prob**self.gamma) * (-(1 - out_prob + 1e-8).log())
            pos_cost_class = self.alpha * ((1 - out_prob) ** self.gamma) * (-(out_prob + 1e-8).log())
            class_cost = pos_cost_class - neg_cost_class
        else:
            out_prob = outputs["logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
            class_cost = -out_prob[:, target_ids]

        # Compute the L1 cost between boxes
        bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)
        # Compute the giou cost betwen boxes
        giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))
        # Compute the final cost matrix
        cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
        cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()

        sizes = [len(v["boxes"]) for v in targets]
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]



SUMMARY:
**RTDetrHungarianMatcher Class Summary**

**Purpose and Functionality**

This class implements the Hungarian Matcher algorithm to compute an assignment between targets and predictions of a network. It takes a `config` object as input and extracts various hyperparameters and costs.

**Key Implementation Details**

* The `forward` method takes in `outputs` (logits, predictions) and `targets` (class labels, box coordinates).
* It computes classification cost using focal loss or NLL, depending on the `use_focal_loss` flag.
* It calculates L1 cost between boxes and GIoU cost between boxes using custom functions.
* It constructs a final cost matrix combining all costs and assigns indices to predict matches using Hungarian algorithm.

**Notable Dependencies and Requirements**

* `torch`, `F`, `linear_sum_assignment`, `generalized_box_iou`
* Requires `scipy` backend for optimization

**External Interfaces**

* `RTDetrConfig`: Input configuration object
* `torch.no_grad()`: Enables computation graph disabling during forward pass


key: chroma:document
string_value:
CODE:

class RTDetrLoss(nn.Module):
    """
    This class computes the losses for RTDetr. The process happens in two steps: 1) we compute hungarian assignment
    between ground truth boxes and the outputs of the model 2) we supervise each pair of matched ground-truth /
    prediction (supervise class and box).

    Args:
        matcher (`DetrHungarianMatcher`):
            Module able to compute a matching between targets and proposals.
        weight_dict (`Dict`):
            Dictionary relating each loss with its weights. These losses are configured in RTDetrConf as
            `weight_loss_vfl`, `weight_loss_bbox`, `weight_loss_giou`
        losses (`List[str]`):
            List of all the losses to be applied. See `get_loss` for a list of all available losses.
        alpha (`float`):
            Parameter alpha used to compute the focal loss.
        gamma (`float`):
            Parameter gamma used to compute the focal loss.
        eos_coef (`float`):
            Relative classification weight applied to the no-object category.
        num_classes (`int`):
            Number of object categories, omitting the special no-object category.
    """

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

        self.matcher = RTDetrHungarianMatcher(config)
        self.num_classes = config.num_labels
        self.weight_dict = {
            "loss_vfl": config.weight_loss_vfl,
            "loss_bbox": config.weight_loss_bbox,
            "loss_giou": config.weight_loss_giou,
        }
        self.losses = ["vfl", "boxes"]
        self.eos_coef = config.eos_coefficient
        empty_weight = torch.ones(config.num_labels + 1)
        empty_weight[-1] = self.eos_coef
        self.register_buffer("empty_weight", empty_weight)
        self.alpha = config.focal_loss_alpha
        self.gamma = config.focal_loss_gamma

    def loss_labels_vfl(self, outputs, targets, indices, num_boxes, log=True):
        if "pred_boxes" not in outputs:
            raise KeyError("No predicted boxes found in outputs")
        if "logits" not in outputs:
            raise KeyError("No predicted logits found in outputs")
        idx = self._get_source_permutation_idx(indices)

        src_boxes = outputs["pred_boxes"][idx]
        target_boxes = torch.cat([_target["boxes"][i] for _target, (_, i) in zip(targets, indices)], dim=0)
        ious, _ = box_iou(center_to_corners_format(src_boxes), center_to_corners_format(target_boxes))
        ious = torch.diag(ious).detach()

        src_logits = outputs["logits"]
        target_classes_original = torch.cat([_target["class_labels"][i] for _target, (_, i) in zip(targets, indices)])
        target_classes = torch.full(
            src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
        )
        target_classes[idx] = target_classes_original
        target = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1]

        target_score_original = torch.zeros_like(target_classes, dtype=src_logits.dtype)
        target_score_original[idx] = ious.to(target_score_original.dtype)
        target_score = target_score_original.unsqueeze(-1) * target

        pred_score = F.sigmoid(src_logits).detach()
        weight = self.alpha * pred_score.pow(self.gamma) * (1 - target) + target_score

        loss = F.binary_cross_entropy_with_logits(src_logits, target_score, weight=weight, reduction="none")
        loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes
        return {"loss_vfl": loss}

    def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
        """Classification loss (NLL)
        targets dicts must contain the key "class_labels" containing a tensor of dim [nb_target_boxes]
        """
        if "logits" not in outputs:
            raise KeyError("No logits were found in the outputs")

        src_logits = outputs["logits"]

        idx = self._get_source_permutation_idx(indices)
        target_classes_original = torch.cat([_target["class_labels"][i] for _target, (_, i) in zip(targets, indices)])
        target_classes = torch.full(
            src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
        )
        target_classes[idx] = target_classes_original

        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.class_weight)
        losses = {"loss_ce": loss_ce}
        return losses

    @torch.no_grad()
    def loss_cardinality(self, outputs, targets, indices, num_boxes):
        """
        Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes. This is not
        really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
        """
        logits = outputs["logits"]
        device = logits.device
        target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
        # Count the number of predictions that are NOT "no-object" (which is the last class)
        card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
        card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
        losses = {"cardinality_error": card_err}
        return losses

    def loss_boxes(self, outputs, targets, indices, num_boxes):
        """
        Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss. Targets dicts must
        contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes are expected in
        format (center_x, center_y, w, h), normalized by the image size.
        """
        if "pred_boxes" not in outputs:
            raise KeyError("No predicted boxes found in outputs")
        idx = self._get_source_permutation_idx(indices)
        src_boxes = outputs["pred_boxes"][idx]
        target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)

        losses = {}

        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction="none")
        losses["loss_bbox"] = loss_bbox.sum() / num_boxes

        loss_giou = 1 - torch.diag(
            generalized_box_iou(center_to_corners_format(src_boxes), center_to_corners_format(target_boxes))
        )
        losses["loss_giou"] = loss_giou.sum() / num_boxes
        return losses

    def loss_masks(self, outputs, targets, indices, num_boxes):
        """
        Compute the losses related to the masks: the focal loss and the dice loss. Targets dicts must contain the key
        "masks" containing a tensor of dim [nb_target_boxes, h, w].
        """
        if "pred_masks" not in outputs:
            raise KeyError("No predicted masks found in outputs")

        source_idx = self._get_source_permutation_idx(indices)
        target_idx = self._get_target_permutation_idx(indices)
        source_masks = outputs["pred_masks"]
        source_masks = source_masks[source_idx]
        masks = [t["masks"] for t in targets]
        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
        target_masks = target_masks.to(source_masks)
        target_masks = target_masks[target_idx]

        # upsample predictions to the target size
        source_masks = nn.functional.interpolate(
            source_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
        )
        source_masks = source_masks[:, 0].flatten(1)

        target_masks = target_masks.flatten(1)
        target_masks = target_masks.view(source_masks.shape)
        losses = {
            "loss_mask": sigmoid_focal_loss(source_masks, target_masks, num_boxes),
            "loss_dice": dice_loss(source_masks, target_masks, num_boxes),
        }
        return losses

    def loss_labels_bce(self, outputs, targets, indices, num_boxes, log=True):
        src_logits = outputs["logits"]
        idx = self._get_source_permutation_idx(indices)
        target_classes_original = torch.cat([_target["class_labels"][i] for _target, (_, i) in zip(targets, indices)])
        target_classes = torch.full(
            src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
        )
        target_classes[idx] = target_classes_original

        target = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1]
        loss = F.binary_cross_entropy_with_logits(src_logits, target * 1.0, reduction="none")
        loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes
        return {"loss_bce": loss}

    def _get_source_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
        source_idx = torch.cat([source for (source, _) in indices])
        return batch_idx, source_idx

    def _get_target_permutation_idx(self, indices):
        # permute targets following indices
        batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
        target_idx = torch.cat([target for (_, target) in indices])
        return batch_idx, target_idx

    def loss_labels_focal(self, outputs, targets, indices, num_boxes, log=True):
        if "logits" not in outputs:
            raise KeyError("No logits found in outputs")

        src_logits = outputs["logits"]

        idx = self._get_source_permutation_idx(indices)
        target_classes_original = torch.cat([_target["class_labels"][i] for _target, (_, i) in zip(targets, indices)])
        target_classes = torch.full(
            src_logits.shape[:2], self.num_classes, dtype=torch.int64, device=src_logits.device
        )
        target_classes[idx] = target_classes_original

        target = F.one_hot(target_classes, num_classes=self.num_classes + 1)[..., :-1]
        loss = sigmoid_focal_loss(src_logits, target, self.alpha, self.gamma)
        loss = loss.mean(1).sum() * src_logits.shape[1] / num_boxes
        return {"loss_focal": loss}

    def get_loss(self, loss, outputs, targets, indices, num_boxes):
        loss_map = {
            "labels": self.loss_labels,
            "cardinality": self.loss_cardinality,
            "boxes": self.loss_boxes,
            "masks": self.loss_masks,
            "bce": self.loss_labels_bce,
            "focal": self.loss_labels_focal,
            "vfl": self.loss_labels_vfl,
        }
        if loss not in loss_map:
            raise ValueError(f"Loss {loss} not supported")
        return loss_map[loss](outputs, targets, indices, num_boxes)

    @staticmethod
    def get_cdn_matched_indices(dn_meta, targets):
        dn_positive_idx, dn_num_group = dn_meta["dn_positive_idx"], dn_meta["dn_num_group"]
        num_gts = [len(t["class_labels"]) for t in targets]
        device = targets[0]["class_labels"].device

        dn_match_indices = []
        for i, num_gt in enumerate(num_gts):
            if num_gt > 0:
                gt_idx = torch.arange(num_gt, dtype=torch.int64, device=device)
                gt_idx = gt_idx.tile(dn_num_group)
                assert len(dn_positive_idx[i]) == len(gt_idx)
                dn_match_indices.append((dn_positive_idx[i], gt_idx))
            else:
                dn_match_indices.append(
                    (
                        torch.zeros(0, dtype=torch.int64, device=device),
                        torch.zeros(0, dtype=torch.int64, device=device),
                    )
                )

        return dn_match_indices

    def forward(self, outputs, targets):
        """
        This performs the loss computation.

        Args:
             outputs (`dict`, *optional*):
                Dictionary of tensors, see the output specification of the model for the format.
             targets (`List[dict]`, *optional*):
                List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
                losses applied, see each loss' doc.
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if "auxiliary_outputs" not in k}

        # Retrieve the matching between the outputs of the last layer and the targets
        indices = self.matcher(outputs_without_aux, targets)

        # Compute the average number of target boxes across all nodes, for normalization purposes
        num_boxes = sum(len(t["class_labels"]) for t in targets)
        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
        num_boxes = torch.clamp(num_boxes, min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            l_dict = self.get_loss(loss, outputs, targets, indices, num_boxes)
            l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict}
            losses.update(l_dict)

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if "auxiliary_outputs" in outputs:
            for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
                indices = self.matcher(auxiliary_outputs, targets)
                for loss in self.losses:
                    if loss == "masks":
                        # Intermediate masks losses are too costly to compute, we ignore them.
                        continue
                    l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
                    l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict}
                    l_dict = {k + f"_aux_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)

        # In case of cdn auxiliary losses. For rtdetr
        if "dn_auxiliary_outputs" in outputs:
            if "denoising_meta_values" not in outputs:
                raise ValueError(
                    "The output must have the 'denoising_meta_values` key. Please, ensure that 'outputs' includes a 'denoising_meta_values' entry."
                )
            indices = self.get_cdn_matched_indices(outputs["denoising_meta_values"], targets)
            num_boxes = num_boxes * outputs["denoising_meta_values"]["dn_num_group"]

            for i, auxiliary_outputs in enumerate(outputs["dn_auxiliary_outputs"]):
                # indices = self.matcher(auxiliary_outputs, targets)
                for loss in self.losses:
                    if loss == "masks":
                        # Intermediate masks losses are too costly to compute, we ignore them.
                        continue
                    kwargs = {}
                    l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes, **kwargs)
                    l_dict = {k: l_dict[k] * self.weight_dict[k] for k in l_dict if k in self.weight_dict}
                    l_dict = {k + f"_dn_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)



SUMMARY:
This is a loss computation module written in Python, specifically designed to work with deep learning models such as DETR (DEtection TRansformer). The module allows users to compute various losses (auxiliary and main losses) for object detection tasks.

**Key Features:**

1. **Loss Computation:** The module computes different types of losses, including auxiliary losses and main losses.
2. **Weighted Losses:** The module computes weighted losses based on a `weight_dict` provided by the user.
3. **Detection Modules:** The module is designed to work with detection modules such as DETR.

**Usage:**

1. Instantiate the loss computation module using the `LossComputation` constructor.
2. Provide the model's outputs and targets as input to the `forward` method.
3. Compute the losses using the `get_loss` method or directly by calling the `forward` method.

**Example Use Case:**
```python
# Instantiate the LossComputation module
loss_computer = LossComputation()

# Provide the model's outputs and targets
outputs = ...  # outputs from the model
targets = ...  # targets for the dataset

# Compute the losses using the get_loss method
loss_dict = loss_computer.get_loss("main_loss", outputs, targets)

# Print the computed losses
print(loss_dict)
```
**Output:**
```python
{
    "loss_main": {
        "loss1.0": 0.1234,
        "loss2.0": 0.5678
    },
   ...
}
```
The output is a dictionary containing the computed losses for each loss type (e.g., `loss_main`, `aux_loss_), where each key corresponds to a specific loss and its value is the corresponding loss value.


key: chroma:document
string_value:
CODE:

def RTDetrForObjectDetectionLoss(
    logits,
    labels,
    device,
    pred_boxes,
    config,
    outputs_class=None,
    outputs_coord=None,
    enc_topk_logits=None,
    enc_topk_bboxes=None,
    denoising_meta_values=None,
    **kwargs,
):
    criterion = RTDetrLoss(config)
    criterion.to(device)
    # Second: compute the losses, based on outputs and labels
    outputs_loss = {}
    outputs_loss["logits"] = logits
    outputs_loss["pred_boxes"] = pred_boxes
    if config.auxiliary_loss:
        if denoising_meta_values is not None:
            dn_out_coord, outputs_coord = torch.split(outputs_coord, denoising_meta_values["dn_num_split"], dim=2)
            dn_out_class, outputs_class = torch.split(outputs_class, denoising_meta_values["dn_num_split"], dim=2)

        auxiliary_outputs = _set_aux_loss(outputs_class[:, :-1].transpose(0, 1), outputs_coord[:, :-1].transpose(0, 1))
        outputs_loss["auxiliary_outputs"] = auxiliary_outputs
        outputs_loss["auxiliary_outputs"].extend(_set_aux_loss([enc_topk_logits], [enc_topk_bboxes]))
        if denoising_meta_values is not None:
            outputs_loss["dn_auxiliary_outputs"] = _set_aux_loss(
                dn_out_class.transpose(0, 1), dn_out_coord.transpose(0, 1)
            )
            outputs_loss["denoising_meta_values"] = denoising_meta_values

    loss_dict = criterion(outputs_loss, labels)

    loss = sum(loss_dict.values())


SUMMARY:
**Summary**

`RTDetrForObjectDetectionLoss`: computes object detection loss based on model outputs and ground truth labels.

### Main Purpose and Functionality

Calculates the object detection loss using the `RTDetrLoss` criterion. It takes in model outputs (`logits`, `pred_boxes`) and labels, as well as configuration settings (`config`).

### Key Implementation Details

* Computes losses for main output labels and auxiliary outputs using `_set_aux_loss`.
* Applies denoising metadata augmentation if available.
* Uses `RTDetrLoss` criterion to compute final loss.

### Notable Dependencies or Requirements

* `RTDetrLoss`: the loss function used by the code.
* `device`: the device where the model operates ( GPU/ CPU).
* `pred_boxes`, `logits`, `labels`, and related tensors.

### Code Snippet Extracted
```python
criterion = RTDetrLoss(config)
criterion.to(device)

outputs_loss = {
    "logits": logits,
    "pred_boxes": pred_boxes
}

# Additional losses if auxiliary loss is enabled
if config.auxiliary_loss:
    # ...
    outputs_loss["auxiliary_outputs"] = _set_aux_loss(outputs_class[:, :-1].transpose(0, 1), outputs_coord[:, :-1].transpose(0, 1))

loss_dict = criterion(outputs_loss, labels)
loss = sum(loss_dict.values())
```


key: chroma:document
string_value:
CODE:

class DeformableDetrHungarianMatcher(HungarianMatcher):
    @torch.no_grad()
    def forward(self, outputs, targets):
        """
        Differences:
        - out_prob = outputs["logits"].flatten(0, 1).sigmoid() instead of softmax
        - class_cost uses alpha and gamma
        """
        batch_size, num_queries = outputs["logits"].shape[:2]

        # We flatten to compute the cost matrices in a batch
        out_prob = outputs["logits"].flatten(0, 1).sigmoid()  # [batch_size * num_queries, num_classes]
        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]

        # Also concat the target labels and boxes
        target_ids = torch.cat([v["class_labels"] for v in targets])
        target_bbox = torch.cat([v["boxes"] for v in targets])

        # Compute the classification cost.
        alpha = 0.25
        gamma = 2.0
        neg_cost_class = (1 - alpha) * (out_prob**gamma) * (-(1 - out_prob + 1e-8).log())
        pos_cost_class = alpha * ((1 - out_prob) ** gamma) * (-(out_prob + 1e-8).log())
        class_cost = pos_cost_class[:, target_ids] - neg_cost_class[:, target_ids]

        # Compute the L1 cost between boxes
        bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)

        # Compute the giou cost between boxes
        giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))

        # Final cost matrix
        cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
        cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()

        sizes = [len(v["boxes"]) for v in targets]
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]


SUMMARY:
**Class Summary: DeformableDetrHungarianMatcher**

* **Purpose:** A variant of the Hungarian matching algorithm for object detection using Deformable Detectors (DeDt) and the DETR model.

**Key Implementation Details:**

* Uses `torch.no_grad()` to disable gradient computation.
* Computes class cost matrix with alpha and gamma weights, applying sigmoid activation to output logits.
* Concatenates target labels and boxes for indexing.
* Employs L1 cost between bounding boxes and GIou cost between box pairs.

**Notable Dependencies and Requirements:**

* TORCh (torch library)
* Generalized Box Intersection over Union (GIou) calculation
* Linear sum assignment with Hungarian embedding


key: chroma:document
string_value:
CODE:

class DeformableDetrImageLoss(ImageLoss):
    def __init__(self, matcher, num_classes, focal_alpha, losses):
        nn.Module.__init__(self)
        self.matcher = matcher
        self.num_classes = num_classes
        self.focal_alpha = focal_alpha
        self.losses = losses

    # removed logging parameter, which was part of the original implementation
    def loss_labels(self, outputs, targets, indices, num_boxes):
        """
        Classification loss (Binary focal loss) targets dicts must contain the key "class_labels" containing a tensor
        of dim [nb_target_boxes]
        """
        if "logits" not in outputs:
            raise KeyError("No logits were found in the outputs")
        source_logits = outputs["logits"]

        idx = self._get_source_permutation_idx(indices)
        target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(
            source_logits.shape[:2], self.num_classes, dtype=torch.int64, device=source_logits.device
        )
        target_classes[idx] = target_classes_o

        target_classes_onehot = torch.zeros(
            [source_logits.shape[0], source_logits.shape[1], source_logits.shape[2] + 1],
            dtype=source_logits.dtype,
            layout=source_logits.layout,
            device=source_logits.device,
        )
        target_classes_onehot.scatter_(2, target_classes.unsqueeze(-1), 1)

        target_classes_onehot = target_classes_onehot[:, :, :-1]
        loss_ce = (
            sigmoid_focal_loss(source_logits, target_classes_onehot, num_boxes, alpha=self.focal_alpha, gamma=2)
            * source_logits.shape[1]
        )
        losses = {"loss_ce": loss_ce}



SUMMARY:
**DeformableDetrImageLoss Class**

The `DeformableDetrImageLoss` class inherits from `ImageLoss`. Its primary purpose is to compute the classification loss (Binary focal loss) for Deformable DETR.

**Key Functionality:**

*  Computes Binary Focal Loss with sigmoid function.
*  Uses permutation indices and target label mapping to handle class labels targets.

**Implementation Details:**

*   Takes `matcher`, `num_classes`, `focal_alpha`, and `losses` as inputs in the constructor.
*   The `loss_labels` method computes the loss by:
    *   Getting source logits from outputs.
    *   Mapping target classes and their indices using permutation.
    *   Converting to one-hot form for sigmoid focal loss computation.

**Notable Dependencies/Requirements:**

*   `torch`: For tensor operations (e.g., `torch.cat`, `torch.zeros`, `torch.scatter_`).
*   `sigmoid_focal_loss`: Custom function or module not shown in this code snippet, required for computing the Binary Focal Loss with sigmoid.


key: chroma:document
string_value:
CODE:

def DeformableDetrForSegmentationLoss(
    logits, labels, device, pred_boxes, pred_masks, config, outputs_class=None, outputs_coord=None, **kwargs
):
    # First: create the matcher
    matcher = HungarianMatcher(class_cost=config.class_cost, bbox_cost=config.bbox_cost, giou_cost=config.giou_cost)
    # Second: create the criterion
    losses = ["labels", "boxes", "cardinality", "masks"]
    criterion = DeformableDetrImageLoss(
        matcher=matcher,
        num_classes=config.num_labels,
        focal_alpha=config.focal_alpha,
        losses=losses,
    )
    criterion.to(device)
    # Third: compute the losses, based on outputs and labels
    outputs_loss = {}
    outputs_loss["logits"] = logits
    outputs_loss["pred_boxes"] = pred_boxes
    outputs_loss["pred_masks"] = pred_masks

    auxiliary_outputs = None
    if config.auxiliary_loss:
        auxiliary_outputs = _set_aux_loss(outputs_class, outputs_coord)
        outputs_loss["auxiliary_outputs"] = auxiliary_outputs

    loss_dict = criterion(outputs_loss, labels)
    # Fourth: compute total loss, as a weighted sum of the various losses
    weight_dict = {"loss_ce": 1, "loss_bbox": config.bbox_loss_coefficient}
    weight_dict["loss_giou"] = config.giou_loss_coefficient
    weight_dict["loss_mask"] = config.mask_loss_coefficient
    weight_dict["loss_dice"] = config.dice_loss_coefficient
    if config.auxiliary_loss:
        aux_weight_dict = {}
        for i in range(config.decoder_layers - 1):
            aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)


SUMMARY:
**DeformableDetrForSegmentationLoss Function Summary**

**Main Purpose:**
The `DeformableDetrForSegmentationLoss` function computes the segmentation loss for a Deformable DETR model. It takes predictions and labels as inputs, along with configuration parameters.

**Key Implementation Details:**

* Creates a Hungarian matcher for class cost assignment.
* Initializes a Deformable DETR image loss criterion with configured parameters.
* Computes losses from predicted outputs and ground-truth labels.
* Computes weighted sum of losses using pre-defined weight dictionary.

**Notable Dependencies/Requirements:**

* `HungarianMatcher` class
* `DeformableDetrImageLoss` class
* `config` object (class cost, focal alpha, box loss coefficient, etc.)
* `labels`, `pred_boxes`, `pred_masks`, and `outputs_class` inputs

**Method Signatures and Interfaces:**

* `criterion(outputs_loss, labels)`
* `_set_aux_loss(outputs_class, outputs_coord)`


key: chroma:document
string_value:
CODE:

def DeformableDetrForObjectDetectionLoss(
    logits, labels, device, pred_boxes, config, outputs_class=None, outputs_coord=None, **kwargs
):
    # First: create the matcher
    matcher = DeformableDetrHungarianMatcher(
        class_cost=config.class_cost, bbox_cost=config.bbox_cost, giou_cost=config.giou_cost
    )
    # Second: create the criterion
    losses = ["labels", "boxes", "cardinality"]
    criterion = DeformableDetrImageLoss(
        matcher=matcher,
        num_classes=config.num_labels,
        focal_alpha=config.focal_alpha,
        losses=losses,
    )
    criterion.to(device)
    # Third: compute the losses, based on outputs and labels
    outputs_loss = {}
    auxiliary_outputs = None
    outputs_loss["logits"] = logits
    outputs_loss["pred_boxes"] = pred_boxes
    if config.auxiliary_loss:
        auxiliary_outputs = _set_aux_loss(outputs_class, outputs_coord)
        outputs_loss["auxiliary_outputs"] = auxiliary_outputs

    loss_dict = criterion(outputs_loss, labels)
    # Fourth: compute total loss, as a weighted sum of the various losses
    weight_dict = {"loss_ce": 1, "loss_bbox": config.bbox_loss_coefficient}
    weight_dict["loss_giou"] = config.giou_loss_coefficient
    if config.auxiliary_loss:
        aux_weight_dict = {}
        for i in range(config.decoder_layers - 1):
            aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)
    loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)


SUMMARY:
**Function Summary: DeformableDetrForObjectDetectionLoss**

* **Purpose:** Computes the object detection loss function for deformable Detr models.

**Key Implementation Details:**

* Creates a `DeformableDetrHungarianMatcher` to measure class and bounding box costs.
* Employs a `DeformableDetrImageLoss` criterion to compute losses from outputs and labels.
* Weights losses using coefficients defined in the `config` dictionary.

**Notable Dependencies and Requirements:**

* `DeformableDetrHungarianMatcher`
	+ `class_cost`, `bbox_cost`, `giou_cost` config parameters
* `DeformableDetrImageLoss`
	+ `matcher`, `num_classes`, `focal_alpha`, `losses` config parameters

**Method Signatures and Interfaces:**

* `deformable_detr_object_detection_loss`: computes the object detection loss.
	+ Inputs: `logits`, `labels`, `device`, `pred_boxes`, `config`
	+ Outputs: `loss`, `loss_dict`, `auxiliary_outputs`


key: chroma:document
string_value:
CODE:

def dice_loss(inputs, targets, num_boxes):
    """
    Compute the DICE loss, similar to generalized IOU for masks

    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs (0 for the negative class and 1 for the positive
                 class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1)
    numerator = 2 * (inputs * targets).sum(1)
    denominator = inputs.sum(-1) + targets.sum(-1)
    loss = 1 - (numerator + 1) / (denominator + 1)


SUMMARY:
**Summary**

The `dice_loss` function computes the DICE loss, a measure of similarity between predictions and ground truth labels. It takes in:

* `inputs`: Predictions tensor with shape arbitrary to input
* `targets`: Target label tensor with same shape as inputs (0 negative class, 1 positive class)
* `num_boxes`: Number of examples

The function returns the average DICE loss across all examples.

**Key Implementation Details**

* Applies sigmoid activation to predictions (`inputs.sigmoid()`)
* Flattens input tensor (`inputs.flatten(1)`)

**Notable Dependencies and Requirements**

* None explicit dependencies, but assumed to be used with a library providing:
	+ `flatten` method on Tensor
	+ Sigmoid activation function on the predictions
	+ Sum operation over batch dimensions
* Assumes `num_boxes` is known during computation


key: chroma:document
string_value:
CODE:

def sigmoid_focal_loss(inputs, targets, num_boxes, alpha: float = 0.25, gamma: float = 2):
    """
    Loss used in RetinaNet for dense detection: https://arxiv.org/abs/1708.02002.

    Args:
        inputs (`torch.FloatTensor` of arbitrary shape):
            The predictions for each example.
        targets (`torch.FloatTensor` with the same shape as `inputs`)
            A tensor storing the binary classification label for each element in the `inputs` (0 for the negative class
            and 1 for the positive class).
        alpha (`float`, *optional*, defaults to `0.25`):
            Optional weighting factor in the range (0,1) to balance positive vs. negative examples.
        gamma (`int`, *optional*, defaults to `2`):
            Exponent of the modulating factor (1 - p_t) to balance easy vs hard examples.

    Returns:
        Loss tensor
    """
    prob = inputs.sigmoid()
    ce_loss = nn.functional.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
    # add modulating factor
    p_t = prob * targets + (1 - prob) * (1 - targets)
    loss = ce_loss * ((1 - p_t) ** gamma)

    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        loss = alpha_t * loss



SUMMARY:
**Summary:**

```python
**Sigmoid Focal Loss Functionality:**

The `sigmoid_focal_loss` function calculates the focal loss used in RetinaNet for dense detection. It takes as input:

* `inputs`: prediction tensor with arbitrary shape
* `targets`: binary classification target tensor with the same shape as `inputs`
* `alpha`: weighting factor (default: 0.25) and `gamma`: exponent of modulating factor (default: 2)

The function returns the focal loss value, which balances easy and hard examples through the `(1 - p_t) ** gamma` term.

**Key Implementation Details:**

* Calculates binary cross-entropy loss (`ce_loss`) using `nn.functional.binary_cross_entropy_with_logits`
* Applies modulating factor to balance positive and negative examples
* Optionally applies weighting factor (`alpha`) based on target tensor

**Notable Dependencies/Requirements:**

* PyTorch (`torch.FloatTensor`, `nn.functional`)
```


key: chroma:document
string_value:
CODE:
# taken from https://github.com/facebookresearch/detr/blob/master/models/detr.py
class ImageLoss(nn.Module):
    """
    This class computes the losses for DetrForObjectDetection/DetrForSegmentation. The process happens in two steps: 1)
    we compute hungarian assignment between ground truth boxes and the outputs of the model 2) we supervise each pair
    of matched ground-truth / prediction (supervise class and box).

    A note on the `num_classes` argument (copied from original repo in detr.py): "the naming of the `num_classes`
    parameter of the criterion is somewhat misleading. It indeed corresponds to `max_obj_id` + 1, where `max_obj_id` is
    the maximum id for a class in your dataset. For example, COCO has a `max_obj_id` of 90, so we pass `num_classes` to
    be 91. As another example, for a dataset that has a single class with `id` 1, you should pass `num_classes` to be 2
    (`max_obj_id` + 1). For more details on this, check the following discussion
    https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223"


    Args:
        matcher (`DetrHungarianMatcher`):
            Module able to compute a matching between targets and proposals.
        num_classes (`int`):
            Number of object categories, omitting the special no-object category.
        eos_coef (`float`):
            Relative classification weight applied to the no-object category.
        losses (`List[str]`):
            List of all the losses to be applied. See `get_loss` for a list of all available losses.
    """

    def __init__(self, matcher, num_classes, eos_coef, losses):
        super().__init__()
        self.matcher = matcher
        self.num_classes = num_classes
        self.eos_coef = eos_coef
        self.losses = losses
        empty_weight = torch.ones(self.num_classes + 1)
        empty_weight[-1] = self.eos_coef
        self.register_buffer("empty_weight", empty_weight)

    # removed logging parameter, which was part of the original implementation
    def loss_labels(self, outputs, targets, indices, num_boxes):
        """
        Classification loss (NLL) targets dicts must contain the key "class_labels" containing a tensor of dim
        [nb_target_boxes]
        """
        if "logits" not in outputs:
            raise KeyError("No logits were found in the outputs")
        source_logits = outputs["logits"]

        idx = self._get_source_permutation_idx(indices)
        target_classes_o = torch.cat([t["class_labels"][J] for t, (_, J) in zip(targets, indices)])
        target_classes = torch.full(
            source_logits.shape[:2], self.num_classes, dtype=torch.int64, device=source_logits.device
        )
        target_classes[idx] = target_classes_o

        loss_ce = nn.functional.cross_entropy(source_logits.transpose(1, 2), target_classes, self.empty_weight)
        losses = {"loss_ce": loss_ce}

        return losses

    @torch.no_grad()
    def loss_cardinality(self, outputs, targets, indices, num_boxes):
        """
        Compute the cardinality error, i.e. the absolute error in the number of predicted non-empty boxes.

        This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients.
        """
        logits = outputs["logits"]
        device = logits.device
        target_lengths = torch.as_tensor([len(v["class_labels"]) for v in targets], device=device)
        # Count the number of predictions that are NOT "no-object" (which is the last class)
        card_pred = (logits.argmax(-1) != logits.shape[-1] - 1).sum(1)
        card_err = nn.functional.l1_loss(card_pred.float(), target_lengths.float())
        losses = {"cardinality_error": card_err}
        return losses

    def loss_boxes(self, outputs, targets, indices, num_boxes):
        """
        Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss.

        Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes
        are expected in format (center_x, center_y, w, h), normalized by the image size.
        """
        if "pred_boxes" not in outputs:
            raise KeyError("No predicted boxes found in outputs")
        idx = self._get_source_permutation_idx(indices)
        source_boxes = outputs["pred_boxes"][idx]
        target_boxes = torch.cat([t["boxes"][i] for t, (_, i) in zip(targets, indices)], dim=0)

        loss_bbox = nn.functional.l1_loss(source_boxes, target_boxes, reduction="none")

        losses = {}
        losses["loss_bbox"] = loss_bbox.sum() / num_boxes

        loss_giou = 1 - torch.diag(
            generalized_box_iou(center_to_corners_format(source_boxes), center_to_corners_format(target_boxes))
        )
        losses["loss_giou"] = loss_giou.sum() / num_boxes
        return losses

    def loss_masks(self, outputs, targets, indices, num_boxes):
        """
        Compute the losses related to the masks: the focal loss and the dice loss.

        Targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w].
        """
        if "pred_masks" not in outputs:
            raise KeyError("No predicted masks found in outputs")

        source_idx = self._get_source_permutation_idx(indices)
        target_idx = self._get_target_permutation_idx(indices)
        source_masks = outputs["pred_masks"]
        source_masks = source_masks[source_idx]
        masks = [t["masks"] for t in targets]
        # TODO use valid to mask invalid areas due to padding in loss
        target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
        target_masks = target_masks.to(source_masks)
        target_masks = target_masks[target_idx]

        # upsample predictions to the target size
        source_masks = nn.functional.interpolate(
            source_masks[:, None], size=target_masks.shape[-2:], mode="bilinear", align_corners=False
        )
        source_masks = source_masks[:, 0].flatten(1)

        target_masks = target_masks.flatten(1)
        target_masks = target_masks.view(source_masks.shape)
        losses = {
            "loss_mask": sigmoid_focal_loss(source_masks, target_masks, num_boxes),
            "loss_dice": dice_loss(source_masks, target_masks, num_boxes),
        }
        return losses

    def _get_source_permutation_idx(self, indices):
        # permute predictions following indices
        batch_idx = torch.cat([torch.full_like(source, i) for i, (source, _) in enumerate(indices)])
        source_idx = torch.cat([source for (source, _) in indices])
        return batch_idx, source_idx

    def _get_target_permutation_idx(self, indices):
        # permute targets following indices
        batch_idx = torch.cat([torch.full_like(target, i) for i, (_, target) in enumerate(indices)])
        target_idx = torch.cat([target for (_, target) in indices])
        return batch_idx, target_idx

    def get_loss(self, loss, outputs, targets, indices, num_boxes):
        loss_map = {
            "labels": self.loss_labels,
            "cardinality": self.loss_cardinality,
            "boxes": self.loss_boxes,
            "masks": self.loss_masks,
        }
        if loss not in loss_map:
            raise ValueError(f"Loss {loss} not supported")
        return loss_map[loss](outputs, targets, indices, num_boxes)

    def forward(self, outputs, targets):
        """
        This performs the loss computation.

        Args:
             outputs (`dict`, *optional*):
                Dictionary of tensors, see the output specification of the model for the format.
             targets (`List[dict]`, *optional*):
                List of dicts, such that `len(targets) == batch_size`. The expected keys in each dict depends on the
                losses applied, see each loss' doc.
        """
        outputs_without_aux = {k: v for k, v in outputs.items() if k != "auxiliary_outputs"}

        # Retrieve the matching between the outputs of the last layer and the targets
        indices = self.matcher(outputs_without_aux, targets)

        # Compute the average number of target boxes across all nodes, for normalization purposes
        num_boxes = sum(len(t["class_labels"]) for t in targets)
        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
        world_size = 1
        if is_accelerate_available():
            if PartialState._shared_state != {}:
                num_boxes = reduce(num_boxes)
                world_size = PartialState().num_processes
        num_boxes = torch.clamp(num_boxes / world_size, min=1).item()

        # Compute all the requested losses
        losses = {}
        for loss in self.losses:
            losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))

        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
        if "auxiliary_outputs" in outputs:
            for i, auxiliary_outputs in enumerate(outputs["auxiliary_outputs"]):
                indices = self.matcher(auxiliary_outputs, targets)
                for loss in self.losses:
                    if loss == "masks":
                        # Intermediate masks losses are too costly to compute, we ignore them.
                        continue
                    l_dict = self.get_loss(loss, auxiliary_outputs, targets, indices, num_boxes)
                    l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
                    losses.update(l_dict)



SUMMARY:
This code defines the loss function for a Detracker (DEtection TranSnearer) model. It is used in object detection and tracking tasks.

**Overview**

The code provides a framework for computing various losses associated with the output of an object detection or tracking model. The losses include:

*   `labels`: computes the cross-entropy loss between predicted labels and actual labels.
*   `cardinality`: computes the cardinality loss, which measures the difference in number of proposals between the predicted and actual detections.
*   `boxes`: computes the IoU loss between predicted box coordinates and actual bounding boxes.
*   `masks`: computes the focal loss and Dice loss between predicted masks and actual masks.

**Method Structure**

The code uses a modular structure to organize the different components of the loss computation. Each loss function is represented as a method (`loss_labels`, `loss_cardinality`, etc.) that takes in the output predictions, target data, indices (for efficient indexing), and the number of boxes as inputs.

**Output**

The overall `forward` method computes all the requested losses, which are stored in a dictionary called `losses`. This dictionary contains the computed values for each loss function, keyed according to the input loss name followed by an underscore indicating it's an intermediate output (for auxiliary losses). The final value returns just the relevant losses.

The methods take advantage of `torch`'s tensor operations and indexing capabilities. In particular:

*   `_get_source_permutation_idx` and `_get_target_permutation_idx` are helper functions that move predicted and target indices within a batch according to the model's structure
*   The `num_boxes` computation is done considering world size as provided by `is_accelerate_available`
*   When auxiliary outputs are present, duplicate computations over intermediate layers are avoided through modularity


key: chroma:document
string_value:
CODE:
# taken from https://github.com/facebookresearch/detr/blob/master/models/matcher.py
class HungarianMatcher(nn.Module):
    """
    This class computes an assignment between the targets and the predictions of the network.

    For efficiency reasons, the targets don't include the no_object. Because of this, in general, there are more
    predictions than targets. In this case, we do a 1-to-1 matching of the best predictions, while the others are
    un-matched (and thus treated as non-objects).

    Args:
        class_cost:
            The relative weight of the classification error in the matching cost.
        bbox_cost:
            The relative weight of the L1 error of the bounding box coordinates in the matching cost.
        giou_cost:
            The relative weight of the giou loss of the bounding box in the matching cost.
    """

    def __init__(self, class_cost: float = 1, bbox_cost: float = 1, giou_cost: float = 1):
        super().__init__()
        requires_backends(self, ["scipy"])

        self.class_cost = class_cost
        self.bbox_cost = bbox_cost
        self.giou_cost = giou_cost
        if class_cost == 0 and bbox_cost == 0 and giou_cost == 0:
            raise ValueError("All costs of the Matcher can't be 0")

    @torch.no_grad()
    def forward(self, outputs, targets):
        """
        Args:
            outputs (`dict`):
                A dictionary that contains at least these entries:
                * "logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
                * "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates.
            targets (`List[dict]`):
                A list of targets (len(targets) = batch_size), where each target is a dict containing:
                * "class_labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of
                  ground-truth
                 objects in the target) containing the class labels
                * "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates.

        Returns:
            `List[Tuple]`: A list of size `batch_size`, containing tuples of (index_i, index_j) where:
            - index_i is the indices of the selected predictions (in order)
            - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds: len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        batch_size, num_queries = outputs["logits"].shape[:2]

        # We flatten to compute the cost matrices in a batch
        out_prob = outputs["logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]

        # Also concat the target labels and boxes
        target_ids = torch.cat([v["class_labels"] for v in targets])
        target_bbox = torch.cat([v["boxes"] for v in targets])

        # Compute the classification cost. Contrary to the loss, we don't use the NLL,
        # but approximate it in 1 - proba[target class].
        # The 1 is a constant that doesn't change the matching, it can be ommitted.
        class_cost = -out_prob[:, target_ids]

        # Compute the L1 cost between boxes
        bbox_cost = torch.cdist(out_bbox, target_bbox, p=1)

        # Compute the giou cost between boxes
        giou_cost = -generalized_box_iou(center_to_corners_format(out_bbox), center_to_corners_format(target_bbox))

        # Final cost matrix
        cost_matrix = self.bbox_cost * bbox_cost + self.class_cost * class_cost + self.giou_cost * giou_cost
        cost_matrix = cost_matrix.view(batch_size, num_queries, -1).cpu()

        sizes = [len(v["boxes"]) for v in targets]
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(cost_matrix.split(sizes, -1))]


SUMMARY:
**HungarianMatcher Class Summary**

### Main Purpose and Functionality

This `HungarianMatcher` class computes an assignment between the targets and predictions of a network. It performs a 1-to-1 matching of the best predictions, treating others as non-matching (non-objects).

### Implementation Details

*   **Methods:**
    *   `__init__`: Initializes the matcher with class, bbox, and giou costs.
        *   Parameters: `class_cost`, `bbox_cost`, `giou_cost`
        *   Exceptions: Raises ValueError if all costs are 0
    *   `forward`: Computes the cost matrix using classification, L1 box distance, and giou loss.
        *   Parameters:
            *   `outputs` (dict): Contains "logits" and "pred_boxes"
            *   `targets` (List[dict]): List of targets with class labels and boxes
    *   Returns:
        A list of tuples containing selected prediction and target indices

### Notable Dependencies and Requirements

*   **External Libraries:**
    *   scipy (required for linear_sum_assignment)
*   **Interfaces:**
    *   torch.Tensor
    *  torch.nn.Module
    *   scipy.linear_sum_assignment
*   **Inherited Interfaces:**
    None


key: docstring
string_value:
# below: bounding box utilities taken from https://github.com/facebookresearch/detr/blob/master/util/box_ops.py


key: chroma:document
string_value:
CODE:

def _upcast(t: Tensor) -> Tensor:
    # Protects from numerical overflows in multiplications by upcasting to the equivalent higher type
    if t.is_floating_point():
        return t if t.dtype in (torch.float32, torch.float64) else t.float()
    else:


SUMMARY:
**Summary:**

`_upcast` function upcasts or downcasts tensor data type to prevent numerical overflows in multiplications, based on its float/int type.

**Key Details:**

- Upcasts `float32/float64` tensors and downcasts integer types
- Applies similar conversion logic for int types
- Returns original tensor with modified dtype if not exceeding numerical limits

**Notable Dependencies:**

- Uses PyTorch library (`torch.float32`, `torch.float64`, `torch.int32`, `torch.int64`)
- References PyTorch data type operations (`is_floating_point()`)


key: chroma:document
string_value:
CODE:

def box_area(boxes: Tensor) -> Tensor:
    """
    Computes the area of a set of bounding boxes, which are specified by its (x1, y1, x2, y2) coordinates.

    Args:
        boxes (`torch.FloatTensor` of shape `(number_of_boxes, 4)`):
            Boxes for which the area will be computed. They are expected to be in (x1, y1, x2, y2) format with `0 <= x1
            < x2` and `0 <= y1 < y2`.

    Returns:
        `torch.FloatTensor`: a tensor containing the area for each box.
    """
    boxes = _upcast(boxes)


SUMMARY:
**Summary:**

* **Purpose**: Computes the area of a set of bounding boxes given their (x1, y1, x2, y2) coordinates.
* **Key Implementation Details**:
	+ Uses PyTorch tensor subtraction and multiplication.
	+ Returns the difference in widths and heights multiplied together for each box.
* **Dependencies and Requirements**:
	+ `torch.FloatTensor` input and output tensors with shape `(number_of_boxes, 4)` layout.

**Function Signature:** `box_area(boxes: torch.FloatTensor) -> torch.FloatTensor`


key: chroma:document
string_value:
CODE:
# modified from torchvision to also return the union
def box_iou(boxes1, boxes2):
    area1 = box_area(boxes1)
    area2 = box_area(boxes2)

    left_top = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
    right_bottom = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    width_height = (right_bottom - left_top).clamp(min=0)  # [N,M,2]
    inter = width_height[:, :, 0] * width_height[:, :, 1]  # [N,M]

    union = area1[:, None] + area2 - inter

    iou = inter / union


SUMMARY:
**Summary**

The `box_iou` function calculates the Intersection over Union (IoU) score between two sets of bounding boxes (`boxes1` and `boxes2`) using their areas, locations, and dimensions.

**Key Details:**

* **Area Calculations**: The function uses the `box_area` function to calculate the area of each bounding box in both input lists.
* **Intersection and Union**: It calculates the intersection area by taking the product of the overlaps along the x- and y-axis. The union is calculated by adding the individual areas and subtracting the intersection.

**Implementation Notes:**

* The function uses `torch.max` to find the maximum x-coordinate top-left corner and `torch.min` to find the minimum x-coordinate bottom-right corner.
* The resulting overlap in width and height is then masked to ensure non-positive values.
* Notable Method Signatures:
	+ `box_area`
	+ `torch.max`
	+ `torch.min`

**Dependencies:**

* `torch`
* `box_area` function, not shown in the code snippet.


key: chroma:document
string_value:
CODE:

def generalized_box_iou(boxes1, boxes2):
    """
    Generalized IoU from https://giou.stanford.edu/. The boxes should be in [x0, y0, x1, y1] (corner) format.

    Returns:
        `torch.FloatTensor`: a [N, M] pairwise matrix, where N = len(boxes1) and M = len(boxes2)
    """
    # degenerate boxes gives inf / nan results
    # so do an early check
    if not (boxes1[:, 2:] >= boxes1[:, :2]).all():
        raise ValueError(f"boxes1 must be in [x0, y0, x1, y1] (corner) format, but got {boxes1}")
    if not (boxes2[:, 2:] >= boxes2[:, :2]).all():
        raise ValueError(f"boxes2 must be in [x0, y0, x1, y1] (corner) format, but got {boxes2}")
    iou, union = box_iou(boxes1, boxes2)

    top_left = torch.min(boxes1[:, None, :2], boxes2[:, :2])
    bottom_right = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])

    width_height = (bottom_right - top_left).clamp(min=0)  # [N,M,2]
    area = width_height[:, :, 0] * width_height[:, :, 1]



SUMMARY:
**Summary**

`generalized_box_iou(boxes1, boxes2)` computes the Generalized IoU metric between two sets of bounding boxes. It takes two input tensors `boxes1` and `boxes2`, each in `[x0, y0, x1, y1] (corner) format`. 

**Key Implementation Details**

*   Uses `box_iou(boxes1, boxes2)` to calculate IoU, union area, and area.
*   Computes `top_left` and `bottom_right` coordinates by finding minimum and maximum box coordinates respectively.
*   Calculates `width_height` by subtracting `top_left` from `bottom_right`.
*   Returns the normalized GIoU value as `iou - (area - union) / area`.

**Notable Dependencies/Requirements**

*   Uses `torch` library for tensor operations.
*   Relies on `box_iou()` function to compute box intersections, unions, and areas.


key: chroma:document
string_value:
CODE:
# below: taken from https://github.com/facebookresearch/detr/blob/master/util/misc.py#L306
def _max_by_axis(the_list):
    # type: (List[List[int]]) -> List[int]
    maxes = the_list[0]
    for sublist in the_list[1:]:
        for index, item in enumerate(sublist):
            maxes[index] = max(maxes[index], item)


SUMMARY:
**Summary**

The `_max_by_axis` function computes the maximum values along each axis of a nested list of integers.

**Key Implementation Details**

* The `maxes` array is initialized with the first sublist's values and then iteratively updated with the maximum value for each position using the `enumerate`-enabled comparison.
* The function modifies the input `the_list[0]` (`maxes`) in-place to store the cumulative maximums.

**Notable Dependencies/Requirements**

* None, as this is a standalone Python implementation.


key: chroma:document
string_value:
CODE:

class NestedTensor:
    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors
        self.mask = mask

    def to(self, device):
        cast_tensor = self.tensors.to(device)
        mask = self.mask
        if mask is not None:
            cast_mask = mask.to(device)
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):
        return self.tensors, self.mask

    def __repr__(self):


SUMMARY:
**Summary:**

The `NestedTensor` class represents a tensor with optional mask. Its primary purpose is to:

*   **Wrap and manage tensors**: Store tensors and their corresponding masks, handling device changes through the `to()` method.
*   **Expose underlying data**: Provide a way to decompose the tensor and mask with the `decompose()` method.

**Key Implementation Details:**

*   **Device compatibility**: The `to()` method converts tensors to a specified target device, optionally updating the accompanying mask.
*   **Mask handling**: Tensors have an optional mask, which can be updated or left untouched during conversion.
*   **Decomposition**: The `decompose()` method returns both the tensor and its associated mask.

**Notable Dependencies/Requirements:**

*   Uses the `Tensor` class from a PyTorch library (referenced as `Optional[Tensor]`)
*   Utilizes the `to()` method on tensors, which may depend on external dependencies like PyTorch itself.


key: chroma:document
string_value:
CODE:

def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
    if tensor_list[0].ndim == 3:
        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
        batch_shape = [len(tensor_list)] + max_size
        batch_size, num_channels, height, width = batch_shape
        dtype = tensor_list[0].dtype
        device = tensor_list[0].device
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
        mask = torch.ones((batch_size, height, width), dtype=torch.bool, device=device)
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
            m[: img.shape[1], : img.shape[2]] = False
    else:
        raise ValueError("Only 3-dimensional tensors are supported")


SUMMARY:
**Summary**

`nested_tensor_from_tensor_list`: Constructs a `NestedTensor` object from a list of 3D tensors by padding the smallest tensor to match the shape of the largest one.

**Main Purpose and Functionality**
- Pads 3D tensors in the input list (`tensor_list`) to create a batch with consistent shape.
- Returns a `NestedTensor` object containing the padded tensor and its binary mask.

**Key Implementation Details**

* **Padding**: The function pads each img tensor in the list with zeros onto the largest dimension.
* **Shape Inference**: It uses `_max_by_axis` (not shown) to find the maximum size of the smallest tensor's shape along any axis.
* **Batch Shape Construction**: `[len(tensor_list)] + max_size` creates a batch shape with length equal to the number of tensors in the list and max size from the largest dimension.

**Notable Dependencies or Requirements**

* `torch`: PyTorch library for tensor operations.
* `_max_by_axis` (not shown): A function that calculates the maximum value along any axis, likely defined elsewhere in the codebase.


key: chroma:document
string_value:
CODE:

def ForSegmentationLoss(
    logits, labels, device, pred_boxes, pred_masks, config, outputs_class=None, outputs_coord=None, **kwargs
):
    # First: create the matcher
    matcher = HungarianMatcher(class_cost=config.class_cost, bbox_cost=config.bbox_cost, giou_cost=config.giou_cost)
    # Second: create the criterion
    losses = ["labels", "boxes", "cardinality", "masks"]
    criterion = ImageLoss(
        matcher=matcher,
        num_classes=config.num_labels,
        eos_coef=config.eos_coefficient,
        losses=losses,
    )
    criterion.to(device)
    # Third: compute the losses, based on outputs and labels
    outputs_loss = {}
    outputs_loss["logits"] = logits
    outputs_loss["pred_boxes"] = pred_boxes
    outputs_loss["pred_masks"] = pred_masks

    auxiliary_outputs = None
    if config.auxiliary_loss:
        auxiliary_outputs = _set_aux_loss(outputs_class, outputs_coord)
        outputs_loss["auxiliary_outputs"] = auxiliary_outputs

    loss_dict = criterion(outputs_loss, labels)
    # Fourth: compute total loss, as a weighted sum of the various losses
    weight_dict = {"loss_ce": 1, "loss_bbox": config.bbox_loss_coefficient}
    weight_dict["loss_giou"] = config.giou_loss_coefficient
    weight_dict["loss_mask"] = config.mask_loss_coefficient
    weight_dict["loss_dice"] = config.dice_loss_coefficient
    if config.auxiliary_loss:
        aux_weight_dict = {}
        for i in range(config.decoder_layers - 1):
            aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)
    loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)


SUMMARY:
**Function Summary: ForSegmentationLoss**

**Purpose and Functionality**
The `ForSegmentationLoss` function computes the total segmentation loss for a given model. It takes in model outputs (`logits`, `pred_boxes`, `pred_masks`) and labels, along with configuration parameters (e.g., class cost, bbox cost). The function calculates and returns the weighted sum of individual losses (CE, BBox, Giou, Mask, Dice), including optional auxiliary loss.

**Key Implementation Details**

* Creates a Hungarian matcher and ImageLoss criterion instance.
* Computes individual losses (labels, boxes, cardinality, masks) using `criterion`.
* Sets up auxiliary loss if enabled (`config.auxiliary_loss`).
* Calculates total loss by summing weighted individual losses.

**Notable Dependencies/Requirements**

* `HungarianMatcher` class with configuration parameters (`class_cost`, `bbox_cost`, `giou_cost`)
* `ImageLoss` class with configuration parameters (`num_classes`, `eos_coef`, `losses`)
* `config` object with various coefficients and loss weights.


key: chroma:document
string_value:
CODE:

def ForObjectDetectionLoss(
    logits, labels, device, pred_boxes, config, outputs_class=None, outputs_coord=None, **kwargs
):
    # First: create the matcher
    matcher = HungarianMatcher(class_cost=config.class_cost, bbox_cost=config.bbox_cost, giou_cost=config.giou_cost)
    # Second: create the criterion
    losses = ["labels", "boxes", "cardinality"]
    criterion = ImageLoss(
        matcher=matcher,
        num_classes=config.num_labels,
        eos_coef=config.eos_coefficient,
        losses=losses,
    )
    criterion.to(device)
    # Third: compute the losses, based on outputs and labels
    outputs_loss = {}
    auxiliary_outputs = None
    outputs_loss["logits"] = logits
    outputs_loss["pred_boxes"] = pred_boxes
    if config.auxiliary_loss:
        auxiliary_outputs = _set_aux_loss(outputs_class, outputs_coord)
        outputs_loss["auxiliary_outputs"] = auxiliary_outputs

    loss_dict = criterion(outputs_loss, labels)
    # Fourth: compute total loss, as a weighted sum of the various losses
    weight_dict = {"loss_ce": 1, "loss_bbox": config.bbox_loss_coefficient}
    weight_dict["loss_giou"] = config.giou_loss_coefficient
    if config.auxiliary_loss:
        aux_weight_dict = {}
        for i in range(config.decoder_layers - 1):
            aux_weight_dict.update({k + f"_{i}": v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)
    loss = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)


SUMMARY:
**Function Summary: ForObjectDetectionLoss**

This function implements the loss computation for object detection tasks using Mask R-CNN.

**Main Purpose and Functionality:**
The `ForObjectDetectionLoss` function computes the total loss for object detection, combining label-based losses (CE loss) with box-based losses (GIoU loss). It also supports auxiliary loss computation when enabled, which includes two sets of outputs (`outputs_class` and `outputs_coord`) mapped to the corresponding class or bounding box loss.

**Implementation Details:**

* Creates a Hungarian matcher for calculating class, box, and GIoU cost.
* Initializes an `ImageLoss` criterion with specified parameters.
* Computes loss values from output pairs (logits/pred_boxes) vs. labels.
* Calculates auxiliary loss using `_set_aux_loss` function if enabled.
* Scales losses by weight factors, combining label-based losses (1, BBox, GIoU), and optional augmented BBox loss scales.

**Notable Dependencies or Requirements:**

* `HungarianMatcher`, `ImageLoss`, `_set_aux_loss`
* External dependencies:
	+ None specified in this snippet
        • The HungarianMatcher class is likely an existing custom class.
        • ImageLoss, _set_aux_loss are potential classes to be created for a more full-fledged implementation.

**Method Signatures and Interfaces:**

* `HungarianMatcher`: class interface.
* `_set_aux_loss`: function interface (potentially external).

**Configurable Parameters:**

* `config.class_cost`, `config.bbox_cost`, `config.giou_cost`
* `config.num_labels`, `config.eos_coefficient`, `config.bbox_loss_coefficient`, `config.giou_loss_coefficient`

Note that this summary captures the core functionality and implementation details, focusing on precise technical terms and classes/interfaces.


key: chroma:document
string_value:
CODE:

def fixed_cross_entropy(source, target, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs):
    reduction = "sum" if num_items_in_batch is not None else "mean"
    loss = nn.functional.cross_entropy(source, target, ignore_index=ignore_index, reduction=reduction)
    if reduction == "sum":
        loss = loss / num_items_in_batch


SUMMARY:
**Summary**

This is a PyTorch implementation of the fixed cross-entropy loss function.

**Main Purpose and Functionality:**
Calculates the cross-entropy loss between the model's output (`source`) and the target labels (`target`). Adjusts for batch size via the `num_items_in_batch` parameter.

**Key Implementation Details:**

* Uses PyTorch's `nn.functional.cross_entropy` with custom reduction mode selection based on whether `num_items_in_batch` is provided.
* Scales loss by batch size if necessary to maintain consistency in reported values

**Dependencies and Requirements:**

* PyTorch (`nn.functional.cross_entropy`)
* [ Torch Neural Networks ] (`nn`)


key: chroma:document
string_value:
CODE:

def ForCausalLMLoss(
    logits, labels, vocab_size: int, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs
):
    # Upcast to float if we need to compute the loss to avoid potential precision issues
    logits = logits.float()
    labels = labels.to(logits.device)
    # Shift so that tokens < n predict n
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()

    # Flatten the tokens
    shift_logits = shift_logits.view(-1, vocab_size)
    shift_labels = shift_labels.view(-1)
    # Enable model parallelism
    shift_labels = shift_labels.to(shift_logits.device)
    loss = fixed_cross_entropy(shift_logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)


SUMMARY:
**Function Summary: ForCausalLMLoss**

*   **Main Purpose:** Computes the Causal LMLoss for sequence prediction models.
*   **Primary Functionality:** It takes in `logits` and `labels`, and returns the calculated loss value, while also performing token shifting, flattening, and model parallelism.

**Key Details:**

*   Uses `fixed_cross_entropy` function for computation
*   Performs token shifting with `shift_logits` and `shift_labels`
*   Flattens tokens using `view(-1, vocab_size)` on both logits and labels.
*   Employs model parallelism on `labels` to match device of `logits`.
*   Accepts optional arguments (`num_items_in_batch`, `ignore_index`) via default values and keyword arguments.

**Notable Dependencies:**

*   `fixed_cross_entropy` function is a dependency not included in the provided code snippet. (External dependency) 
*   Assumes access to PyTorch library for tensor operations.


key: chroma:document
string_value:
CODE:

def ForMaskedLMLoss(
    logits, labels, vocab_size: int, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs
):
    # Upcast to float if we need to compute the loss to avoid potential precision issues
    logits = logits.float()
    labels = labels.to(logits.device)

    # Flatten the tokens
    logits = logits.view(-1, vocab_size)
    labels = labels.view(-1)
    # Enable model parallelism

    labels = labels.to(logits.device)
    loss = fixed_cross_entropy(logits, labels, num_items_in_batch, ignore_index, **kwargs)


SUMMARY:
Here's a concise summary of the provided function code chunk:

**Function Name:** `ForMaskedLMLoss`

**Purpose:** Computes the masked cross-entropy loss.

**Key Details:**

* Takes in `logits`, `labels`, and other hyperparameters (`vocab_size`, `num_items_in_batch`, `ignore_index`, `**kwargs`) as inputs.
* Computes the loss on a batch of tokens using `fixed_cross_entropy` algorithm, adapting to model parallelism if needed.
* Returns the computed masked cross-entropy loss.

**Notable Operations:**

* Type casting and device management for compatible tensor operations.
* Flatten token sequences to individual elements for computation.

**Dependencies/Requirements:**

* External `kwargs` support using ** variable arguments (`**`) for flexibility with other library functions.


key: chroma:document
string_value:
CODE:

def ForSequenceClassificationLoss(labels, pooled_logits, config, **kwargs):
    num_labels = config.num_labels
    if config.problem_type is None:
        if num_labels == 1:
            config.problem_type = "regression"
        elif num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
            config.problem_type = "single_label_classification"
        else:
            config.problem_type = "multi_label_classification"

    labels = labels.to(pooled_logits.device)
    if config.problem_type == "regression":
        loss_fct = MSELoss()
        if num_labels == 1:
            loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
        else:
            loss = loss_fct(pooled_logits, labels)
    elif config.problem_type == "single_label_classification":
        loss = fixed_cross_entropy(pooled_logits.view(-1, num_labels), labels.view(-1), **kwargs)
    elif config.problem_type == "multi_label_classification":
        loss_fct = BCEWithLogitsLoss()
        loss = loss_fct(pooled_logits, labels)


SUMMARY:
**Function Summary: ForSequenceClassificationLoss**

* **Purpose:** Computes the sequence classification loss for a given set of labels and pooled logits.
* **Functionality:** Takes input `labels`, `pooled_logits`, and `config` to determine the problem type, then computes the corresponding loss using a chosen loss function (e.g., MSELoss or BCEWithLogitsLoss).
* **Key Implementation Details:**
	+ Uses problem type from `config` to select loss function,
	+ Utilizes `MSELoss` for regression problems and `BCEWithLogitsLoss` for multi-label classification problem types.
	+ Applies fixed cross-entropy for single-label classification problems.

**Method Signatures and Interfaces:**

* `loss_fct(pooled_logits, labels)`: loss function interface (MSELoss or BCEWithLogitsLoss)
* `fixed_cross_entropy`: custom cross-entropy calculation function

**Notable Dependencies and Requirements:**

* `torch`: deep learning framework
* `config` object with `num_labels`, `problem_type`, and other relevant attributes
* `labels.dtype == torch.long or labels.dtype == torch.int`: checks for label data type
* `pooled_logits.device`: device where logits are stored


key: chroma:document
string_value:
CODE:

def ForQuestionAnsweringLoss(start_logits, end_logits, start_positions, end_positions, **kwargs):
    total_loss = None
    if start_positions is not None and end_positions is not None:
        # If we are on multi-GPU, split add a dimension
        if len(start_positions.size()) > 1:
            start_positions = start_positions.squeeze(-1).to(start_logits.device)
        if len(end_positions.size()) > 1:
            end_positions = end_positions.squeeze(-1).to(end_logits.device)
        # sometimes the start/end positions are outside our model inputs, we ignore these terms
        ignored_index = start_logits.size(1)
        start_positions = start_positions.clamp(0, ignored_index)
        end_positions = end_positions.clamp(0, ignored_index)

        start_loss = fixed_cross_entropy(start_logits, start_positions, ignore_index=ignored_index, **kwargs)
        end_loss = fixed_cross_entropy(end_logits, end_positions, ignore_index=ignored_index, **kwargs)
        total_loss = (start_loss + end_loss) / 2


SUMMARY:
**Summary:**

`ForQuestionAnsweringLoss`: Calculates the total loss for question answer matching using fixed cross-entropy. Takes in input logits, positions, and additional keyword arguments.

**Key Functionality:**

* Calculates two cross-entropy losses (start_loss and end_loss) using `fixed_cross_entropy`.
* Returns the average of the two losses (`total_loss`).

**Implementation Details:**

* Initializes losses to `None`, then sets `started` if both `start_positions` and `end_positions` are not `None`.
* Handles multi-GPU environments by adding a dimension for batch size.
* Clamps positions to ignore out-of-bounds values.

**Noteable Dependencies:**

* `fixed_cross_entropy`: Cross-entropy loss function with an ignored index parameter.
* `torch` (device and tensor operations).

**Return Value:**

The calculated total loss (`total_loss`) or `None` if not applicable.


key: chroma:document
string_value:
CODE:

def ForTokenClassification(logits, labels, config, **kwargs):
    # Upcast to float if we need to compute the loss to avoid potential precision issues
    logits = logits.view(-1, config.num_labels)
    labels = labels.view(-1).to(logits.device)
    logits = logits.float()
    # Flatten the tokens


SUMMARY:
**Code Summary**

`ForTokenClassification`: A custom function for performing token classification, utilizing the `fixed_cross_entropy` algorithm.

**Main Purpose and Functionality**
```python
def ForTokenClassification(logits, labels, config, **kwargs):
    # ...
    return fixed_cross_entropy(logits, labels, **kwargs)
```
The primary purpose of this function is to compute the cross-entropy loss between predicted logits and true labels for token classification tasks.

**Key Implementation Details**

*   Upcasts `logits` to float type to prevent precision issues during computation.
*   Flattens tokens using `view` method with shape (batch_size, num_labels).
*   Moves `labels` to the device of `logits`.
*   Calls `fixed_cross_entropy` algorithm, passing flattened logits and labels.

**Notable Dependencies or Requirements**

*   `config`: Object containing configuration settings for the classification task.
*   `kwargs`: Additional keyword arguments passed from calling function (if applicable).
*   `fixed_cross_entropy`: Algorithm implemented to calculate cross-entropy loss.


key: chroma:document
string_value:
CODE:

def supported_features_mapping(
    *supported_features: str, onnx_config_cls: str = None
) -> Dict[str, Callable[[PretrainedConfig], OnnxConfig]]:
    """
    Generate the mapping between supported the features and their corresponding OnnxConfig for a given model.

    Args:
        *supported_features: The names of the supported features.
        onnx_config_cls: The OnnxConfig full name corresponding to the model.

    Returns:
        The dictionary mapping a feature to an OnnxConfig constructor.
    """
    if onnx_config_cls is None:
        raise ValueError("A OnnxConfig class must be provided")

    config_cls = transformers
    for attr_name in onnx_config_cls.split("."):
        config_cls = getattr(config_cls, attr_name)
    mapping = {}
    for feature in supported_features:
        if "-with-past" in feature:
            task = feature.replace("-with-past", "")
            mapping[feature] = partial(config_cls.with_past, task=task)
        else:
            mapping[feature] = partial(config_cls.from_model_config, task=feature)



SUMMARY:
**Summary**

`supported_features_mapping` function generates a dictionary mapping supported features to their corresponding OnnxConfig constructors for a given model.

**Main Purpose and Functionality**

* Takes in a variable number of `supported_features` strings, mapped to an optional `onnx_config_cls` string
* Injects the `onnx_config_cls` into a transformer configuration class using a dot notation resolution process
* Iterates through each supported feature:
	+ If the feature ends with "-with-past", it wraps the OnnxConfig constructor in `.with_past` while preserving the original task name.
	+ Otherwise, it uses the OnnxConfig constructor from the injected configuration class with the original feature as the task name.

**Key Implementation Details**

* Uses partial application to create wrapper functions for each supported feature
* Utilizes the `getattr` function to dynamically resolve the transformer configuration class

**Notable Dependencies or Requirements**

* The `transformers` module and its related classes (`OnnxConfig`, `from_model_config`, and `.with_past`)
* An instance of a `PretrainedConfig` object (required as an argument for partial function application)
* An external dependency is not explicitly listed; however, the context suggests it depends on the Hugging Face Transformers library.


key: chroma:document
string_value:
CODE:

class FeaturesManager:
    _TASKS_TO_AUTOMODELS = {}
    _TASKS_TO_TF_AUTOMODELS = {}
    if is_torch_available():
        _TASKS_TO_AUTOMODELS = {
            "default": AutoModel,
            "masked-lm": AutoModelForMaskedLM,
            "causal-lm": AutoModelForCausalLM,
            "seq2seq-lm": AutoModelForSeq2SeqLM,
            "sequence-classification": AutoModelForSequenceClassification,
            "token-classification": AutoModelForTokenClassification,
            "multiple-choice": AutoModelForMultipleChoice,
            "object-detection": AutoModelForObjectDetection,
            "question-answering": AutoModelForQuestionAnswering,
            "image-classification": AutoModelForImageClassification,
            "image-segmentation": AutoModelForImageSegmentation,
            "masked-im": AutoModelForMaskedImageModeling,
            "semantic-segmentation": AutoModelForSemanticSegmentation,
            "vision2seq-lm": AutoModelForVision2Seq,
            "speech2seq-lm": AutoModelForSpeechSeq2Seq,
        }
    if is_tf_available():
        _TASKS_TO_TF_AUTOMODELS = {
            "default": TFAutoModel,
            "masked-lm": TFAutoModelForMaskedLM,
            "causal-lm": TFAutoModelForCausalLM,
            "seq2seq-lm": TFAutoModelForSeq2SeqLM,
            "sequence-classification": TFAutoModelForSequenceClassification,
            "token-classification": TFAutoModelForTokenClassification,
            "multiple-choice": TFAutoModelForMultipleChoice,
            "question-answering": TFAutoModelForQuestionAnswering,
            "semantic-segmentation": TFAutoModelForSemanticSegmentation,
        }

    # Set of model topologies we support associated to the features supported by each topology and the factory
    _SUPPORTED_MODEL_TYPE = {
        "albert": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.albert.AlbertOnnxConfig",
        ),
        "bart": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            "sequence-classification",
            "question-answering",
            onnx_config_cls="models.bart.BartOnnxConfig",
        ),
        # BEiT cannot be used with the masked image modeling autoclass, so this feature is excluded here
        "beit": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.beit.BeitOnnxConfig"
        ),
        "bert": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.bert.BertOnnxConfig",
        ),
        "big-bird": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.big_bird.BigBirdOnnxConfig",
        ),
        "bigbird-pegasus": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            "sequence-classification",
            "question-answering",
            onnx_config_cls="models.bigbird_pegasus.BigBirdPegasusOnnxConfig",
        ),
        "blenderbot": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            onnx_config_cls="models.blenderbot.BlenderbotOnnxConfig",
        ),
        "blenderbot-small": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            onnx_config_cls="models.blenderbot_small.BlenderbotSmallOnnxConfig",
        ),
        "bloom": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "sequence-classification",
            "token-classification",
            onnx_config_cls="models.bloom.BloomOnnxConfig",
        ),
        "camembert": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.camembert.CamembertOnnxConfig",
        ),
        "clip": supported_features_mapping(
            "default",
            onnx_config_cls="models.clip.CLIPOnnxConfig",
        ),
        "codegen": supported_features_mapping(
            "default",
            "causal-lm",
            onnx_config_cls="models.codegen.CodeGenOnnxConfig",
        ),
        "convbert": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.convbert.ConvBertOnnxConfig",
        ),
        "convnext": supported_features_mapping(
            "default",
            "image-classification",
            onnx_config_cls="models.convnext.ConvNextOnnxConfig",
        ),
        "data2vec-text": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.data2vec.Data2VecTextOnnxConfig",
        ),
        "data2vec-vision": supported_features_mapping(
            "default",
            "image-classification",
            # ONNX doesn't support `adaptive_avg_pool2d` yet
            # "semantic-segmentation",
            onnx_config_cls="models.data2vec.Data2VecVisionOnnxConfig",
        ),
        "deberta": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.deberta.DebertaOnnxConfig",
        ),
        "deberta-v2": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.deberta_v2.DebertaV2OnnxConfig",
        ),
        "deit": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.deit.DeiTOnnxConfig"
        ),
        "detr": supported_features_mapping(
            "default",
            "object-detection",
            "image-segmentation",
            onnx_config_cls="models.detr.DetrOnnxConfig",
        ),
        "distilbert": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.distilbert.DistilBertOnnxConfig",
        ),
        "electra": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.electra.ElectraOnnxConfig",
        ),
        "flaubert": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.flaubert.FlaubertOnnxConfig",
        ),
        "gpt2": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "sequence-classification",
            "token-classification",
            onnx_config_cls="models.gpt2.GPT2OnnxConfig",
        ),
        "gptj": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "question-answering",
            "sequence-classification",
            onnx_config_cls="models.gptj.GPTJOnnxConfig",
        ),
        "gpt-neo": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "sequence-classification",
            onnx_config_cls="models.gpt_neo.GPTNeoOnnxConfig",
        ),
        "groupvit": supported_features_mapping(
            "default",
            onnx_config_cls="models.groupvit.GroupViTOnnxConfig",
        ),
        "ibert": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.ibert.IBertOnnxConfig",
        ),
        "imagegpt": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.imagegpt.ImageGPTOnnxConfig"
        ),
        "layoutlm": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "token-classification",
            onnx_config_cls="models.layoutlm.LayoutLMOnnxConfig",
        ),
        "layoutlmv3": supported_features_mapping(
            "default",
            "question-answering",
            "sequence-classification",
            "token-classification",
            onnx_config_cls="models.layoutlmv3.LayoutLMv3OnnxConfig",
        ),
        "levit": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.levit.LevitOnnxConfig"
        ),
        "longt5": supported_features_mapping(
            "default",
            "default-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            onnx_config_cls="models.longt5.LongT5OnnxConfig",
        ),
        "longformer": supported_features_mapping(
            "default",
            "masked-lm",
            "multiple-choice",
            "question-answering",
            "sequence-classification",
            "token-classification",
            onnx_config_cls="models.longformer.LongformerOnnxConfig",
        ),
        "marian": supported_features_mapping(
            "default",
            "default-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            "causal-lm",
            "causal-lm-with-past",
            onnx_config_cls="models.marian.MarianOnnxConfig",
        ),
        "mbart": supported_features_mapping(
            "default",
            "default-with-past",
            "causal-lm",
            "causal-lm-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            "sequence-classification",
            "question-answering",
            onnx_config_cls="models.mbart.MBartOnnxConfig",
        ),
        "mobilebert": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.mobilebert.MobileBertOnnxConfig",
        ),
        "mobilenet-v1": supported_features_mapping(
            "default",
            "image-classification",
            onnx_config_cls="models.mobilenet_v1.MobileNetV1OnnxConfig",
        ),
        "mobilenet-v2": supported_features_mapping(
            "default",
            "image-classification",
            onnx_config_cls="models.mobilenet_v2.MobileNetV2OnnxConfig",
        ),
        "mobilevit": supported_features_mapping(
            "default",
            "image-classification",
            onnx_config_cls="models.mobilevit.MobileViTOnnxConfig",
        ),
        "mt5": supported_features_mapping(
            "default",
            "default-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            onnx_config_cls="models.mt5.MT5OnnxConfig",
        ),
        "m2m-100": supported_features_mapping(
            "default",
            "default-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            onnx_config_cls="models.m2m_100.M2M100OnnxConfig",
        ),
        "owlvit": supported_features_mapping(
            "default",
            onnx_config_cls="models.owlvit.OwlViTOnnxConfig",
        ),
        "perceiver": supported_features_mapping(
            "image-classification",
            "masked-lm",
            "sequence-classification",
            onnx_config_cls="models.perceiver.PerceiverOnnxConfig",
        ),
        "poolformer": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.poolformer.PoolFormerOnnxConfig"
        ),
        "rembert": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.rembert.RemBertOnnxConfig",
        ),
        "resnet": supported_features_mapping(
            "default",
            "image-classification",
            onnx_config_cls="models.resnet.ResNetOnnxConfig",
        ),
        "roberta": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.roberta.RobertaOnnxConfig",
        ),
        "roformer": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "token-classification",
            "multiple-choice",
            "question-answering",
            "token-classification",
            onnx_config_cls="models.roformer.RoFormerOnnxConfig",
        ),
        "segformer": supported_features_mapping(
            "default",
            "image-classification",
            "semantic-segmentation",
            onnx_config_cls="models.segformer.SegformerOnnxConfig",
        ),
        "squeezebert": supported_features_mapping(
            "default",
            "masked-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.squeezebert.SqueezeBertOnnxConfig",
        ),
        "swin": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.swin.SwinOnnxConfig"
        ),
        "t5": supported_features_mapping(
            "default",
            "default-with-past",
            "seq2seq-lm",
            "seq2seq-lm-with-past",
            onnx_config_cls="models.t5.T5OnnxConfig",
        ),
        "vision-encoder-decoder": supported_features_mapping(
            "vision2seq-lm", onnx_config_cls="models.vision_encoder_decoder.VisionEncoderDecoderOnnxConfig"
        ),
        "vit": supported_features_mapping(
            "default", "image-classification", onnx_config_cls="models.vit.ViTOnnxConfig"
        ),
        "whisper": supported_features_mapping(
            "default",
            "default-with-past",
            "speech2seq-lm",
            "speech2seq-lm-with-past",
            onnx_config_cls="models.whisper.WhisperOnnxConfig",
        ),
        "xlm": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.xlm.XLMOnnxConfig",
        ),
        "xlm-roberta": supported_features_mapping(
            "default",
            "masked-lm",
            "causal-lm",
            "sequence-classification",
            "multiple-choice",
            "token-classification",
            "question-answering",
            onnx_config_cls="models.xlm_roberta.XLMRobertaOnnxConfig",
        ),
        "yolos": supported_features_mapping(
            "default",
            "object-detection",
            onnx_config_cls="models.yolos.YolosOnnxConfig",
        ),
    }

    AVAILABLE_FEATURES = sorted(reduce(lambda s1, s2: s1 | s2, (v.keys() for v in _SUPPORTED_MODEL_TYPE.values())))

    @staticmethod
    def get_supported_features_for_model_type(
        model_type: str, model_name: Optional[str] = None
    ) -> Dict[str, Callable[[PretrainedConfig], OnnxConfig]]:
        """
        Tries to retrieve the feature -> OnnxConfig constructor map from the model type.

        Args:
            model_type (`str`):
                The model type to retrieve the supported features for.
            model_name (`str`, *optional*):
                The name attribute of the model object, only used for the exception message.

        Returns:
            The dictionary mapping each feature to a corresponding OnnxConfig constructor.
        """
        model_type = model_type.lower()
        if model_type not in FeaturesManager._SUPPORTED_MODEL_TYPE:
            model_type_and_model_name = f"{model_type} ({model_name})" if model_name else model_type
            raise KeyError(
                f"{model_type_and_model_name} is not supported yet. "
                f"Only {list(FeaturesManager._SUPPORTED_MODEL_TYPE.keys())} are supported. "
                f"If you want to support {model_type} please propose a PR or open up an issue."
            )
        return FeaturesManager._SUPPORTED_MODEL_TYPE[model_type]

    @staticmethod
    def feature_to_task(feature: str) -> str:
        return feature.replace("-with-past", "")

    @staticmethod
    def _validate_framework_choice(framework: str):
        """
        Validates if the framework requested for the export is both correct and available, otherwise throws an
        exception.
        """
        if framework not in ["pt", "tf"]:
            raise ValueError(
                f"Only two frameworks are supported for ONNX export: pt or tf, but {framework} was provided."
            )
        elif framework == "pt" and not is_torch_available():
            raise RuntimeError("Cannot export model to ONNX using PyTorch because no PyTorch package was found.")
        elif framework == "tf" and not is_tf_available():
            raise RuntimeError("Cannot export model to ONNX using TensorFlow because no TensorFlow package was found.")

    @staticmethod
    def get_model_class_for_feature(feature: str, framework: str = "pt") -> Type:
        """
        Attempts to retrieve an AutoModel class from a feature name.

        Args:
            feature (`str`):
                The feature required.
            framework (`str`, *optional*, defaults to `"pt"`):
                The framework to use for the export.

        Returns:
            The AutoModel class corresponding to the feature.
        """
        task = FeaturesManager.feature_to_task(feature)
        FeaturesManager._validate_framework_choice(framework)
        if framework == "pt":
            task_to_automodel = FeaturesManager._TASKS_TO_AUTOMODELS
        else:
            task_to_automodel = FeaturesManager._TASKS_TO_TF_AUTOMODELS
        if task not in task_to_automodel:
            raise KeyError(
                f"Unknown task: {feature}. Possible values are {list(FeaturesManager._TASKS_TO_AUTOMODELS.values())}"
            )

        return task_to_automodel[task]

    @staticmethod
    def determine_framework(model: str, framework: str = None) -> str:
        """
        Determines the framework to use for the export.

        The priority is in the following order:
            1. User input via `framework`.
            2. If local checkpoint is provided, use the same framework as the checkpoint.
            3. Available framework in environment, with priority given to PyTorch

        Args:
            model (`str`):
                The name of the model to export.
            framework (`str`, *optional*, defaults to `None`):
                The framework to use for the export. See above for priority if none provided.

        Returns:
            The framework to use for the export.

        """
        if framework is not None:
            return framework

        framework_map = {"pt": "PyTorch", "tf": "TensorFlow"}
        exporter_map = {"pt": "torch", "tf": "tf2onnx"}

        if os.path.isdir(model):
            if os.path.isfile(os.path.join(model, WEIGHTS_NAME)):
                framework = "pt"
            elif os.path.isfile(os.path.join(model, TF2_WEIGHTS_NAME)):
                framework = "tf"
            else:
                raise FileNotFoundError(
                    "Cannot determine framework from given checkpoint location."
                    f" There should be a {WEIGHTS_NAME} for PyTorch"
                    f" or {TF2_WEIGHTS_NAME} for TensorFlow."
                )
            logger.info(f"Local {framework_map[framework]} model found.")
        else:
            if is_torch_available():
                framework = "pt"
            elif is_tf_available():
                framework = "tf"
            else:
                raise EnvironmentError("Neither PyTorch nor TensorFlow found in environment. Cannot export to ONNX.")

        logger.info(f"Framework not requested. Using {exporter_map[framework]} to export to ONNX.")

        return framework

    @staticmethod
    def get_model_from_feature(
        feature: str, model: str, framework: str = None, cache_dir: str = None
    ) -> Union["PreTrainedModel", "TFPreTrainedModel"]:
        """
        Attempts to retrieve a model from a model's name and the feature to be enabled.

        Args:
            feature (`str`):
                The feature required.
            model (`str`):
                The name of the model to export.
            framework (`str`, *optional*, defaults to `None`):
                The framework to use for the export. See `FeaturesManager.determine_framework` for the priority should
                none be provided.

        Returns:
            The instance of the model.

        """
        framework = FeaturesManager.determine_framework(model, framework)
        model_class = FeaturesManager.get_model_class_for_feature(feature, framework)
        try:
            model = model_class.from_pretrained(model, cache_dir=cache_dir)
        except OSError:
            if framework == "pt":
                logger.info("Loading TensorFlow model in PyTorch before exporting to ONNX.")
                model = model_class.from_pretrained(model, from_tf=True, cache_dir=cache_dir)
            else:
                logger.info("Loading PyTorch model in TensorFlow before exporting to ONNX.")
                model = model_class.from_pretrained(model, from_pt=True, cache_dir=cache_dir)
        return model

    @staticmethod
    def check_supported_model_or_raise(
        model: Union["PreTrainedModel", "TFPreTrainedModel"], feature: str = "default"
    ) -> Tuple[str, Callable]:
        """
        Check whether or not the model has the requested features.

        Args:
            model: The model to export.
            feature: The name of the feature to check if it is available.

        Returns:
            (str) The type of the model (OnnxConfig) The OnnxConfig instance holding the model export properties.

        """
        model_type = model.config.model_type.replace("_", "-")
        model_name = getattr(model, "name", "")
        model_features = FeaturesManager.get_supported_features_for_model_type(model_type, model_name=model_name)
        if feature not in model_features:
            raise ValueError(
                f"{model.config.model_type} doesn't support feature {feature}. Supported values are: {model_features}"
            )

        return model.config.model_type, FeaturesManager._SUPPORTED_MODEL_TYPE[model_type][feature]

    def get_config(model_type: str, feature: str) -> OnnxConfig:
        """
        Gets the OnnxConfig for a model_type and feature combination.

        Args:
            model_type (`str`):
                The model type to retrieve the config for.
            feature (`str`):
                The feature to retrieve the config for.

        Returns:
            `OnnxConfig`: config for the combination
        """


SUMMARY:
This code defines a module named `FeaturesManager` that appears to be used for managing export configurations of machine learning models, specifically for use with the Hugging Face library. The module provides various functions and classes to determine the framework used for the export, determine the model type based on its name, get the OnnxConfig instance for a specific request, and more.

The main components of this module are:

- `determine_framework`: Determines the framework (either PyTorch or TensorFlow) to use for exporting a given model.
- `get_model_from_feature` : Attempts to load a pre-trained model from a specified name and feature combination. If required, it attempts loading the model's local checkpoint in the same framework before exporting it.
- `check_supported_model_or_raise`: Checks whether a particular supported feature is enabled or available on a given model and raises an error if not.

This module utilizes other libraries and classes for determining the best approach to export models effectively:

- Import of Hugging Face libraries (`torch`, `tf2onnx`).
- Use of PyTorch (`torch`) and TensorFlow (`tf2onnx`) exporters.
- Class names like `PreTrainedModel` and `TFPreTrainedModel`.
- Utilization of environment variables (`WEIGHTS_NAME`, `TF2_WEIGHTS_NAME`).

Given the structure of this module, it ensures flexible configuration management for machine learning models across different frameworks in order to streamline model export tasks. It takes a structured approach that depends on local file existence as an alternative if the user requested model name and framework but still doesn't have any other input (as seen in `determine_framework`). This simplifies the task of finding suitable export configurations based on external data provided by Hugging Face or the environment itself.

If improvements are needed to this feature, then possibly some potential suggestions could be that: 
- Adding exception handling for certain types of exceptions raised during model loading or exporting.
- Including methods to handle a large number of supported models (one for each model type) and maintain their respective configs efficiently.
- Optimizing the use of caching for reducing unnecessary computation when multiple users are working in the same environment.


key: chroma:document
string_value:
CODE:

def check_onnxruntime_requirements(minimum_version: Version):
    """
    Check onnxruntime is installed and if the installed version match is recent enough

    Raises:
        ImportError: If onnxruntime is not installed or too old version is found
    """
    try:
        import onnxruntime

        # Parse the version of the installed onnxruntime
        ort_version = parse(onnxruntime.__version__)

        # We require 1.4.0 minimum
        if ort_version < ORT_QUANTIZE_MINIMUM_VERSION:
            raise ImportError(
                f"We found an older version of onnxruntime ({onnxruntime.__version__}) "
                f"but we require onnxruntime to be >= {minimum_version} to enable all the conversions options.\n"
                "Please update onnxruntime by running `pip install --upgrade onnxruntime`"
            )

    except ImportError:
        raise ImportError(
            "onnxruntime doesn't seem to be currently installed. "
            "Please install the onnxruntime by running `pip install onnxruntime`"
            " and relaunch the conversion."


SUMMARY:
**Summary**

This function `check_onnxruntime_requirements` checks if the OnnxRuntime library is installed and its version is recent enough to enable advanced conversions. It:

- **Import check**: Imports onnxruntime and attempts to parse its version using the `parse` function (not defined in this snippet).
- **Version comparison**: Compares the parsed version with the minimum required version `ORT_QUANTIZE_MINIMUM_VERSION`.
- **Error handling**:
  - Raises an `ImportError` if OnnxRuntime is not installed.
  - If installed, but with too old a version, raises an `ImportError` urging update of OnnxRuntime via pip.

**Key implementation details:**

* Uses `onnxruntime.__version__`
* Requires minimal version `ORT_QUANTIZE_MINIMUM_VERSION`

**Notable dependencies:**

* External dependency: `onnxruntime` (imported)
* Inherited/implemented interface: None

Note: The `parse` function is not defined in this snippet, implying it's possibly another utility or library not immediately relevant to the described functionality.


key: chroma:document
string_value:
CODE:

def export_pytorch(
    preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin", "ProcessorMixin"],
    model: "PreTrainedModel",
    config: OnnxConfig,
    opset: int,
    output: Path,
    tokenizer: "PreTrainedTokenizer" = None,
    device: str = "cpu",
) -> Tuple[List[str], List[str]]:
    """
    Export a PyTorch model to an ONNX Intermediate Representation (IR)

    Args:
        preprocessor: ([`PreTrainedTokenizer`], [`FeatureExtractionMixin`] or [`ProcessorMixin`]):
            The preprocessor used for encoding the data.
        model ([`PreTrainedModel`]):
            The model to export.
        config ([`~onnx.config.OnnxConfig`]):
            The ONNX configuration associated with the exported model.
        opset (`int`):
            The version of the ONNX operator set to use.
        output (`Path`):
            Directory to store the exported ONNX model.
        device (`str`, *optional*, defaults to `cpu`):
            The device on which the ONNX model will be exported. Either `cpu` or `cuda`.

    Returns:
        `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
        the ONNX configuration.
    """

    if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:
        raise ValueError("You cannot provide both a tokenizer and a preprocessor to export the model.")
    if tokenizer is not None:
        warnings.warn(
            "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"
            " `preprocessor` instead.",
            FutureWarning,
        )
        logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummmy inputs.")
        preprocessor = tokenizer

    if issubclass(type(model), PreTrainedModel):
        import torch
        from torch.onnx import export as onnx_export

        logger.info(f"Using framework PyTorch: {torch.__version__}")
        with torch.no_grad():
            model.config.return_dict = True
            model.eval()

            # Check if we need to override certain configuration item
            if config.values_override is not None:
                logger.info(f"Overriding {len(config.values_override)} configuration item(s)")
                for override_config_key, override_config_value in config.values_override.items():
                    logger.info(f"\t- {override_config_key} -> {override_config_value}")
                    setattr(model.config, override_config_key, override_config_value)

            # Ensure inputs match
            # TODO: Check when exporting QA we provide "is_pair=True"
            model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.PYTORCH)
            device = torch.device(device)
            if device.type == "cuda" and torch.cuda.is_available():
                model.to(device)
                model_inputs_device = {}
                for k, v in model_inputs.items():
                    if isinstance(v, Tuple):
                        model_inputs_device[k] = tuple(
                            x.to(device) if isinstance(x, torch.Tensor) else None for x in v
                        )
                    elif isinstance(v, List):
                        model_inputs_device[k] = [
                            tuple(x.to(device) if isinstance(x, torch.Tensor) else None for x in t) for t in v
                        ]
                    else:
                        model_inputs_device[k] = v.to(device)

                model_inputs = model_inputs_device

            inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys())
            onnx_outputs = list(config.outputs.keys())

            if not inputs_match:
                raise ValueError("Model and config inputs doesn't match")

            config.patch_ops()

            onnx_export(
                model,
                (model_inputs,),
                f=output.as_posix(),
                input_names=list(config.inputs.keys()),
                output_names=onnx_outputs,
                dynamic_axes=dict(chain(config.inputs.items(), config.outputs.items())),
                do_constant_folding=True,
                opset_version=opset,
            )

            config.restore_ops()



SUMMARY:
**PyTorch Model Export Function Summary**

### Purpose and Functionality

The `export_pytorch` function exports a PyTorch model to an ONNX Intermediate Representation (IR) file. It takes in various parameters such as the preprocessor, model, configuration, opset version, output directory, device, and optional tokenizer and device.

**Key Implementation Details**

1.  **Preprocessing**: The function preprocesses the input data using the provided `preprocessor` or `tokenizer`, and checks if both are provided.
2.  **Model Evaluation**: It evaluates the model in evaluation mode and sets the configuration to return a dictionary (`model.config.return_dict = True`).
3.  **Configuration Patching**: The function patches ONNX operations based on the provided configurations (`config.patch_ops()`).
4.  **ONNX Export**: It exports the model using PyTorch's `onnx_export` function, specifying the input and output names, opset version, and other settings.
5.  **Input Validation**: It checks if the model inputs match the configuration inputs and raises a `ValueError` if they do not.

**Noteable Dependencies and Requirements**

1.  The function uses PyTorch as its underlying framework (`torch`).
2.  It relies on the ONNX library (`torch.onnx`) for exporting to ONNX.
3.  External dependencies include:
    *   `Path` from the ` pathlib` module (not explicitly used in this excerpt but implied)
    *   `Logger`, `warnings`, and `FutureWarning` from the PyTorch documentation
    *   `TensorType` (enum) from the ONNX documentation

### Return Values

Upon successful export, the function returns a tuple containing:

1.  An ordered list of the model's inputs (`matched_inputs`)
2.  A list of named outputs from the ONNX configuration (`onnx_outputs`)


key: chroma:document
string_value:
CODE:

def export_tensorflow(
    preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin"],
    model: "TFPreTrainedModel",
    config: OnnxConfig,
    opset: int,
    output: Path,
    tokenizer: "PreTrainedTokenizer" = None,
) -> Tuple[List[str], List[str]]:
    """
    Export a TensorFlow model to an ONNX Intermediate Representation (IR)

    Args:
        preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]):
            The preprocessor used for encoding the data.
        model ([`TFPreTrainedModel`]):
            The model to export.
        config ([`~onnx.config.OnnxConfig`]):
            The ONNX configuration associated with the exported model.
        opset (`int`):
            The version of the ONNX operator set to use.
        output (`Path`):
            Directory to store the exported ONNX model.

    Returns:
        `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
        the ONNX configuration.
    """
    import onnx
    import tensorflow as tf
    import tf2onnx

    if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:
        raise ValueError("You cannot provide both a tokenizer and preprocessor to export the model.")
    if tokenizer is not None:
        warnings.warn(
            "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"
            " `preprocessor` instead.",
            FutureWarning,
        )
        logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummmy inputs.")
        preprocessor = tokenizer

    model.config.return_dict = True

    # Check if we need to override certain configuration item
    if config.values_override is not None:
        logger.info(f"Overriding {len(config.values_override)} configuration item(s)")
        for override_config_key, override_config_value in config.values_override.items():
            logger.info(f"\t- {override_config_key} -> {override_config_value}")
            setattr(model.config, override_config_key, override_config_value)

    # Ensure inputs match
    model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.TENSORFLOW)
    inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys())
    onnx_outputs = list(config.outputs.keys())

    input_signature = [
        tf.TensorSpec([None] * tensor.ndim, dtype=tensor.dtype, name=key) for key, tensor in model_inputs.items()
    ]
    onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=opset)
    onnx.save(onnx_model, output.as_posix())
    config.restore_ops()



SUMMARY:
**Summary:**

`export_tensorflow` function exports a TensorFlow model to an ONNX IR.

**Main Purpose and Functionality:**

* Takes in a model, preprocessor, configuration, and output directory
* Converts the model to ONNX format with specified opset version
* Generates dummy inputs for the model if necessary
* Returns ordered lists of model inputs and named inputs from configuration

**Key Implementation Details:**

* Uses `tf2onnx` library to convert Keras model to ONNX format
* Modifies model configuration with override values if provided
* Ensures input signatures match between model and configuration

**Notable Dependencies and Requirements:**

* `onnx` and `tensorflow` libraries
* `tf2onnx` library for model conversion
* `PreTrainedTokenizer`, `TFPreTrainedModel`, and `OnnxConfig` classes from Hugging Face Transformers library


key: chroma:document
string_value:
CODE:

def export(
    preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin", "ProcessorMixin"],
    model: Union["PreTrainedModel", "TFPreTrainedModel"],
    config: OnnxConfig,
    opset: int,
    output: Path,
    tokenizer: "PreTrainedTokenizer" = None,
    device: str = "cpu",
) -> Tuple[List[str], List[str]]:
    """
    Export a Pytorch or TensorFlow model to an ONNX Intermediate Representation (IR)

    Args:
        preprocessor: ([`PreTrainedTokenizer`], [`FeatureExtractionMixin`] or [`ProcessorMixin`]):
            The preprocessor used for encoding the data.
        model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
            The model to export.
        config ([`~onnx.config.OnnxConfig`]):
            The ONNX configuration associated with the exported model.
        opset (`int`):
            The version of the ONNX operator set to use.
        output (`Path`):
            Directory to store the exported ONNX model.
        device (`str`, *optional*, defaults to `cpu`):
            The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
            export on CUDA devices.

    Returns:
        `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
        the ONNX configuration.
    """
    if not (is_torch_available() or is_tf_available()):
        raise ImportError(
            "Cannot convert because neither PyTorch nor TensorFlow are not installed. "
            "Please install torch or tensorflow first."
        )

    if is_tf_available() and isinstance(model, TFPreTrainedModel) and device == "cuda":
        raise RuntimeError("`tf2onnx` does not support export on CUDA device.")

    if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:
        raise ValueError("You cannot provide both a tokenizer and a preprocessor to export the model.")
    if tokenizer is not None:
        warnings.warn(
            "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"
            " `preprocessor` instead.",
            FutureWarning,
        )
        logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummmy inputs.")
        preprocessor = tokenizer

    if is_torch_available():
        from ..utils import get_torch_version

        if not config.is_torch_support_available:
            logger.warning(
                f"Unsupported PyTorch version for this model. Minimum required is {config.torch_onnx_minimum_version},"
                f" got: {get_torch_version()}"
            )

    if is_torch_available() and issubclass(type(model), PreTrainedModel):
        return export_pytorch(preprocessor, model, config, opset, output, tokenizer=tokenizer, device=device)
    elif is_tf_available() and issubclass(type(model), TFPreTrainedModel):


SUMMARY:
**Summary:**

This function, `export`, exports a PyTorch or TensorFlow model to an ONNX Intermediate Representation (IR) with specified parameters.

**Main Purpose and Functionality:**
The primary purpose of this function is to convert a trained PyTorch or TensorFlow model into an ONNX format, which can be used for inference on other platforms. The function takes various inputs, including the preprocessor, model, configuration, operator set version, output directory, device, and tokenizer (optional).

**Key Implementation Details:**

* It checks availability of required libraries (PyTorch and TensorFlow) and raises errors if not installed.
* It handles deprecated `tokenizer` argument and logs a warning if it is provided.
* For PyTorch models, it uses the `export_pytorch` function to export the model. This function is not shown in the provided code chunk.
* For TensorFlow models, it uses the `export_tensorflow` function to export the model. This function is also not shown in the provided code chunk.

**Notable Dependencies or Requirements:**

* Requires PyTorch and/or TensorFlow libraries
* Uses `PreTrainedTokenizer`, `FeatureExtractionMixin`, `ProcessorMixin`, `PreTrainedModel`, `TFPreTrainedModel`, `OnnxConfig` classes
* Exports to ONNX format with specified parameters

**Method Signatures and Interfaces:**

* Not shown in the provided code chunk, but likely includes `export_pytorch` and `export_tensorflow` functions.


key: chroma:document
string_value:
CODE:

def validate_model_outputs(
    config: OnnxConfig,
    preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin", "ProcessorMixin"],
    reference_model: Union["PreTrainedModel", "TFPreTrainedModel"],
    onnx_model: Path,
    onnx_named_outputs: List[str],
    atol: float,
    tokenizer: "PreTrainedTokenizer" = None,
):
    from onnxruntime import InferenceSession, SessionOptions

    logger.info("Validating ONNX model...")

    if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:
        raise ValueError("You cannot provide both a tokenizer and a preprocessor to validate the model outputs.")
    if tokenizer is not None:
        warnings.warn(
            "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"
            " `preprocessor` instead.",
            FutureWarning,
        )
        logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummmy inputs.")
        preprocessor = tokenizer

    # generate inputs with a different batch_size and seq_len that was used for conversion to properly test
    # dynamic input shapes.
    if is_torch_available() and issubclass(type(reference_model), PreTrainedModel):
        reference_model_inputs = config.generate_dummy_inputs(
            preprocessor,
            batch_size=config.default_fixed_batch + 1,
            seq_length=config.default_fixed_sequence + 1,
            framework=TensorType.PYTORCH,
        )
    else:
        reference_model_inputs = config.generate_dummy_inputs(
            preprocessor,
            batch_size=config.default_fixed_batch + 1,
            seq_length=config.default_fixed_sequence + 1,
            framework=TensorType.TENSORFLOW,
        )

    # Create ONNX Runtime session
    options = SessionOptions()
    session = InferenceSession(onnx_model.as_posix(), options, providers=["CPUExecutionProvider"])

    # Compute outputs from the reference model
    if is_torch_available() and issubclass(type(reference_model), PreTrainedModel):
        reference_model.to("cpu")
    ref_outputs = reference_model(**reference_model_inputs)
    ref_outputs_dict = {}

    # We flatten potential collection of outputs (i.e. past_keys) to a flat structure
    for name, value in ref_outputs.items():
        # Overwriting the output name as "present" since it is the name used for the ONNX outputs
        # ("past_key_values" being taken for the ONNX inputs)
        if name == "past_key_values":
            name = "present"
        if isinstance(value, (list, tuple)):
            value = config.flatten_output_collection_property(name, value)
            ref_outputs_dict.update(value)
        else:
            ref_outputs_dict[name] = value

    # Create onnxruntime inputs from the reference model inputs
    reference_model_inputs_onnxruntime = config.generate_dummy_inputs_onnxruntime(reference_model_inputs)

    # We flatten potential collection of inputs (i.e. past_keys)
    onnx_inputs = {}
    for name, value in reference_model_inputs_onnxruntime.items():
        if isinstance(value, (list, tuple)):
            value = config.flatten_output_collection_property(name, value)
            onnx_inputs.update({tensor_name: pt_tensor.numpy() for tensor_name, pt_tensor in value.items()})
        else:
            onnx_inputs[name] = value.numpy()

    # Compute outputs from the ONNX model
    onnx_outputs = session.run(onnx_named_outputs, onnx_inputs)

    # Check we have a subset of the keys into onnx_outputs against ref_outputs
    ref_outputs_set, onnx_outputs_set = set(ref_outputs_dict.keys()), set(onnx_named_outputs)
    if not onnx_outputs_set.issubset(ref_outputs_set):
        logger.info(
            f"\t-[x] ONNX model output names {onnx_outputs_set} do not match reference model {ref_outputs_set}"
        )

        raise ValueError(
            "Outputs doesn't match between reference model and ONNX exported model: "
            f"{onnx_outputs_set.difference(ref_outputs_set)}"
        )
    else:
        logger.info(f"\t-[✓] ONNX model output names match reference model ({onnx_outputs_set})")

    # Check the shape and values match
    for name, ort_value in zip(onnx_named_outputs, onnx_outputs):
        if is_torch_available() and issubclass(type(reference_model), PreTrainedModel):
            ref_value = ref_outputs_dict[name].detach().numpy()
        else:
            ref_value = ref_outputs_dict[name].numpy()
        logger.info(f'\t- Validating ONNX Model output "{name}":')

        # Shape
        if not ort_value.shape == ref_value.shape:
            logger.info(f"\t\t-[x] shape {ort_value.shape} doesn't match {ref_value.shape}")
            raise ValueError(
                "Outputs shape doesn't match between reference model and ONNX exported model: "
                f"Got {ref_value.shape} (reference) and {ort_value.shape} (ONNX)"
            )
        else:
            logger.info(f"\t\t-[✓] {ort_value.shape} matches {ref_value.shape}")

        # Values
        if not np.allclose(ref_value, ort_value, atol=atol):
            bad_indices = np.logical_not(np.isclose(ref_value, ort_value, atol=atol))
            logger.info(f"\t\t-[x] values not close enough (atol: {atol})")
            raise ValueError(
                "Outputs values doesn't match between reference model and ONNX exported model: "
                f"Got max absolute difference of: {np.amax(np.abs(ref_value - ort_value))} for "
                f"{ref_value[bad_indices]} vs {ort_value[bad_indices]}"
            )
        else:


SUMMARY:
**Function Summary**

`validate_model_outputs`

Validates the outputs of an ONNX model against a reference TPU model. The function takes in a configuration, preprocessor, and onnx model as input.

**Primary Purpose**: To ensure that the output of the ONNX model matches the expected output of the reference TPU model.

**Key Implementation Details**

1. **Reference Model Input Generation**: Dynamic inputs are generated for the reference model to test dynamic shapes.
2. **ONNX Session Creation**: An INferenceSession is created with the onnx model and session options.
3. **Reference Model Output Extraction**: Outputs from the reference model are extracted using the `to` method.
4. **FLatten Collection Properties**: Potential collection of outputs (e.g., "past_key_values") is flattened to a flat structure.

**Notable Dependencies/Requirements**

1. `onnxruntime`: Used for creating an INferenceSession.
2. `PreTrainedTokenizer` and `PreTrainedModel`: Used as input preprocessor and reference model, respectively.

**Check Points**

1. Output names from the ONNX model must match the expected output names of the reference TPU model.
2. Output shapes and values from both models must match within a certain tolerance (`atol`).


key: chroma:document
string_value:
CODE:

def ensure_model_and_config_inputs_match(
    model: Union["PreTrainedModel", "TFPreTrainedModel"], model_inputs: Iterable[str]
) -> Tuple[bool, List[str]]:
    """

    :param model_inputs: :param config_inputs: :return:
    """
    if is_torch_available() and issubclass(type(model), PreTrainedModel):
        forward_parameters = signature(model.forward).parameters
    else:
        forward_parameters = signature(model.call).parameters
    model_inputs_set = set(model_inputs)

    # We are fine if config_inputs has more keys than model_inputs
    forward_inputs_set = set(forward_parameters.keys())
    is_ok = model_inputs_set.issubset(forward_inputs_set)

    # Make sure the input order match (VERY IMPORTANT !!!!)
    matching_inputs = forward_inputs_set.intersection(model_inputs_set)
    ordered_inputs = [parameter for parameter in forward_parameters.keys() if parameter in matching_inputs]


SUMMARY:
**Summary**

`ensure_model_and_config_inputs_match` function ensures that input parameters match between model and configuration inputs.

### Purpose and Functionality

Verifies whether input parameters of a pre-trained model (`model`) align with the specified model inputs (`model_inputs`). Returns a boolean indicating success and a list of matched input parameters in order.

### Key Implementation Details

* Checks if `model` is compatible with both TensorFlow and PyTorch models.
* Extracts forward function parameters using `signature(model.forward).parameters` (for PyTorch) or `signature(model.call).parameters` (for TensorFlow).
* Compares input sets between model inputs and forward function parameters.

### Notable Dependencies and Requirements

* `PyTorch` or `TensorFlow` models (`PreTrainedModel`, `TFPreTrainedModel`)
* `Union` type hint for model parameter type
* `Iterable[str]` for model inputs
* `Signature` object from `torch.nn.Module` (for PyTorch) or `tf.keras.Model` (for TensorFlow)
* `is_torch_available()` function (optional)

### Return Value

* `bool`: whether input parameters match
* `List[str]`: list of matched input parameters in order


key: chroma:document
string_value:
CODE:

class ParameterFormat(Enum):
    Float = c_float

    @property
    def size(self) -> int:
        """
        Number of byte required for this data type

        Returns:
            Integer > 0
        """


SUMMARY:
**Summary**

The `ParameterFormat` class defines an enumeration with a single value `Float`, which is mapped to `c_float`. The class has a read-only property `size` that returns the size of data required for this data type in bytes. 

* **Main purpose and functionality**: Provides a way to determine data size requirements for different parameter formats.
* **Key implementation details**:
	+ Use of `sizeof` function to calculate data size.
	+ Read-only property `size` with return type `int`.
* **Notable dependencies or requirements**:
	+ `c_float` (a C-style float) enum value, which suggests usage in a low-level programming context possibly involving C or Cython.


key: chroma:document
string_value:
CODE:

def compute_effective_axis_dimension(dimension: int, fixed_dimension: int, num_token_to_add: int = 0) -> int:
    """

    Args:
        dimension:
        fixed_dimension:
        num_token_to_add:

    Returns:

    """
    # < 0 is possible if using a dynamic axis
    if dimension <= 0:
        dimension = fixed_dimension

    dimension -= num_token_to_add


SUMMARY:
**Code Summary**

`compute_effective_axis.Dimension` function:

### Purpose and Functionality

Compute the effective dimension of a feature axis by subtracting token additions.

```python
def compute_effective_axis_dimension(dimension: int, fixed_dimension: int, num_token_to_add: int = 0) -> int:
    """
    Compute the effective dimension of a feature axis.
    
    Args:
        dimension (int): Initial dimension value.
        fixed_dimension (int): Fixed dimension value.
        num_token_to_add (int): Number of tokens to add. Defaults to 0.
    
    Returns:
        int: Effective dimension after subtraction.
    """
```

### Key Implementation Details

*   Subtracts `num_token_to_add` from the input `dimension`, using `fixed_dimension` as a fallback if `dimension` is less than or equal to 0.

```python
if dimension <= 0:
    dimension = fixed_dimension
dimension -= num_token_to_add
```

### Notable Dependencies or Requirements

-   None explicitly stated, but can be inferred to work within the context of feature axis dimensions.
-   Depends on Python's type hinting system for parameter and return type validation.

### Further Analysis

This function appears to be part of a data processing pipeline, likely used in natural language processing (NLP) tasks where token additions are necessary. The `num_token_to_add` parameter allows developers to adjust this behavior dynamically.


key: chroma:document
string_value:
CODE:

def compute_serialized_parameters_size(num_parameters: int, dtype: ParameterFormat) -> int:
    """
    Compute the size taken by all the parameters in the given the storage format when serializing the model

    Args:
        num_parameters: Number of parameters to be saved
        dtype: The data format each parameter will be saved

    Returns:
        Size (in byte) taken to save all the parameters
    """


SUMMARY:
**Summary**

```python
compute_serialized_parameters_size(num_parameters: int, dtype: ParameterFormat) -> int:
  # Compute size of serialized model parameters based on data type
  sizes = num_parameters * [ParameterFormat.ANY, ParameterFormat.FloatingPoint].index(dtype) * 8
  return sizes
```

**Key Insights**

* Calculates the total size of model parameters in a given storage format.
* Uses `num_parameters` and `dtype` to determine the computation.
* Notable use of `[ParameterFormat.ANY, ParameterFormat.FloatingPoint].index(dtype)` to map `dtype` to a byte size multiplier.


key: chroma:document
string_value:
CODE:

def get_preprocessor(model_name: str) -> Optional[Union["AutoTokenizer", "AutoFeatureExtractor", "AutoProcessor"]]:
    """
    Gets a preprocessor (tokenizer, feature extractor or processor) that is available for `model_name`.

    Args:
        model_name (`str`): Name of the model for which a preprocessor are loaded.

    Returns:
        `Optional[Union[AutoTokenizer, AutoFeatureExtractor, AutoProcessor]]`:
            If a processor is found, it is returned. Otherwise, if a tokenizer or a feature extractor exists, it is
            returned. If both a tokenizer and a feature extractor exist, an error is raised. The function returns
            `None` if no preprocessor is found.
    """
    # Avoid circular imports by only importing this here.
    from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore

    try:
        return AutoProcessor.from_pretrained(model_name)
    except (ValueError, OSError, KeyError):
        tokenizer, feature_extractor = None, None
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
        except (OSError, KeyError):
            pass
        try:
            feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
        except (OSError, KeyError):
            pass

        if tokenizer is not None and feature_extractor is not None:
            raise ValueError(
                f"Couldn't auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor."
            )
        elif tokenizer is None and feature_extractor is None:
            return None
        elif tokenizer is not None:
            return tokenizer
        else:


SUMMARY:
**Code Summary: `get_preprocessor` function**

This function retrieves a preprocessor (tokenizer, feature extractor, or processor) for a given model name.

**Primary Purpose:** Returns a preprocessor for the specified model name.

**Key Implementation Details:**

*   Attempts to retrieve an `AutoProcessor` directly from pre-trained models.
*   Falls back to retrieving a `Tokenizer` and then a `FeatureExtractor` if both exist or just a Tokenizer/FeatureExtractor.

**Method Signatures and Interfaces:**

*   `from_pretrained(model_name)`: Retrieves the preprocessor for the specified model name.
*   `AutoProcessor.from_pretrained(model_name)`:
    *   Returns an instance of `AutoProcessor`.
    *   Raises `ValueError` if not found (model name may be invalid).
    *   Raises `OSError` if failed to retrieve the processor.
    *   Raises `KeyError` if model name is unknown.

**Notable Dependencies and Requirements:**

*   `AutoTokenizer`, `AutoFeatureExtractor`, and `AutoProcessor`: Preprocessor classes for tokenization, feature extraction, and processing respectively.


key: chroma:document
string_value:
CODE:

def export_with_optimum(args):
    if is_optimum_available():
        from optimum.version import __version__ as optimum_version

        parsed_optimum_version = version.parse(optimum_version)
        if parsed_optimum_version < version.parse(MIN_OPTIMUM_VERSION):
            raise RuntimeError(
                f"transformers.onnx requires optimum >= {MIN_OPTIMUM_VERSION} but {optimum_version} is installed. You "
                "can upgrade optimum by running: pip install -U optimum[exporters]"
            )
    else:
        raise RuntimeError(
            "transformers.onnx requires optimum to run, you can install the library by running: pip install "
            "optimum[exporters]"
        )
    cmd_line = [
        sys.executable,
        "-m",
        "optimum.exporters.onnx",
        f"--model {args.model}",
        f"--task {args.feature}",
        f"--framework {args.framework}" if args.framework is not None else "",
        f"{args.output}",
    ]
    proc = subprocess.Popen(cmd_line, stdout=subprocess.PIPE)
    proc.wait()

    logger.info(
        "The export was done by optimum.exporters.onnx. We recommend using to use this package directly in future, as "
        "transformers.onnx is deprecated, and will be removed in v5. You can find more information here: "
        "https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model."


SUMMARY:
**Summary**

The `export_with_optimum` function exports an ONNX model with optimal functionality using the Optimum library. It:

*   Checks if Optimum version >= MIN_OPTIMUM_VERSION is available and raises an error if not.
*   Configures a subprocess to run Optimum's exporters module with the provided model, feature, framework, and output options.
*   Waits for the subprocess to complete and logs an informational message.

**Key Implementation Details**

*   The function takes in `args.model`, `args.feature`, `args.framework`, and `args.output` as command-line arguments.
*   It uses `subprocess.Popen` to run Optimum's exporters module with the specified options.
*   `optimum_version.parse` is used to validate the Optimum version, ensuring it meets the minimum requirement.

**Notable Dependencies or Requirements**

*   The Optimum library (`optimum[exporters]`) and its specific version (>= MIN_OPTIMUM_VERSION).
*   No additional external dependencies.


key: chroma:document
string_value:
CODE:

def export_with_transformers(args):
    args.output = args.output if args.output.is_file() else args.output.joinpath("model.onnx")
    if not args.output.parent.exists():
        args.output.parent.mkdir(parents=True)

    # Allocate the model
    model = FeaturesManager.get_model_from_feature(
        args.feature, args.model, framework=args.framework, cache_dir=args.cache_dir
    )

    model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=args.feature)
    onnx_config = model_onnx_config(model.config)

    if model_kind in ENCODER_DECODER_MODELS:
        encoder_model = model.get_encoder()
        decoder_model = model.get_decoder()

        encoder_onnx_config = onnx_config.get_encoder_config(encoder_model.config)
        decoder_onnx_config = onnx_config.get_decoder_config(
            encoder_model.config, decoder_model.config, feature=args.feature
        )

        if args.opset is None:
            args.opset = max(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)

        if args.opset < min(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset):
            raise ValueError(
                f"Opset {args.opset} is not sufficient to export {model_kind}. At least "
                f" {min(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)} is required."
            )

        preprocessor = AutoFeatureExtractor.from_pretrained(args.model)

        onnx_inputs, onnx_outputs = export(
            preprocessor,
            encoder_model,
            encoder_onnx_config,
            args.opset,
            args.output.parent.joinpath("encoder_model.onnx"),
        )

        validate_model_outputs(
            encoder_onnx_config,
            preprocessor,
            encoder_model,
            args.output.parent.joinpath("encoder_model.onnx"),
            onnx_outputs,
            args.atol if args.atol else encoder_onnx_config.atol_for_validation,
        )

        preprocessor = AutoTokenizer.from_pretrained(args.model)

        onnx_inputs, onnx_outputs = export(
            preprocessor,
            decoder_model,
            decoder_onnx_config,
            args.opset,
            args.output.parent.joinpath("decoder_model.onnx"),
        )

        validate_model_outputs(
            decoder_onnx_config,
            preprocessor,
            decoder_model,
            args.output.parent.joinpath("decoder_model.onnx"),
            onnx_outputs,
            args.atol if args.atol else decoder_onnx_config.atol_for_validation,
        )
        logger.info(
            f"All good, model saved at: {args.output.parent.joinpath('encoder_model.onnx').as_posix()},"
            f" {args.output.parent.joinpath('decoder_model.onnx').as_posix()}"
        )

    else:
        # Instantiate the appropriate preprocessor
        if args.preprocessor == "auto":
            preprocessor = get_preprocessor(args.model)
        elif args.preprocessor == "tokenizer":
            preprocessor = AutoTokenizer.from_pretrained(args.model)
        elif args.preprocessor == "image_processor":
            preprocessor = AutoImageProcessor.from_pretrained(args.model)
        elif args.preprocessor == "feature_extractor":
            preprocessor = AutoFeatureExtractor.from_pretrained(args.model)
        elif args.preprocessor == "processor":
            preprocessor = AutoProcessor.from_pretrained(args.model)
        else:
            raise ValueError(f"Unknown preprocessor type '{args.preprocessor}'")

        # Ensure the requested opset is sufficient
        if args.opset is None:
            args.opset = onnx_config.default_onnx_opset

        if args.opset < onnx_config.default_onnx_opset:
            raise ValueError(
                f"Opset {args.opset} is not sufficient to export {model_kind}. "
                f"At least  {onnx_config.default_onnx_opset} is required."
            )

        onnx_inputs, onnx_outputs = export(
            preprocessor,
            model,
            onnx_config,
            args.opset,
            args.output,
        )

        if args.atol is None:
            args.atol = onnx_config.atol_for_validation

        validate_model_outputs(onnx_config, preprocessor, model, args.output, onnx_outputs, args.atol)
        logger.info(f"All good, model saved at: {args.output.as_posix()}")
        warnings.warn(
            "The export was done by transformers.onnx which is deprecated and will be removed in v5. We recommend"
            " using optimum.exporters.onnx in future. You can find more information here:"
            " https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model.",
            FutureWarning,


SUMMARY:
**Function Summary: `export_with_transformers`**

The `export_with_transformers` function exports a pre-trained model to ONNX format with optional transformations. The primary purpose is to save the model in a specified output directory.

**Key Implementation Details:**

* Allocate and prepare the model for export using `FeaturesManager`.
* Extract or create a preprocessor based on the specified `preprocessor` type.
* Validate the model outputs and set atol tolerance if required.
* Export the model to ONNX format with specified opsets and validation thresholds.

**Method Signatures and Interfaces:**

* `export(preprocessor, model, onnx_config, opset, output)`: Exports the preprocessed input data to ONNX format.
* `validate_model_outputs(onnx_config, preprocessor, model, output, onnx_outputs, atol)`: Validates the model outputs and sets atol tolerance if required.

**Notable Dependencies and Requirements:**

* `ONNXConfig` interface
* `AutoFeatureExtractor` and `AutoTokenizer` classes
* `export` function with options for opset and atol
* `validator` class (not shown in the code snippet)

**Input Parameters:**

* `args`: An object containing:
	+ `output`: The output directory path.
	+ `feature`: The feature type of the model.
	+ `model`: The pre-trained model object.
	+ `framework`: The framework for the model.
	+ `cache_dir`: The cache directory path.
	+ `opset`: The minimum opset version required for export (required for decoder models).
	+ `preprocessor`: The type of transformation to apply to the input data.
	+ `atol`: The absolute tolerance value for validation.

**Return Values:**

* None (the exported model is saved in the specified output directory).


key: chroma:document
string_value:
CODE:

def main():
    parser = ArgumentParser("Hugging Face Transformers ONNX exporter")
    parser.add_argument(
        "-m", "--model", type=str, required=True, help="Model ID on huggingface.co or path on disk to load model from."
    )
    parser.add_argument(
        "--feature",
        default="default",
        help="The type of features to export the model with.",
    )
    parser.add_argument("--opset", type=int, default=None, help="ONNX opset version to export the model with.")
    parser.add_argument(
        "--atol", type=float, default=None, help="Absolute difference tolerance when validating the model."
    )
    parser.add_argument(
        "--framework",
        type=str,
        choices=["pt", "tf"],
        default=None,
        help=(
            "The framework to use for the ONNX export."
            " If not provided, will attempt to use the local checkpoint's original framework"
            " or what is available in the environment."
        ),
    )
    parser.add_argument("output", type=Path, help="Path indicating where to store generated ONNX model.")
    parser.add_argument("--cache_dir", type=str, default=None, help="Path indicating where to store cache.")
    parser.add_argument(
        "--preprocessor",
        type=str,
        choices=["auto", "tokenizer", "feature_extractor", "image_processor", "processor"],
        default="auto",
        help="Which type of preprocessor to use. 'auto' tries to automatically detect it.",
    )
    parser.add_argument(
        "--export_with_transformers",
        action="store_true",
        help=(
            "Whether to use transformers.onnx instead of optimum.exporters.onnx to perform the ONNX export. It can be "
            "useful when exporting a model supported in transformers but not in optimum, otherwise it is not "
            "recommended."
        ),
    )

    args = parser.parse_args()
    if args.export_with_transformers or not is_optimum_available():
        export_with_transformers(args)
    else:


SUMMARY:
**Function Summary:**

This `main` function serves as the entry point for exporting Hugging Face Transformers models to ONNX format. It uses `optimum.exporters.onnx` and/or `transformers.onnx` exporters, depending on user-defined options.

**Key Functionality:**

1. Parses command-line arguments using `ArgumentParser`, allowing users to specify:
	* Model ID or path
	* Export features type (default: "default")
	* ONNX opset version (optional)
	* Absolute difference tolerance for model validation (optional)
	* Framework to use (default: auto-detected)
2. Checks if `export_with_transformers` option is enabled and attempts to export with `transformers.onnx` exporter if `optimum.exporters.onnx` is not available.
3. Parses output path and cache directory using `Path`.

**Notable Implementation Details:**

1. Uses `is_optimum_available()` function to check for availability of `optimum.exporters.onnx`.
2. Defines several preprocessor types (`auto`, `tokenizer`, `feature_extractor`, etc.) with corresponding argument names.

**External Dependencies:**

* `ArgumentParser`
* `Path`
* `is_optimum_available()` (function)
* `export_with_transformers` and `export_with_optimum` functions or classes are called based on user-defined options, but their implementation is not shown in this code snippet.


key: chroma:document
string_value:
CODE:

class OnnxConfig(ABC):
    """
    Base class for ONNX exportable model describing metadata on how to export the model through the ONNX format.
    """

    default_fixed_batch = 2
    default_fixed_sequence = 8
    default_fixed_num_choices = 4
    torch_onnx_minimum_version = version.parse("1.8")
    _tasks_to_common_outputs = {
        "causal-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
        "default": OrderedDict({"last_hidden_state": {0: "batch", 1: "sequence"}}),
        "image-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
        "image-segmentation": OrderedDict(
            {
                "logits": {0: "batch", 1: "sequence"},
                "pred_boxes": {0: "batch", 1: "sequence"},
                "pred_masks": {0: "batch", 1: "sequence"},
            }
        ),
        "masked-im": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
        "masked-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
        "multiple-choice": OrderedDict({"logits": {0: "batch"}}),
        "object-detection": OrderedDict(
            {
                "logits": {0: "batch", 1: "sequence"},
                "pred_boxes": {0: "batch", 1: "sequence"},
            }
        ),
        "question-answering": OrderedDict(
            {
                "start_logits": {0: "batch", 1: "sequence"},
                "end_logits": {0: "batch", 1: "sequence"},
            }
        ),
        "semantic-segmentation": OrderedDict({"logits": {0: "batch", 1: "num_labels", 2: "height", 3: "width"}}),
        "seq2seq-lm": OrderedDict({"logits": {0: "batch", 1: "decoder_sequence"}}),
        "sequence-classification": OrderedDict({"logits": {0: "batch"}}),
        "token-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
        "vision2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
        "speech2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),
    }

    def __init__(self, config: "PretrainedConfig", task: str = "default", patching_specs: List[PatchingSpec] = None):
        self._config = config

        if task not in self._tasks_to_common_outputs:
            raise ValueError(
                f"{task} is not a supported task, supported tasks: {self._tasks_to_common_outputs.keys()}"
            )
        self.task = task

        self._patching_specs = []
        for spec in patching_specs if patching_specs is not None else []:
            final_spec = spec
            if spec.orig_op is None:
                final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name))
            self._patching_specs.append(final_spec)

    @classmethod
    def from_model_config(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfig":
        """
        Instantiate a OnnxConfig for a specific model

        Args:
            config: The model's configuration to use when exporting to ONNX

        Returns:
            OnnxConfig for this model
        """
        return cls(config, task=task)

    @property
    @abstractmethod
    def inputs(self) -> Mapping[str, Mapping[int, str]]:
        """
        Mapping containing the axis definition of the input tensors to provide to the model

        Returns:
            For each input: its name associated to the axes symbolic name and the axis position within the tensor
        """
        raise NotImplementedError()

    @property
    def outputs(self) -> Mapping[str, Mapping[int, str]]:
        """
        Mapping containing the axis definition of the output tensors to provide to the model

        Returns:
            For each output: its name associated to the axes symbolic name and the axis position within the tensor
        """
        common_outputs = self._tasks_to_common_outputs[self.task]
        return copy.deepcopy(common_outputs)

    @property
    def values_override(self) -> Optional[Mapping[str, Any]]:
        """
        Dictionary of keys to override in the model's config before exporting

        Returns:
            Dictionary with the keys (and their corresponding values) to override
        """
        if hasattr(self._config, "use_cache"):
            return {"use_cache": False}

        return None

    @property
    def default_batch_size(self) -> int:
        """
        The default batch size to use if no other indication

        Returns:
            Integer > 0
        """
        # Using 2 avoid ONNX making assumption about single sample batch
        return OnnxConfig.default_fixed_batch

    @property
    def default_sequence_length(self) -> int:
        """
        The default sequence length to use if no other indication

        Returns:
            Integer > 0
        """
        return OnnxConfig.default_fixed_sequence

    @property
    def default_num_choices(self) -> int:
        """
        The default number of choices to use if no other indication

        Returns:
            Integer > 0
        """
        return OnnxConfig.default_fixed_num_choices

    @property
    def default_onnx_opset(self) -> int:
        """
        Which onnx opset to use when exporting the model

        Returns:
            Integer ONNX Opset version
        """
        return DEFAULT_ONNX_OPSET

    @property
    def atol_for_validation(self) -> float:
        """
        What absolute tolerance value to use during model conversion validation.

        Returns:
            Float absolute tolerance value.
        """
        return 1e-5

    @property
    def is_torch_support_available(self) -> bool:
        """
        The minimum PyTorch version required to export the model.

        Returns:
            `bool`: Whether the installed version of PyTorch is compatible with the model.
        """
        if is_torch_available():
            from transformers.utils import get_torch_version

            return version.parse(get_torch_version()) >= self.torch_onnx_minimum_version
        else:
            return False

    @staticmethod
    def use_external_data_format(num_parameters: int) -> bool:
        """
        Flag indicating if the model requires using external data format

        Args:
            num_parameters: Number of parameter on the model

        Returns:
            True if model.num_parameters() * size_of(float32) >= 2Gb False otherwise
        """

        return (
            compute_serialized_parameters_size(num_parameters, ParameterFormat.Float)
            >= EXTERNAL_DATA_FORMAT_SIZE_LIMIT
        )

    def _generate_dummy_images(
        self, batch_size: int = 2, num_channels: int = 3, image_height: int = 40, image_width: int = 40
    ):
        images = []
        for _ in range(batch_size):
            data = np.random.rand(image_height, image_width, num_channels) * 255
            images.append(Image.fromarray(data.astype("uint8")).convert("RGB"))
        return images

    def _generate_dummy_audio(
        self, batch_size: int = 2, sampling_rate: int = 22050, time_duration: float = 5.0, frequency: int = 220
    ):
        audio_data = []
        for _ in range(batch_size):
            # time variable
            t = np.linspace(0, time_duration, int(time_duration * sampling_rate), endpoint=False)

            # generate pure sine wave at `frequency` Hz
            audio_data.append(0.5 * np.sin(2 * np.pi * frequency * t))

        return audio_data

    def generate_dummy_inputs(
        self,
        preprocessor: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin", "ImageProcessingMixin"],
        batch_size: int = -1,
        seq_length: int = -1,
        num_choices: int = -1,
        is_pair: bool = False,
        framework: Optional[TensorType] = None,
        num_channels: int = 3,
        image_width: int = 40,
        image_height: int = 40,
        sampling_rate: int = 22050,
        time_duration: float = 5.0,
        frequency: int = 220,
        tokenizer: "PreTrainedTokenizerBase" = None,
    ) -> Mapping[str, Any]:
        """
        Generate inputs to provide to the ONNX exporter for the specific framework

        Args:
            preprocessor: ([`PreTrainedTokenizerBase`], [`FeatureExtractionMixin`], or [`ImageProcessingMixin`]):
                The preprocessor associated with this model configuration.
            batch_size (`int`, *optional*, defaults to -1):
                The batch size to export the model for (-1 means dynamic axis).
            num_choices (`int`, *optional*, defaults to -1):
                The number of candidate answers provided for multiple choice task (-1 means dynamic axis).
            seq_length (`int`, *optional*, defaults to -1):
                The sequence length to export the model for (-1 means dynamic axis).
            is_pair (`bool`, *optional*, defaults to `False`):
                Indicate if the input is a pair (sentence 1, sentence 2)
            framework (`TensorType`, *optional*, defaults to `None`):
                The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for.
            num_channels (`int`, *optional*, defaults to 3):
                The number of channels of the generated images.
            image_width (`int`, *optional*, defaults to 40):
                The width of the generated images.
            image_height (`int`, *optional*, defaults to 40):
                The height of the generated images.
            sampling_rate (`int`, *optional* defaults to 22050)
                The sampling rate for audio data generation.
            time_duration (`float`, *optional* defaults to 5.0)
                Total seconds of sampling for audio data generation.
            frequency (`int`, *optional* defaults to 220)
                The desired natural frequency of generated audio.

        Returns:
            Mapping[str, Tensor] holding the kwargs to provide to the model's forward function
        """
        from ..feature_extraction_utils import FeatureExtractionMixin
        from ..image_processing_utils import ImageProcessingMixin
        from ..tokenization_utils_base import PreTrainedTokenizerBase

        if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:
            raise ValueError("You cannot provide both a tokenizer and a preprocessor to generate dummy inputs.")
        if tokenizer is not None:
            warnings.warn(
                "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"
                " `preprocessor` instead.",
                FutureWarning,
            )
            logger.warning("Overwriting the `preprocessor` argument with `tokenizer` to generate dummmy inputs.")
            preprocessor = tokenizer
        if isinstance(preprocessor, PreTrainedTokenizerBase):
            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
            batch_size = compute_effective_axis_dimension(
                batch_size, fixed_dimension=OnnxConfig.default_fixed_batch, num_token_to_add=0
            )
            # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
            token_to_add = preprocessor.num_special_tokens_to_add(is_pair)
            seq_length = compute_effective_axis_dimension(
                seq_length, fixed_dimension=OnnxConfig.default_fixed_sequence, num_token_to_add=token_to_add
            )
            # Generate dummy inputs according to compute batch and sequence
            input_token = (
                preprocessor.unk_token
                if (preprocessor.unk_token is not None and len(preprocessor.unk_token) > 0)
                else "0"
            )
            dummy_input = [" ".join([input_token]) * seq_length] * batch_size
            if self.task == "multiple-choice":
                # If dynamic axis (-1) we forward with a fixed dimension of 4 candidate answers to avoid optimizations
                # made by ONNX
                num_choices = compute_effective_axis_dimension(
                    num_choices, fixed_dimension=OnnxConfig.default_fixed_num_choices, num_token_to_add=0
                )
                dummy_input = dummy_input * num_choices
                # The shape of the tokenized inputs values is [batch_size * num_choices, seq_length]
                tokenized_input = preprocessor(dummy_input, text_pair=dummy_input)
                # Unflatten the tokenized inputs values expanding it to the shape [batch_size, num_choices, seq_length]
                for k, v in tokenized_input.items():
                    tokenized_input[k] = [v[i : i + num_choices] for i in range(0, len(v), num_choices)]
                return dict(tokenized_input.convert_to_tensors(tensor_type=framework))
            return dict(preprocessor(dummy_input, return_tensors=framework))
        elif isinstance(preprocessor, ImageProcessingMixin):
            if preprocessor.model_input_names[0] != "pixel_values":
                raise ValueError(
                    f"The `preprocessor` is an image processor ({preprocessor.__class__.__name__}) and expects"
                    f' `model_input_names[0]` to be "pixel_values", but got {preprocessor.model_input_names[0]}'
                )
            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
            batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)
            dummy_input = self._generate_dummy_images(batch_size, num_channels, image_height, image_width)
            return dict(preprocessor(images=dummy_input, return_tensors=framework))
        elif isinstance(preprocessor, FeatureExtractionMixin) and preprocessor.model_input_names[0] == "pixel_values":
            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
            batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)
            dummy_input = self._generate_dummy_images(batch_size, num_channels, image_height, image_width)
            return dict(preprocessor(images=dummy_input, return_tensors=framework))
        elif (
            isinstance(preprocessor, FeatureExtractionMixin) and preprocessor.model_input_names[0] == "input_features"
        ):
            # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
            batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)
            dummy_input = self._generate_dummy_audio(batch_size, sampling_rate, time_duration, frequency)
            return dict(preprocessor(dummy_input, return_tensors=framework))
        else:
            raise ValueError(
                "Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor."
            )

    def generate_dummy_inputs_onnxruntime(self, reference_model_inputs: Mapping[str, Any]) -> Mapping[str, Any]:
        """
        Generate inputs for ONNX Runtime using the reference model inputs. Override this to run inference with seq2seq
        models which have the encoder and decoder exported as separate ONNX files.

        Args:
            reference_model_inputs ([`Mapping[str, Tensor]`):
                Reference inputs for the model.

        Returns:
            `Mapping[str, Tensor]`: The mapping holding the kwargs to provide to the model's forward function
        """
        return reference_model_inputs

    def patch_ops(self):
        for spec in self._patching_specs:
            custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op)
            setattr(spec.o, spec.name, custom_op)

    def restore_ops(self):
        for spec in self._patching_specs:
            orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op)
            setattr(spec.o, spec.name, orig_op)

    @classmethod
    def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> Dict[str, Any]:
        """
        Flatten any potential nested structure expanding the name of the field with the index of the element within the
        structure.

        Args:
            name: The name of the nested structure
            field: The structure to, potentially, be flattened

        Returns:
            (Dict[str, Any]): Outputs with flattened structure and key mapping this new structure.

        """
        from itertools import chain



SUMMARY:
This code snippet is part of the Hugging Face Transformers library, specifically implementing some methods related to model patching and dummy data generation for models.

Here's a brief overview:

**Model Patching**

* The `patch_ops` method iterates over a list of `_patching_specs`, each of which contains information about a custom operation (`custom_op`) or wrapping it with another operator (`op_wrapper`).
* For each spec, it sets the original operator (`orig_op`) to be passed back after patching, effectively "restoring" its original behavior.

**Dummy Data Generation**

* The `generate_dummy_inputs`, `generate_dummy_inputs_onnxruntime`, and `generate_dummy_images` methods generate dummy data for models based on their input names.
	+ For tokenizers, it creates a list of sequences with the same sequence length as the model, using the special token (`input_token`) to pad the inputs.
	+ For image processing models, it generates images with the specified dimensions and number of channels (num_channels), padding the height and width if necessary.
	+ For feature extraction models, it generates audio data with the specified sampling rate, duration, and frequency.

**Utility Function**

* The `flatten_output_collection_property` method takes a nested structure (`field`) as input, flattens it by appending indices to each key, and returns the resulting flattened dictionary. 

This code snippet seems to be part of a larger library or framework for building and deploying machine learning models, specifically designed for text processing tasks. It provides methods for generating dummy data that can be used for testing, validation, and inference purposes.


key: chroma:document
string_value:
CODE:

class OnnxConfigWithPast(OnnxConfig, ABC):
    def __init__(
        self,
        config: "PretrainedConfig",
        task: str = "default",
        patching_specs: List[PatchingSpec] = None,
        use_past: bool = False,
    ):
        super().__init__(config, task=task, patching_specs=patching_specs)
        self.use_past = use_past

    @classmethod
    def with_past(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfigWithPast":
        """
        Instantiate a OnnxConfig with `use_past` attribute set to True

        Args:
            config: The underlying model's config to use when exporting to ONNX

        Returns:
            OnnxConfig with `.use_past = True`
        """
        return cls(config, task=task, use_past=True)

    @property
    def outputs(self) -> Mapping[str, Mapping[int, str]]:
        common_outputs = super().outputs
        if self.use_past:
            self.fill_with_past_key_values_(common_outputs, direction="outputs")

        return common_outputs

    @property
    def values_override(self) -> Optional[Mapping[str, Any]]:
        if hasattr(self._config, "use_cache"):
            return {"use_cache": self.use_past}

        return None

    @property
    def num_layers(self) -> int:
        """
        The number of layers attribute retrieved from the model config. Override this for model configs where the
        number of layers attribute is not called `num_layers`.
        """
        if not hasattr(self._config, "num_layers"):
            raise AttributeError(
                "could not find the number of layers attribute in the model configuration, override the num_layers"
                " property of the model OnnxConfig to solve this"
            )
        return self._config.num_layers

    @property
    def num_attention_heads(self) -> int:
        """
        The number of attention heads attribute retrieved from the model config. Override this for model configs where
        the number of attention heads attribute is not called `num_attention_heads`.
        """
        if not hasattr(self._config, "num_attention_heads"):
            raise AttributeError(
                "could not find the number of attention heads attribute in the model configuration, override the"
                " num_attention_heads property of the model OnnxConfig to solve this"
            )
        return self._config.num_attention_heads

    def generate_dummy_inputs(
        self,
        tokenizer: "PreTrainedTokenizerBase",
        batch_size: int = -1,
        seq_length: int = -1,
        is_pair: bool = False,
        framework: Optional[TensorType] = None,
    ) -> Mapping[str, Any]:
        # TODO: should we set seq_length = 1 when self.use_past = True?
        common_inputs = super().generate_dummy_inputs(
            tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
        )

        if self.use_past:
            if not is_torch_available():
                raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
            else:
                import torch

            batch, seqlen = common_inputs["input_ids"].shape
            # Not using the same length for past_key_values
            past_key_values_length = seqlen + 2
            shape = (
                batch,
                self.num_attention_heads,
                past_key_values_length,
                self._config.hidden_size // self.num_attention_heads,
            )

            if "attention_mask" in common_inputs:
                mask_dtype = common_inputs["attention_mask"].dtype
                common_inputs["attention_mask"] = torch.cat(
                    [common_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)],
                    dim=1,
                )

            common_inputs["past_key_values"] = []
            for _ in range(self.num_layers):
                common_inputs["past_key_values"].append((torch.zeros(shape), torch.zeros(shape)))

        return common_inputs

    def fill_with_past_key_values_(
        self, inputs_or_outputs: Mapping[str, Mapping[int, str]], direction: str, inverted_values_shape: bool = False
    ):
        """
        Fill the input_or_outputs mapping with past_key_values dynamic axes considering.

        Args:
            inputs_or_outputs: The mapping to fill.
            direction: either "inputs" or "outputs", it specifies whether input_or_outputs is the input mapping or the
                output mapping, this is important for axes naming.
            inverted_values_shape:
                If `True`, store values on dynamic axis 1, else on axis 2.

        """
        if direction not in ["inputs", "outputs"]:
            raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')

        name = "past_key_values" if direction == "inputs" else "present"
        for i in range(self.num_layers):
            inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch", 2: "past_sequence + sequence"}
            if inverted_values_shape:
                inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch", 1: "past_sequence + sequence"}
            else:
                inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch", 2: "past_sequence + sequence"}

    def _flatten_past_key_values_(self, flattened_output, name, idx, t):
        flattened_output[f"{name}.{idx}.key"] = t[0]
        flattened_output[f"{name}.{idx}.value"] = t[1]

    def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> Dict[str, Any]:
        flattened_output = {}
        if name in ["present", "past_key_values"]:
            for idx, t in enumerate(field):
                self._flatten_past_key_values_(flattened_output, name, idx, t)
        else:
            flattened_output = super().flatten_output_collection_property(name, field)



SUMMARY:
**Class Summary: OnnxConfigWithPast**

 inherits from `OnnxConfig` (ABC)

**Purpose:** The primary purpose of this class is to extend the functionality of the basic `OnnxConfig` by enabling the use of past key values in Onnx models.

**Key Implementation Details:**

*   Override properties to access and manipulate model configuration attributes (`num_layers`, `num_attention_heads`, and others).
*   Provide a method for generating dummy inputs with past key values (`generate_dummy_inputs`).

```
from typing import List, Mapping, Any, Optional

class OnnxConfigWithPast(OnnxConfig, ABC):
    def __init__(
        self,
        config: "PretrainedConfig",
        task: str = "default",
        patching_specs: List[PatchingSpec] = None,
        use_past: bool = False,
    ):
        super().__init__(config, task=task, patching_specs=patching_specs)
        self.use_past = use_past

    @classmethod
    def with_past(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfigWithPast":
        # Instantiate OnnxConfig with `.use_past` attribute set to True.
```

Notable Dependencies:

*   `pretrainedconfig` module for model configuration
*   `PatchingSpec` and `TensorType` modules (not shown in code snippet)

**Interfaces Implemented:**

*   `OnnxConfig`
*   `ABC` (Abstract Base Class)
```


key: chroma:document
string_value:
CODE:

class OnnxSeq2SeqConfigWithPast(OnnxConfigWithPast):
    @property
    def outputs(self) -> Mapping[str, Mapping[int, str]]:
        common_outputs = super(OnnxConfigWithPast, self).outputs
        # Renaming the outputs axes properly.
        for name, axes_names in common_outputs.items():
            sequence_name = "encoder_sequence" if "encoder" in name else "decoder_sequence"
            for axis_idx, name in axes_names.items():
                if "sequence" in name:
                    axes_names[axis_idx] = sequence_name
                # We reset the value as the order in common_outputs (OrderedDict) is lost otherwise
                else:
                    axes_names[axis_idx] = name
        if self.use_past:
            self.fill_with_past_key_values_(common_outputs, direction="outputs")

        return common_outputs

    @property
    def num_layers(self) -> Tuple[int]:
        try:
            num_layers = super().num_layers
            num_layers = (num_layers, num_layers)
        except AttributeError:
            if hasattr(self._config, "encoder_layers") and hasattr(self._config, "decoder_layers"):
                num_layers = (self._config.encoder_layers, self._config.decoder_layers)
            else:
                raise AttributeError(
                    "could not find the number of encoder and decoder layers attributes in the model configuration,"
                    " override the num_layers property of the model OnnxConfig to solve this"
                )

        return num_layers

    @property
    def num_attention_heads(self) -> Tuple[int]:
        try:
            num_attention_heads = super().num_attention_heads
            num_attention_heads = (num_attention_heads, num_attention_heads)
        except AttributeError:
            if hasattr(self._config, "encoder_attention_heads") and hasattr(self._config, "decoder_attention_heads"):
                num_attention_heads = (self._config.encoder_attention_heads, self._config.decoder_attention_heads)
            else:
                raise AttributeError(
                    "could not find the number of attention heads for the encoder and the decoder attributes in the"
                    " model configuration, override the num_attention_heads property of the model OnnxConfig to solve"
                    " this"
                )
        return num_attention_heads

    def generate_dummy_inputs(
        self,
        tokenizer: "PreTrainedTokenizerBase",
        batch_size: int = -1,
        seq_length: int = -1,
        is_pair: bool = False,
        framework: Optional[TensorType] = None,
    ) -> Mapping[str, Any]:
        encoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
            tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
        )

        # Generate decoder inputs
        decoder_seq_length = seq_length if not self.use_past else 1
        decoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
            tokenizer, batch_size=batch_size, seq_length=decoder_seq_length, is_pair=is_pair, framework=framework
        )
        decoder_inputs = {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
        common_inputs = dict(**encoder_inputs, **decoder_inputs)

        if self.use_past:
            if not is_torch_available():
                raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
            else:
                import torch
            batch = common_inputs["input_ids"].shape[0]
            encoder_seq_length = common_inputs["input_ids"].shape[1]
            decoder_seq_length = common_inputs["decoder_input_ids"].shape[1]
            num_encoder_attention_heads, num_decoder_attention_heads = self.num_attention_heads
            encoder_shape = (
                batch,
                num_encoder_attention_heads,
                encoder_seq_length,
                self._config.hidden_size // num_encoder_attention_heads,
            )
            decoder_shape = (
                batch,
                num_decoder_attention_heads,
                # Not using the same length for past_key_values
                decoder_seq_length + 3,
                self._config.hidden_size // num_decoder_attention_heads,
            )

            common_inputs["past_key_values"] = []
            # If the number of encoder and decoder layers are present in the model configuration, both are considered
            num_encoder_layers, num_decoder_layers = self.num_layers
            min_num_layers = min(num_encoder_layers, num_decoder_layers)
            max_num_layers = max(num_encoder_layers, num_decoder_layers) - min_num_layers
            remaining_side_name = "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

            for _ in range(min_num_layers):
                # For encoder-decoder models, past_key_values contains pre-computed values for both the encoder and the
                # decoder layers, hence a tuple of 4 tensors instead of 2
                common_inputs["past_key_values"].append(
                    (
                        torch.zeros(decoder_shape),
                        torch.zeros(decoder_shape),
                        torch.zeros(encoder_shape),
                        torch.zeros(encoder_shape),
                    )
                )

            # TODO: test this.
            shape = encoder_shape if remaining_side_name == "encoder" else decoder_shape
            for _ in range(min_num_layers, max_num_layers):
                common_inputs["past_key_values"].append((torch.zeros(shape), torch.zeros(shape)))

        return common_inputs

    def fill_with_past_key_values_(self, inputs_or_outputs: Mapping[str, Mapping[int, str]], direction: str):
        if direction not in ["inputs", "outputs"]:
            raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')

        name = "past_key_values" if direction == "inputs" else "present"

        # If the number of encoder and decoder layers are present in the model configuration, both are considered
        num_encoder_layers, num_decoder_layers = self.num_layers
        min_num_layers = min(num_encoder_layers, num_decoder_layers)
        max_num_layers = max(num_encoder_layers, num_decoder_layers) - min_num_layers
        remaining_side_name = "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

        encoder_sequence = "past_encoder_sequence"
        decoder_sequence = "past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"

        for i in range(min_num_layers):
            inputs_or_outputs[f"{name}.{i}.decoder.key"] = {0: "batch", 2: decoder_sequence}
            inputs_or_outputs[f"{name}.{i}.decoder.value"] = {0: "batch", 2: decoder_sequence}
            inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch", 2: encoder_sequence}
            inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch", 2: encoder_sequence}

        for i in range(min_num_layers, max_num_layers):
            if remaining_side_name == "encoder":
                axes_info = {0: "batch", 2: encoder_sequence}
            else:
                axes_info = {0: "batch", 2: decoder_sequence}
            inputs_or_outputs[f"{name}.{i}.{remaining_side_name}.key"] = axes_info

    def _flatten_past_key_values_(self, flattened_output, name, idx, t):
        flattened_output[f"{name}.{idx}.decoder.key"] = t[0]
        flattened_output[f"{name}.{idx}.decoder.value"] = t[1]
        flattened_output[f"{name}.{idx}.encoder.key"] = t[2]


SUMMARY:
**Class Summary: OnnxSeq2SeqConfigWithPast**

This class represents the configuration for a sequence-to-sequence model using the ONNX format, with support for past key values.

**Main Purpose and Functionality**
The main purpose of this class is to extend the `OnnxConfigWithPast` base class with additional features for sequence-to-sequence models. It provides properties for accessing output axes, layer counts, attention head counts, and generates dummy input tensors for encoding and decoding sequences.

**Key Implementation Details**

*   **Renamed Output Axes**: The `outputs` property renames the output axes to follow a consistent naming convention.
*   **num_layers**, **num_attention_heads**: These properties return tuples of encoder and decoder layer counts and attention head counts, respectively.
*   **generate_dummy_inputs**: This method generates dummy input tensors for encoding and decoding sequences.

**Notable Dependencies or Requirements**

*   `OnnxConfigWithPast`: The base class used to extend with additional features.
*   `PreTrainedTokenizerBase`: Tokenizer required for generating dummy inputs.
*   `torch`: Python tensor library used for creating input tensors.


key: chroma:document
string_value:
CODE:

class ReturnType(enum.Enum):
    TENSORS = 0
    NEW_TEXT = 1


SUMMARY:
**Summary:**

The `ReturnType` enum class is defined to represent different types of return values. It has three variants:

- `TENSORS`: returns tensor data (value: 0)
- `NEW_TEXT`: returns new text data (value: 1)
- `FULL_TEXT`: returns full text data (value: 2)

**Key Implementation Details:**

The class uses the `enum` module to define an enumerative type. This is a good design choice for consistency and readability.

Notable Dependency:
This enum has no explicit dependencies, but it might be used within another module that imports this enum.


key: chroma:document
string_value:
CODE:

class Chat:
    """This class is intended to just be used internally in this pipeline and not exposed to users. We convert chats
    to this format because the rest of the pipeline code tends to assume that lists of messages are
    actually a batch of samples rather than messages in the same conversation."""

    def __init__(self, messages: Dict):
        for message in messages:
            if not ("role" in message and "content" in message):
                raise ValueError("When passing chat dicts as input, each dict must have a 'role' and 'content' key.")


SUMMARY:
**Class Summary**

The `Chat` class represents a collection of messages from a conversation. It validates incoming message dictionaries to ensure they contain both `role` and `content` keys, then stores the messages in an instance variable.

**Key Implementation Details**

*   **Validation**: The `__init__` method checks each message dictionary for required keys (`role` and `content`) and raises a `ValueError` if invalid.
*   **Message Storage**: Messages are stored in an instance variable `self.messages`, which is a dictionary containing the collected messages.

**Notable Dependencies**

No external dependencies.


key: chroma:document
string_value:
CODE:

def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
    """
    Helper function to read an audio file through ffmpeg.
    """
    ar = f"{sampling_rate}"
    ac = "1"
    format_for_conversion = "f32le"
    ffmpeg_command = [
        "ffmpeg",
        "-i",
        "pipe:0",
        "-ac",
        ac,
        "-ar",
        ar,
        "-f",
        format_for_conversion,
        "-hide_banner",
        "-loglevel",
        "quiet",
        "pipe:1",
    ]

    try:
        ffmpeg_process = subprocess.Popen(ffmpeg_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    except FileNotFoundError:
        raise ValueError("ffmpeg was not found but is required to load audio files from filename")
    output_stream = ffmpeg_process.communicate(bpayload)
    out_bytes = output_stream[0]

    audio = np.frombuffer(out_bytes, np.float32)
    if audio.shape[0] == 0:
        raise ValueError("Malformed soundfile")


SUMMARY:
**Code Summary**

### Overview

*   Reads an audio file using FFmpeg for conversion to NumPy arrays.
*   Parameters: `bpayload` (bytes input), `sampling_rate` (audio sampling rate).

### Function Signature
```python
def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
```

### Primary Purpose

Reads audio data from a file using FFmpeg, converting it to a NumPy array.

### Key Implementation Details

*   Uses `subprocess` to run an FFmpeg command with specified input and output streams.
*   Converting output to a NumPy array of floating-point values (`f32le` format).
*   Handles potential `ValueError` exceptions for malformed audio files or missing FFmpeg installation.

### Notable Dependencies

*   External dependency: `ffmpeg` (FFmpeg executable) must be installed on the system.
*   Related classes/functions:
    *   `subprocess`
    *   `np.array`, `np.float32`

### Return Value
```python
-> np.array
```

The function returns a NumPy array representing the audio data.


key: chroma:document
string_value:
CODE:

class ReturnType(enum.Enum):
    TENSORS = 0


SUMMARY:
**Summary:**

This `ReturnType` enum defines a set of named values used to specify the output type. It has two values:

* `TENSORS`: represents tensor data (value: 0)
* `TEXT`: represents text data (value: 1)

Notable details:

* Uses Python's built-in `enum.Enum` class.
* No external dependencies are listed, assuming it is self-contained within the module.


key: chroma:document
string_value:
CODE:

class TokenClassificationArgumentHandler(ArgumentHandler):
    """
    Handles arguments for token classification.
    """

    def __call__(self, inputs: Union[str, List[str]], **kwargs):
        if inputs is not None and isinstance(inputs, (list, tuple)) and len(inputs) > 0:
            inputs = list(inputs)
            batch_size = len(inputs)
        elif isinstance(inputs, str):
            inputs = [inputs]
            batch_size = 1
        elif Dataset is not None and isinstance(inputs, Dataset) or isinstance(inputs, types.GeneratorType):
            return inputs, None
        else:
            raise ValueError("At least one input is required.")

        offset_mapping = kwargs.get("offset_mapping")
        if offset_mapping:
            if isinstance(offset_mapping, list) and isinstance(offset_mapping[0], tuple):
                offset_mapping = [offset_mapping]
            if len(offset_mapping) != batch_size:
                raise ValueError("offset_mapping should have the same batch size as the input")


SUMMARY:
**TokenClassificationArgumentHandler Class Summary**

### Purpose and Functionality
The `TokenClassificationArgumentHandler` class is responsible for handling arguments for token classification tasks. It takes an input (string or list of strings) and provides the necessary configuration (`offset_mapping`) to process the data.

### Key Implementation Details

* **Input Processing**: The input is converted to a list, and its length is used to determine the batch size.
* **Configuration Validation**: The `offset_mapping` parameter is validated to ensure it matches the batch size of the input.

### Notable Dependencies and Requirements

* `ArgumentHandler`: Parent class
* `Dataset`: Optional parent class or generator type
* `types.GeneratorType`: Conditional check type
* `kwargs`: Configurations, including `offset_mapping`

**Notes:**

* The class returns two values: input data (`inputs`) and `offset_mapping` configuration.
* Validation checks are in place to ensure proper handling of input types.


key: chroma:document
string_value:
CODE:

class AggregationStrategy(ExplicitEnum):
    """All the valid aggregation strategies for TokenClassificationPipeline"""

    NONE = "none"
    SIMPLE = "simple"
    FIRST = "first"
    AVERAGE = "average"


SUMMARY:
**Summary:**

The `AggregationStrategy` class represents an enumeration of valid aggregation strategies for a `TokenClassificationPipeline`. It defines five distinct strategies:
	- `NONE`
	- `SIMPLE`
	- `FIRST`
	- `AVERAGE`
	- `MAX`

**Implementation Details:**

No key implementation details are provided; this is an explicit enumeration class.

**Notable Dependencies/Requirements:**

* No external dependencies or requirements listed.
* Interfaces: Implied by the class's use in a `TokenClassificationPipeline`.


key: file_path
string_value:
/home/gaby/projects/2025/transformers-4.48.1/src/transformers/pipelines/automatic_speech_recognition.py


key: chroma:document
string_value:
CODE:

def rescale_stride(stride, ratio):
    """
    Rescales the stride values from audio space to tokens/logits space.

    (160_000, 16_000, 16_000) -> (2000, 200, 200) for instance.
    """
    # Shape is [B, SEQ] for tokens
    # [B, SEQ, V] for logits

    new_strides = []
    for input_n, left, right in stride:
        token_n = int(round(input_n * ratio))
        left = int(round(left / input_n * token_n))
        right = int(round(right / input_n * token_n))
        new_stride = (token_n, left, right)
        new_strides.append(new_stride)



SUMMARY:
**Function Summary: `rescale_stride`**

**Purpose and Functionality**
Rescales stride values from audio space to tokens/logits space.

**Key Implementation Details**

* Calculates new stride dimensions based on input stride dimensions and scaling ratio (`ratio`)
* Uses integer arithmetic to round intermediate results
* Appends transformed strides to a list (`new_strides`)

**Notable Dependencies/Requirements**

* Input `stride`: expected shape `[B, SEQ]`, or `[B, SEQ, V]`
* Scaling ratio: `ratio` (no specific type or range indicated)
* External dependency: `int.round()` and arithmetic operations


key: chroma:document
string_value:
CODE:

def chunk_iter(inputs, feature_extractor, chunk_len, stride_left, stride_right, dtype=None):
    inputs_len = inputs.shape[0]
    step = chunk_len - stride_left - stride_right
    for chunk_start_idx in range(0, inputs_len, step):
        chunk_end_idx = chunk_start_idx + chunk_len
        chunk = inputs[chunk_start_idx:chunk_end_idx]
        processed = feature_extractor(chunk, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt")
        if dtype is not None:
            processed = processed.to(dtype=dtype)
        _stride_left = 0 if chunk_start_idx == 0 else stride_left
        is_last = chunk_end_idx >= inputs_len
        _stride_right = 0 if is_last else stride_right

        chunk_len = chunk.shape[0]
        stride = (chunk_len, _stride_left, _stride_right)
        if chunk.shape[0] > _stride_left:
            yield {"is_last": is_last, "stride": stride, **processed}
        if is_last:


SUMMARY:
**Function Summary:**

`chunk_iter`: An iterator function that processes input data in chunks based on user-defined parameters. It leverages a `feature_extractor` to extract features from each chunk and yields intermediate results.

**Key Details:**

* Inputs are processed in chunks of length `chunk_len`, with striding applied by `stride_left` and `stride_right`.
* The function extracts audio features using `feature_extractor` and returns tensors.
* It maintains an optional output data type `dtype`.
* Yields dictionaries containing feature values, along with stride information and processed tensors.

**Notable Dependencies/Requirements:**

* `feature_extractor`
* `sampling_rate` attribute (from `feature_extractor`)
* Output tensor shape and data type (`processed`)


key: chroma:document
string_value:
CODE:

def _fast_find_longest_common_sequence(sequence_left, sequence_right):
    seq_len_left = len(sequence_left)
    seq_len_right = len(sequence_right)
    counter = [[0] * (seq_len_right + 1) for _ in range(seq_len_left + 1)]
    longest = 0
    for i in range(seq_len_left):
        for j in range(seq_len_right):
            if sequence_left[i] == sequence_right[j]:
                previous_counter = counter[i][j] + 1
                counter[i + 1][j + 1] = previous_counter
                if previous_counter > longest:
                    longest = previous_counter

    counter = np.array(counter)
    # we return the idx of the first element of the longest common sequence in the left sequence
    index_left = np.argwhere(counter == longest)[-1][0] - longest if longest != 0 else -1
    index_right = np.argwhere(counter == longest)[-1][1] - longest if longest != 0 else -1


SUMMARY:
**Function Summary: `_fast_find_longest_common_sequence`**

This function finds the longest common subsequence (LCS) between two input sequences and returns its starting indices and length in both sequences.

**Purposes and Functionality:**

* Calculates the LCS using dynamic programming with a nested counter matrix.
* Compares characters at corresponding positions in both sequences to build the LCS.
* Returns:
	+ `index_left`: Index of the first element of the LCS in the `sequence_left`.
	+ `index_right`: Index of the first element of the LCS in the `sequence_right`.
	+ `longest`: Length of the longest common subsequence.

**Key Implementation Details:**

* Uses a 2D counter matrix (`counter`) to store intermediate results.
* Failsafe checks to handle cases where the input sequences are empty or have no matching characters.
* Utilizes NumPy's `argwhere` function to find the indices of the maximum value in the counter matrix.

**Notable Dependencies and Requirements:**

* Python 3.x compatibility
* NumPy library for efficient array operations
* No external dependencies on third-party libraries or APIs.


key: chroma:document
string_value:
CODE:

def _find_longest_common_sequence(sequences, tokenizer):
    # TODO  Use a faster algorithm this can probably be done in O(n)
    # using suffix array.
    # It might be tedious to do because of fault tolerance.
    # We actually have a really good property which is that the total sequence
    # MUST be those subsequences in order.
    # Also the algorithm should be more tolerant to errors.
    sequence = [tok_id for tok_id in sequences[0][0].tolist() if tok_id not in tokenizer.all_special_ids]
    for new_seq in sequences[1:]:
        new_sequence = [tok_id for tok_id in new_seq[0].tolist() if tok_id not in tokenizer.all_special_ids]

        index = 0
        max_ = 0.0
        for i in range(1, len(new_sequence) + 1):
            # epsilon to favor long perfect matches
            eps = i / 10000.0
            matches = np.sum(np.array(sequence[-i:]) == np.array(new_sequence[:i]))
            matching = matches / i + eps
            if matches > 1 and matching > max_:
                index = i
                max_ = matching
        sequence.extend(new_sequence[index:])


SUMMARY:
**Summary:**

### `_find_longest_common_sequence` Function 

#### Main Purpose & Functionality:

This function finds the longest common subsequence (LCS) among input sequences. It filters out special IDs from tokenized sequences and uses dynamic programming with epsilon optimization to measure match tolerance.

#### Key Implementation Details:

* **Sequence Filtering:** Removes special IDs from each sequence.
* **Dynamic Programming w/ Epsilon Optimizations:**
	+ Measures matches using Euclidean distance (`np.sum(np.array(sequence[-i:]) == np.array(new_sequence[:i]))`)
	+ Introduces epsilon for favorable long perfect matches (`matches / i + eps`)

#### Notable Dependencies & Requirements:

* `tokenizer`: instance with `all_special_ids`
* `numpy` (alias: `np`) for numerical computations

#### Method Signatures & Interfaces:

* None specific, but operates on sequences treated as numpy arrays.
* Extends sequence through LCP indexing.


key: chroma:document
string_value:
CODE:

class AutomaticSpeechRecognitionPipeline(ChunkPipeline):
    """
    Pipeline that aims at extracting spoken text contained within some audio.

    The input can be either a raw waveform or a audio file. In case of the audio file, ffmpeg should be installed for
    to support multiple audio formats

    Example:

    ```python
    >>> from transformers import pipeline

    >>> transcriber = pipeline(model="openai/whisper-base")
    >>> transcriber("https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/1.flac")
    {'text': ' He hoped there would be stew for dinner, turnips and carrots and bruised potatoes and fat mutton pieces to be ladled out in thick, peppered flour-fatten sauce.'}
    ```

    Learn more about the basics of using a pipeline in the [pipeline tutorial](../pipeline_tutorial)

    Arguments:
        model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
            The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from
            [`PreTrainedModel`] for PyTorch and [`TFPreTrainedModel`] for TensorFlow.
        feature_extractor ([`SequenceFeatureExtractor`]):
            The feature extractor that will be used by the pipeline to encode waveform for the model.
        tokenizer ([`PreTrainedTokenizer`]):
            The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from
            [`PreTrainedTokenizer`].
        decoder (`pyctcdecode.BeamSearchDecoderCTC`, *optional*):
            [PyCTCDecode's
            BeamSearchDecoderCTC](https://github.com/kensho-technologies/pyctcdecode/blob/2fd33dc37c4111417e08d89ccd23d28e9b308d19/pyctcdecode/decoder.py#L180)
            can be passed for language model boosted decoding. See [`Wav2Vec2ProcessorWithLM`] for more information.
        chunk_length_s (`float`, *optional*, defaults to 0):
            The input length for in each chunk. If `chunk_length_s = 0` then chunking is disabled (default).

            <Tip>

            For more information on how to effectively use `chunk_length_s`, please have a look at the [ASR chunking
            blog post](https://huggingface.co/blog/asr-chunking).

            </Tip>

        stride_length_s (`float`, *optional*, defaults to `chunk_length_s / 6`):
            The length of stride on the left and right of each chunk. Used only with `chunk_length_s > 0`. This enables
            the model to *see* more context and infer letters better than without this context but the pipeline
            discards the stride bits at the end to make the final reconstitution as perfect as possible.

            <Tip>

            For more information on how to effectively use `stride_length_s`, please have a look at the [ASR chunking
            blog post](https://huggingface.co/blog/asr-chunking).

            </Tip>

        framework (`str`, *optional*):
            The framework to use, either `"pt"` for PyTorch or `"tf"` for TensorFlow. The specified framework must be
            installed. If no framework is specified, will default to the one currently installed. If no framework is
            specified and both frameworks are installed, will default to the framework of the `model`, or to PyTorch if
            no model is provided.
        device (Union[`int`, `torch.device`], *optional*):
            Device ordinal for CPU/GPU supports. Setting this to `None` will leverage CPU, a positive will run the
            model on the associated CUDA device id.
        torch_dtype (Union[`int`, `torch.dtype`], *optional*):
            The data-type (dtype) of the computation. Setting this to `None` will use float32 precision. Set to
            `torch.float16` or `torch.bfloat16` to use half-precision in the respective dtypes.

    """

    def __init__(
        self,
        model: "PreTrainedModel",
        feature_extractor: Union["SequenceFeatureExtractor", str] = None,
        tokenizer: Optional[PreTrainedTokenizer] = None,
        decoder: Optional[Union["BeamSearchDecoderCTC", str]] = None,
        device: Union[int, "torch.device"] = None,
        torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
        **kwargs,
    ):
        # set the model type so we can check we have the right pre- and post-processing parameters
        if model.config.model_type == "whisper":
            self.type = "seq2seq_whisper"
        elif model.__class__.__name__ in MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES.values():
            self.type = "seq2seq"
        elif (
            feature_extractor._processor_class
            and feature_extractor._processor_class.endswith("WithLM")
            and decoder is not None
        ):
            self.decoder = decoder
            self.type = "ctc_with_lm"
        else:
            self.type = "ctc"

        super().__init__(model, tokenizer, feature_extractor, device=device, torch_dtype=torch_dtype, **kwargs)

    def __call__(
        self,
        inputs: Union[np.ndarray, bytes, str],
        **kwargs,
    ):
        """
        Transcribe the audio sequence(s) given as inputs to text. See the [`AutomaticSpeechRecognitionPipeline`]
        documentation for more information.

        Args:
            inputs (`np.ndarray` or `bytes` or `str` or `dict`):
                The inputs is either :
                    - `str` that is either the filename of a local audio file, or a public URL address to download the
                      audio file. The file will be read at the correct sampling rate to get the waveform using
                      *ffmpeg*. This requires *ffmpeg* to be installed on the system.
                    - `bytes` it is supposed to be the content of an audio file and is interpreted by *ffmpeg* in the
                      same way.
                    - (`np.ndarray` of shape (n, ) of type `np.float32` or `np.float64`)
                        Raw audio at the correct sampling rate (no further check will be done)
                    - `dict` form can be used to pass raw audio sampled at arbitrary `sampling_rate` and let this
                      pipeline do the resampling. The dict must be in the format `{"sampling_rate": int, "raw":
                      np.array}` with optionally a `"stride": (left: int, right: int)` than can ask the pipeline to
                      treat the first `left` samples and last `right` samples to be ignored in decoding (but used at
                      inference to provide more context to the model). Only use `stride` with CTC models.
            return_timestamps (*optional*, `str` or `bool`):
                Only available for pure CTC models (Wav2Vec2, HuBERT, etc) and the Whisper model. Not available for
                other sequence-to-sequence models.

                For CTC models, timestamps can take one of two formats:
                    - `"char"`: the pipeline will return timestamps along the text for every character in the text. For
                        instance, if you get `[{"text": "h", "timestamp": (0.5, 0.6)}, {"text": "i", "timestamp": (0.7,
                        0.9)}]`, then it means the model predicts that the letter "h" was spoken after `0.5` and before
                        `0.6` seconds.
                    - `"word"`: the pipeline will return timestamps along the text for every word in the text. For
                        instance, if you get `[{"text": "hi ", "timestamp": (0.5, 0.9)}, {"text": "there", "timestamp":
                        (1.0, 1.5)}]`, then it means the model predicts that the word "hi" was spoken after `0.5` and
                        before `0.9` seconds.

                For the Whisper model, timestamps can take one of two formats:
                    - `"word"`: same as above for word-level CTC timestamps. Word-level timestamps are predicted
                        through the *dynamic-time warping (DTW)* algorithm, an approximation to word-level timestamps
                        by inspecting the cross-attention weights.
                    - `True`: the pipeline will return timestamps along the text for *segments* of words in the text.
                        For instance, if you get `[{"text": " Hi there!", "timestamp": (0.5, 1.5)}]`, then it means the
                        model predicts that the segment "Hi there!" was spoken after `0.5` and before `1.5` seconds.
                        Note that a segment of text refers to a sequence of one or more words, rather than individual
                        words as with word-level timestamps.
            generate_kwargs (`dict`, *optional*):
                The dictionary of ad-hoc parametrization of `generate_config` to be used for the generation call. For a
                complete overview of generate, check the [following
                guide](https://huggingface.co/docs/transformers/en/main_classes/text_generation).

        Return:
            `Dict`: A dictionary with the following keys:
                - **text** (`str`): The recognized text.
                - **chunks** (*optional(, `List[Dict]`)
                    When using `return_timestamps`, the `chunks` will become a list containing all the various text
                    chunks identified by the model, *e.g.* `[{"text": "hi ", "timestamp": (0.5, 0.9)}, {"text":
                    "there", "timestamp": (1.0, 1.5)}]`. The original full text can roughly be recovered by doing
                    `"".join(chunk["text"] for chunk in output["chunks"])`.
        """
        return super().__call__(inputs, **kwargs)

    def _sanitize_parameters(
        self,
        chunk_length_s=None,
        stride_length_s=None,
        ignore_warning=None,
        decoder_kwargs=None,
        return_timestamps=None,
        return_language=None,
        generate_kwargs=None,
        max_new_tokens=None,
    ):
        # No parameters on this pipeline right now
        preprocess_params = {}
        if chunk_length_s is not None:
            if self.type == "seq2seq" and not ignore_warning:
                logger.warning(
                    "Using `chunk_length_s` is very experimental with seq2seq models. The results will not necessarily"
                    " be entirely accurate and will have caveats. More information:"
                    " https://github.com/huggingface/transformers/pull/20104. Ignore this warning with pipeline(...,"
                    " ignore_warning=True)"
                )
            preprocess_params["chunk_length_s"] = chunk_length_s
        if stride_length_s is not None:
            preprocess_params["stride_length_s"] = stride_length_s

        forward_params = defaultdict(dict)
        if max_new_tokens is not None:
            warnings.warn(
                "`max_new_tokens` is deprecated and will be removed in version 4.49 of Transformers. To remove this warning, pass `max_new_tokens` as a key inside `generate_kwargs` instead.",
                FutureWarning,
            )
            forward_params["max_new_tokens"] = max_new_tokens
        if generate_kwargs is not None:
            if max_new_tokens is not None and "max_new_tokens" in generate_kwargs:
                raise ValueError(
                    "`max_new_tokens` is defined both as an argument and inside `generate_kwargs` argument, please use"
                    " only 1 version"
                )
            forward_params.update(generate_kwargs)

        postprocess_params = {}
        if decoder_kwargs is not None:
            postprocess_params["decoder_kwargs"] = decoder_kwargs
        if return_timestamps is not None:
            # Check whether we have a valid setting for return_timestamps and throw an error before we perform a forward pass
            if self.type == "seq2seq" and return_timestamps:
                raise ValueError("We cannot return_timestamps yet on non-CTC models apart from Whisper!")
            if self.type == "ctc_with_lm" and return_timestamps != "word":
                raise ValueError("CTC with LM can only predict word level timestamps, set `return_timestamps='word'`")
            if self.type == "ctc" and return_timestamps not in ["char", "word"]:
                raise ValueError(
                    "CTC can either predict character level timestamps, or word level timestamps. "
                    "Set `return_timestamps='char'` or `return_timestamps='word'` as required."
                )
            if self.type == "seq2seq_whisper" and return_timestamps == "char":
                raise ValueError(
                    "Whisper cannot return `char` timestamps, only word level or segment level timestamps. "
                    "Use `return_timestamps='word'` or `return_timestamps=True` respectively."
                )
            forward_params["return_timestamps"] = return_timestamps
            postprocess_params["return_timestamps"] = return_timestamps
        if return_language is not None:
            if self.type != "seq2seq_whisper":
                raise ValueError("Only Whisper can return language for now.")
            postprocess_params["return_language"] = return_language

        if self.assistant_model is not None:
            forward_params["assistant_model"] = self.assistant_model
        if self.assistant_tokenizer is not None:
            forward_params["tokenizer"] = self.tokenizer
            forward_params["assistant_tokenizer"] = self.assistant_tokenizer

        return preprocess_params, forward_params, postprocess_params

    def preprocess(self, inputs, chunk_length_s=0, stride_length_s=None):
        if isinstance(inputs, str):
            if inputs.startswith("http://") or inputs.startswith("https://"):
                # We need to actually check for a real protocol, otherwise it's impossible to use a local file
                # like http_huggingface_co.png
                inputs = requests.get(inputs).content
            else:
                with open(inputs, "rb") as f:
                    inputs = f.read()

        if isinstance(inputs, bytes):
            inputs = ffmpeg_read(inputs, self.feature_extractor.sampling_rate)

        stride = None
        extra = {}
        if isinstance(inputs, dict):
            stride = inputs.pop("stride", None)
            # Accepting `"array"` which is the key defined in `datasets` for
            # better integration
            if not ("sampling_rate" in inputs and ("raw" in inputs or "array" in inputs)):
                raise ValueError(
                    "When passing a dictionary to AutomaticSpeechRecognitionPipeline, the dict needs to contain a "
                    '"raw" key containing the numpy array representing the audio and a "sampling_rate" key, '
                    "containing the sampling_rate associated with that array"
                )

            _inputs = inputs.pop("raw", None)
            if _inputs is None:
                # Remove path which will not be used from `datasets`.
                inputs.pop("path", None)
                _inputs = inputs.pop("array", None)
            in_sampling_rate = inputs.pop("sampling_rate")
            extra = inputs
            inputs = _inputs
            if in_sampling_rate != self.feature_extractor.sampling_rate:
                if is_torchaudio_available():
                    from torchaudio import functional as F
                else:
                    raise ImportError(
                        "torchaudio is required to resample audio samples in AutomaticSpeechRecognitionPipeline. "
                        "The torchaudio package can be installed through: `pip install torchaudio`."
                    )

                inputs = F.resample(
                    torch.from_numpy(inputs), in_sampling_rate, self.feature_extractor.sampling_rate
                ).numpy()
                ratio = self.feature_extractor.sampling_rate / in_sampling_rate
            else:
                ratio = 1
            if stride is not None:
                if stride[0] + stride[1] > inputs.shape[0]:
                    raise ValueError("Stride is too large for input")

                # Stride needs to get the chunk length here, it's going to get
                # swallowed by the `feature_extractor` later, and then batching
                # can add extra data in the inputs, so we need to keep track
                # of the original length in the stride so we can cut properly.
                stride = (inputs.shape[0], int(round(stride[0] * ratio)), int(round(stride[1] * ratio)))
        if not isinstance(inputs, np.ndarray):
            raise TypeError(f"We expect a numpy ndarray as input, got `{type(inputs)}`")
        if len(inputs.shape) != 1:
            raise ValueError("We expect a single channel audio input for AutomaticSpeechRecognitionPipeline")

        if chunk_length_s:
            if stride_length_s is None:
                stride_length_s = chunk_length_s / 6

            if isinstance(stride_length_s, (int, float)):
                stride_length_s = [stride_length_s, stride_length_s]

            # XXX: Carefuly, this variable will not exist in `seq2seq` setting.
            # Currently chunking is not possible at this level for `seq2seq` so
            # it's ok.
            align_to = getattr(self.model.config, "inputs_to_logits_ratio", 1)
            chunk_len = int(round(chunk_length_s * self.feature_extractor.sampling_rate / align_to) * align_to)
            stride_left = int(round(stride_length_s[0] * self.feature_extractor.sampling_rate / align_to) * align_to)
            stride_right = int(round(stride_length_s[1] * self.feature_extractor.sampling_rate / align_to) * align_to)

            if chunk_len < stride_left + stride_right:
                raise ValueError("Chunk length must be superior to stride length")

            for item in chunk_iter(
                inputs, self.feature_extractor, chunk_len, stride_left, stride_right, self.torch_dtype
            ):
                yield {**item, **extra}
        else:
            if self.type == "seq2seq_whisper" and inputs.shape[0] > self.feature_extractor.n_samples:
                processed = self.feature_extractor(
                    inputs,
                    sampling_rate=self.feature_extractor.sampling_rate,
                    truncation=False,
                    padding="longest",
                    return_tensors="pt",
                    return_attention_mask=True,
                )
            else:
                if self.type == "seq2seq_whisper" and stride is None:
                    processed = self.feature_extractor(
                        inputs,
                        sampling_rate=self.feature_extractor.sampling_rate,
                        return_tensors="pt",
                        return_token_timestamps=True,
                        return_attention_mask=True,
                    )
                    extra["num_frames"] = processed.pop("num_frames")
                else:
                    processed = self.feature_extractor(
                        inputs,
                        sampling_rate=self.feature_extractor.sampling_rate,
                        return_tensors="pt",
                        return_attention_mask=True,
                    )
            if self.torch_dtype is not None:
                processed = processed.to(dtype=self.torch_dtype)
            if stride is not None:
                if self.type == "seq2seq":
                    raise ValueError("Stride is only usable with CTC models, try removing it !")

                processed["stride"] = stride
            yield {"is_last": True, **processed, **extra}

    def _forward(self, model_inputs, return_timestamps=False, **generate_kwargs):
        attention_mask = model_inputs.pop("attention_mask", None)
        stride = model_inputs.pop("stride", None)
        num_frames = model_inputs.pop("num_frames", None)
        is_last = model_inputs.pop("is_last")

        if stride is not None and num_frames is not None:
            raise ValueError("num_frames must be used only when stride is None")

        if self.type in {"seq2seq", "seq2seq_whisper"}:
            # Consume values so we can let extra information flow freely through
            # the pipeline (important for `partial` in microphone)
            if "input_features" in model_inputs:
                inputs = model_inputs.pop("input_features")
            elif "input_values" in model_inputs:
                inputs = model_inputs.pop("input_values")
            else:
                raise ValueError(
                    "Seq2Seq speech recognition model requires either a "
                    f"`input_features` or `input_values` key, but only has {model_inputs.keys()}"
                )

            # custom processing for Whisper timestamps and word-level timestamps
            if return_timestamps and self.type == "seq2seq_whisper":
                generate_kwargs["return_timestamps"] = return_timestamps
                if return_timestamps == "word":
                    generate_kwargs["return_token_timestamps"] = True
                    generate_kwargs["return_segments"] = True

                    if stride is not None:
                        if isinstance(stride, tuple):
                            generate_kwargs["num_frames"] = stride[0] // self.feature_extractor.hop_length
                        else:
                            generate_kwargs["num_frames"] = [s[0] // self.feature_extractor.hop_length for s in stride]
                    else:
                        generate_kwargs["num_frames"] = num_frames

            # User-defined `generation_config` passed to the pipeline call take precedence
            if "generation_config" not in generate_kwargs:
                generate_kwargs["generation_config"] = self.generation_config

            tokens = self.model.generate(
                inputs=inputs,
                attention_mask=attention_mask,
                **generate_kwargs,
            )
            # whisper longform generation stores timestamps in "segments"
            if return_timestamps == "word" and self.type == "seq2seq_whisper":
                if "segments" not in tokens:
                    out = {"tokens": tokens["sequences"], "token_timestamps": tokens["token_timestamps"]}
                else:
                    token_timestamps = [
                        torch.cat([segment["token_timestamps"] for segment in segment_list])
                        for segment_list in tokens["segments"]
                    ]
                    out = {"tokens": tokens["sequences"], "token_timestamps": token_timestamps}
            else:
                out = {"tokens": tokens}
            if self.type == "seq2seq_whisper":
                if stride is not None:
                    out["stride"] = stride

        else:
            inputs = {
                self.model.main_input_name: model_inputs.pop(self.model.main_input_name),
                "attention_mask": attention_mask,
            }
            outputs = self.model(**inputs)
            logits = outputs.logits

            if self.type == "ctc_with_lm":
                out = {"logits": logits}
            else:
                out = {"tokens": logits.argmax(dim=-1)}
            if stride is not None:
                # Send stride to `postprocess`.
                # it needs to be handled there where
                # the pieces are to be concatenated.
                ratio = 1 / self.model.config.inputs_to_logits_ratio
                if isinstance(stride, tuple):
                    out["stride"] = rescale_stride([stride], ratio)[0]
                else:
                    out["stride"] = rescale_stride(stride, ratio)
        # Leftover
        extra = model_inputs
        return {"is_last": is_last, **out, **extra}

    def postprocess(
        self, model_outputs, decoder_kwargs: Optional[Dict] = None, return_timestamps=None, return_language=None
    ):
        # Optional return types
        optional = {}

        final_items = []
        key = "logits" if self.type == "ctc_with_lm" else "tokens"
        stride = None
        for outputs in model_outputs:
            if self.framework == "pt" and outputs[key].dtype in (torch.bfloat16, torch.float16):
                items = outputs[key].to(torch.float32).numpy()
            else:
                items = outputs[key].numpy()
            stride = outputs.get("stride", None)
            if stride is not None and self.type in {"ctc", "ctc_with_lm"}:
                total_n, left, right = stride
                # Total_n might be < logits.shape[1]
                # because of padding, that's why
                # we need to reconstruct this information
                # This won't work with left padding (which doesn't exist right now)
                right_n = total_n - right
                items = items[:, left:right_n]
            final_items.append(items)

        if stride and self.type == "seq2seq":
            items = _find_longest_common_sequence(final_items, self.tokenizer)
        elif self.type == "seq2seq_whisper":
            time_precision = self.feature_extractor.chunk_length / self.model.config.max_source_positions
            # Send the chunking back to seconds, it's easier to handle in whisper
            sampling_rate = self.feature_extractor.sampling_rate
            for output in model_outputs:
                if "stride" in output:
                    chunk_len, stride_left, stride_right = output["stride"]
                    # Go back in seconds
                    chunk_len /= sampling_rate
                    stride_left /= sampling_rate
                    stride_right /= sampling_rate
                    output["stride"] = chunk_len, stride_left, stride_right

            text, optional = self.tokenizer._decode_asr(
                model_outputs,
                return_timestamps=return_timestamps,
                return_language=return_language,
                time_precision=time_precision,
            )
        else:
            items = np.concatenate(final_items, axis=1)
            items = items.squeeze(0)

        if self.type == "ctc_with_lm":
            if decoder_kwargs is None:
                decoder_kwargs = {}
            beams = self.decoder.decode_beams(items, **decoder_kwargs)
            text = beams[0][0]
            if return_timestamps:
                # Simply cast from pyctcdecode format to wav2vec2 format to leverage
                # pre-existing code later
                chunk_offset = beams[0][2]
                offsets = []
                for word, (start_offset, end_offset) in chunk_offset:
                    offsets.append({"word": word, "start_offset": start_offset, "end_offset": end_offset})
        elif self.type != "seq2seq_whisper":
            skip_special_tokens = self.type != "ctc"
            text = self.tokenizer.decode(items, skip_special_tokens=skip_special_tokens)
            if return_timestamps:
                offsets = self.tokenizer.decode(
                    items, skip_special_tokens=skip_special_tokens, output_char_offsets=True
                )["char_offsets"]
                if return_timestamps == "word":
                    offsets = self.tokenizer._get_word_offsets(offsets, self.tokenizer.replace_word_delimiter_char)

        if return_timestamps and self.type not in {"seq2seq", "seq2seq_whisper"}:
            chunks = []
            for item in offsets:
                start = item["start_offset"] * self.model.config.inputs_to_logits_ratio
                start /= self.feature_extractor.sampling_rate

                stop = item["end_offset"] * self.model.config.inputs_to_logits_ratio
                stop /= self.feature_extractor.sampling_rate

                chunks.append({"text": item[return_timestamps], "timestamp": (start, stop)})
            optional["chunks"] = chunks

        extra = defaultdict(list)
        for output in model_outputs:
            output.pop("tokens", None)
            output.pop("logits", None)
            output.pop("is_last", None)
            output.pop("stride", None)
            output.pop("token_timestamps", None)
            for k, v in output.items():
                extra[k].append(v)


SUMMARY:
This is a Python class implementation of an Automatic Speech Recognition (ASR) system using the Kaldi framework. It appears to be part of a larger library or application for speech recognition and post-processing.

The class has several methods:

1. `__init__`: Initializes the ASR system with various parameters, such as the model configuration, input streams, and output types.
2. `run_transcription`: Performs ASR transcription on a given audio input using a pre-trained model or a custom model.
3. `postprocess`: Post-processes the ASR output by decoding the transcriptions, handling timestamps and language information, and reformatting the output.

The class seems to be designed for flexibility and customization, with options to use various models, decoding algorithms, and output formats.

Some notable features of this implementation include:

* Integration with the Kaldi framework for speech recognition
* Support for multiple ASR models, including pre-trained and custom models
* Flexible output formats, including text-based transcripts and timestamped data
* Customization options for decoding algorithms and language handling
* Error detection and handling mechanisms for post-processing errors


key: chroma:document
string_value:
CODE:

def _find_timestamp_sequence(sequences, tokenizer, feature_extractor, max_source_positions):
    """
    Computes the final sequences by merging the end of the nth sequence with the beginning of the n+1th sequence. Since
    `WhisperForConditionalGeneration` produces the timestamps pairwise, we filter the consecutive timestamps and only
    iterate over them. We keep track of the `time` which indicates the actual starting time of the chunk that is
    processed. We need to make sure to offset the timestamps tokens by the `time` in order for the tokenizer to
    properly compute the final `offset`.
    """
    # index of the first timestamp token
    timestamp_begin = tokenizer.convert_tokens_to_ids("<|notimestamps|>") + 1
    items = []
    # approximation of the token to time ratio : ~0.2seconds
    time_precision = feature_extractor.chunk_length / max_source_positions
    time = 0
    for seq_idx, item in enumerate(sequences):
        sequence, stride = item
        if isinstance(sequence, list):
            sequence = np.array(sequence)
        chunk_len, stride_left, stride_right = stride
        sequence = sequence.squeeze(0)
        # get rid of the `forced_decoder_idx` that are use to parametrize the generation
        begin_idx = np.where(sequence == timestamp_begin)[0][0] if timestamp_begin in sequence else 0
        sequence = sequence[begin_idx:]

        timestamp_tokens = sequence >= timestamp_begin
        if seq_idx != 0 and sum(timestamp_tokens) > 0:
            consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1
            last_timestamp = np.where(timestamp_tokens)[0][-1]
            consecutive = np.append(consecutive, last_timestamp) if last_timestamp not in consecutive else consecutive
            time -= stride_left + stride_right
            offset = int((time / feature_extractor.sampling_rate) / time_precision)
            overlap_time = int((stride_left / feature_extractor.sampling_rate) / time_precision)
            # relevant timestamps are in the overlapping part
            relevant_timestamp = np.where(sequence[consecutive] >= timestamp_begin + overlap_time)[0]
            if relevant_timestamp.shape[0] > 0:
                relevant_timestamp = (
                    consecutive[relevant_timestamp[0] - 1] if relevant_timestamp[0] > 0 else consecutive[0]
                )
                # if a big stride is used, we need to check some of the previous items for the best overlap
                best_match = 0
                sliced_sequence = []
                for idx, previous_sequence in enumerate(reversed(items)):
                    previous_tokens = previous_sequence[1:-1]
                    if previous_sequence[0] < (timestamp_begin + offset - overlap_time) and idx != 0:
                        break  # the previous sequence is too far in the past
                    if len(previous_tokens) > 0:
                        # find the longest common sequence between the overlapping parts
                        index_left, index_right, match_length = _fast_find_longest_common_sequence(
                            sequence[1:relevant_timestamp], previous_tokens
                        )
                        # don't do anything if only 1 token was matched
                        if match_length > 1 and match_length > best_match:
                            best_match = match_length
                            best_idx = idx
                            end_of_curr_sequence_idx = (
                                np.where(sequence[index_left + 1 :] >= timestamp_begin)[0][0] + 1
                            )
                            end_of_curr_sequence_idx = end_of_curr_sequence_idx + 1 + index_left
                            # if all the tokens are matched, suffix
                            if index_left == 0 and match_length == len(previous_tokens):
                                sliced_sequence = np.insert(
                                    sequence[index_left + 1 : end_of_curr_sequence_idx], 0, previous_sequence[0]
                                )
                                sliced_sequence[-1] = previous_sequence[-1]
                            # if part of the previous sequence is not taken
                            elif index_left >= 0:
                                sliced_sequence = sequence[index_left + 1 : end_of_curr_sequence_idx]
                                # let's insert the missing part of the previous sequence
                                previous_slice = (
                                    previous_sequence[: index_right + 1] if index_right > 0 else [previous_sequence[0]]
                                )
                                sliced_sequence = np.insert(sliced_sequence, 0, previous_slice)
                                sliced_sequence[-1] += offset

                if len(sliced_sequence) > 0:
                    items[len(items) - best_idx - 1] = sliced_sequence
                    items = items[: len(items) - best_idx]
                    sequence = sequence[end_of_curr_sequence_idx:]

        # sequence might have changed
        timestamp_tokens = sequence >= timestamp_begin
        consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1
        if sum(timestamp_tokens) > 0:
            last_timestamp = np.where(timestamp_tokens)[0][-1]
            consecutive = (
                np.append(consecutive, last_timestamp + 1) if last_timestamp not in consecutive else consecutive
            )

        if len(consecutive) > 0:
            last_slice = 0
            for current_slice in consecutive:
                actual_offset = items[-1][-1] if seq_idx != 0 or last_slice != 0 else sequence[0]
                sliced_tokens = sequence[last_slice:current_slice]
                duration = sliced_tokens[-1] - sliced_tokens[0]
                sliced_tokens[0] = actual_offset
                sliced_tokens[-1] = actual_offset + duration
                items.append(sliced_tokens)
                last_slice = current_slice

        time += chunk_len
    result = []
    for i in range(len(items)):
        result += items[i].tolist()


SUMMARY:
**Function Summary: `_find_timestamp_sequence`**

**Main Purpose:** Merges timestamp sequences by finding overlapping parts between consecutive timestamps.

**Primary Functionality:**

1. Computes final sequences by merging end of nth sequence with beginning of n+1th sequence, offsetting tokens by time.
2. Filters out non-consecutive timestamps and finds overlapping parts.
3. Finds longest common sequences between overlapping parts using `_fast_find_longest_common_sequence`.
4. Inserts matching token to overlap and adjusts token offsets accordingly.

**Key Design Choices:**

* Uses `concurrent` feature extraction for timestamp sequence merging.
* Inverts `timestamp_begin` index value for correct lookup.
* Utilizes `time_precision` approximation when finding overlapping parts.

**Dependences/Requirements:**

* `tokenizer`
* `feature_extractor` class with:
	+ `sampling_rate`
	+ `chunk_length`
* `_fast_find_longest_common_sequence` function (not provided in the code snippet).

**Data Structures and Variables:**

* `sequences`: input list of timestamp sequences.
* `items`: output list of merged timestamp sequences.
* `timestamp_begin`: index value for converting timestamps to IDs.
* `time`: current time offset.

Note that this summary captures the main purpose, key design choices, dependencies, and data structures used in the function, while omitting specific implementation details that may not add value to understanding the overall functionality.


key: chroma:document
string_value:
CODE:

def get_supported_tasks() -> List[str]:
    """
    Returns a list of supported task strings.
    """


SUMMARY:
**Function Summary**

*   `get_supported_tasks()` Function
    *   Purpose: Retrieves and returns a list of supported tasks as task strings from the `pipeline_registry` using the `get_supported_tasks()` method on it.
*   Implementation Details:
    *   The function utilizes the `PIPELINE_REGISTRY` class to retrieve a list of supported tasks by calling its `get_supported_tasks()` method.
        +   Interface: `get_supported_tasks() -> List[str]`
        +   Method: `get_supported_tasks()`
*   Dependencies:
    *   **External Dependency:** None
    *   **Class/Function Reference:** `PIPELINE_REGISTRY` (a class that provides registry functionality)

**Summary in RAG Format:** 

R|---`get_supported_tasks`: Retrieves and returns a list of supported tasks as task strings from the `pipeline_registry`.|
A|  Implementation Details: Uses `pipeline_REGISTRY.get_supported_tasks()` to retrieve supported tasks.|
I|  Dependencies: `PIPELINE Registry` class |
D|  Signature: `get_supported_tasks() -> List[str]`


key: chroma:document
string_value:
CODE:

def get_task(model: str, token: Optional[str] = None, **deprecated_kwargs) -> str:
    use_auth_token = deprecated_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 is_offline_mode():
        raise RuntimeError("You cannot infer task automatically within `pipeline` when using offline mode")
    try:
        info = model_info(model, token=token)
    except Exception as e:
        raise RuntimeError(f"Instantiating a pipeline without a task set raised an error: {e}")
    if not info.pipeline_tag:
        raise RuntimeError(
            f"The model {model} does not seem to have a correct `pipeline_tag` set to infer the task automatically"
        )
    if getattr(info, "library_name", "transformers") not in {"transformers", "timm"}:
        raise RuntimeError(f"This model is meant to be used with {info.library_name} not with transformers")
    task = info.pipeline_tag


SUMMARY:
**Concise Summary:**

The `get_task` function retrieves and returns the pipeline tag for a given model. It:

1. Validates input parameters, specifically handling deprecated arguments and token usage.
2. Checks if the model is running in offline mode and raises an error if so.
3. Tries to instantiate a pipeline without manually setting the task, and catches any exceptions.
4. Verifies that the library name matches "transformers" or "timm", raising an error otherwise.
5. Returns the pipeline tag, which is assumed to be set in the model's configuration.

**Key Implementation Details:**

* Handles deprecated `use_auth_token` argument and warns about potential removal.
* Checks for offline mode and raises error if applicable.
* Verifies library name consistency before proceeding.

**Notable Dependencies/Requirements:**

* `token`: an optional string parameter used as authentication token.
* `deprecated_kwargs`: a dictionary containing deprecated arguments (e.g., `use_auth_token`).
* `model_info`, `is_offline_mode`, and `warnings` modules are not explicitly imported, but implied for their usage.


key: chroma:document
string_value:
CODE:

def check_task(task: str) -> Tuple[str, Dict, Any]:
    """
    Checks an incoming task string, to validate it's correct and return the default Pipeline and Model classes, and
    default models if they exist.

    Args:
        task (`str`):
            The task defining which pipeline will be returned. Currently accepted tasks are:

            - `"audio-classification"`
            - `"automatic-speech-recognition"`
            - `"conversational"`
            - `"depth-estimation"`
            - `"document-question-answering"`
            - `"feature-extraction"`
            - `"fill-mask"`
            - `"image-classification"`
            - `"image-feature-extraction"`
            - `"image-segmentation"`
            - `"image-to-text"`
            - `"image-to-image"`
            - `"object-detection"`
            - `"question-answering"`
            - `"summarization"`
            - `"table-question-answering"`
            - `"text2text-generation"`
            - `"text-classification"` (alias `"sentiment-analysis"` available)
            - `"text-generation"`
            - `"text-to-audio"` (alias `"text-to-speech"` available)
            - `"token-classification"` (alias `"ner"` available)
            - `"translation"`
            - `"translation_xx_to_yy"`
            - `"video-classification"`
            - `"visual-question-answering"` (alias `"vqa"` available)
            - `"zero-shot-classification"`
            - `"zero-shot-image-classification"`
            - `"zero-shot-object-detection"`

    Returns:
        (normalized_task: `str`, task_defaults: `dict`, task_options: (`tuple`, None)) The normalized task name
        (removed alias and options). The actual dictionary required to initialize the pipeline and some extra task
        options for parametrized tasks like "translation_XX_to_YY"


    """


SUMMARY:
**Function Summary: `check_task`**

**Main Purpose:** Validates an incoming task string against predefined tasks and returns the default Pipeline and Model classes along with extra options for parametrized tasks.

**Key Implementation Details:**

* The function uses a predefined `PIPELINE_REGISTRY` object to perform validation.
* It checks if the input `task` is one of the accepted tasks and normalizes it by removing aliases.
* Returns a tuple containing:
	+ `normalized_task`: the normalized task name
	+ `task_defaults`: a dictionary required to initialize the pipeline
	+ `task_options`: an optional string with extra options for parametrized tasks

**Notable Dependencies/Requirements:**

* `PIPELINE_REGISTRY` class, not defined in this snippet
* `Pipeline`, `Model`, and other related classes
* External dependencies not listed, assuming they are included in the `PIPELINE_REGISTRY` or elsewhere.


key: chroma:document
string_value:
CODE:

def clean_custom_task(task_info):
    import transformers

    if "impl" not in task_info:
        raise RuntimeError("This model introduces a custom pipeline without specifying its implementation.")
    pt_class_names = task_info.get("pt", ())
    if isinstance(pt_class_names, str):
        pt_class_names = [pt_class_names]
    task_info["pt"] = tuple(getattr(transformers, c) for c in pt_class_names)
    tf_class_names = task_info.get("tf", ())
    if isinstance(tf_class_names, str):
        tf_class_names = [tf_class_names]
    task_info["tf"] = tuple(getattr(transformers, c) for c in tf_class_names)


SUMMARY:
**Function Summary:**

`clean_custom_task` function cleans and normalizes custom task information for use with the Hugging Face Transformers library.

**Main Purpose and Functionality:**

Checks if a custom pipeline implementation is specified and updates the `task_info` dictionary to include necessary transformers classes.

**Key Implementation Details:**

* Handles cases where pipeline names are passed as strings
* Uses `getattr(transformers, c)` to dynamically import transformer classes
* Returns updated `task_info` with loaded transformer classes

**Notable Dependencies or Requirements:**

* Hugging Face Transformers library (`transformers`)
* Python runtime environment


key: chroma:document
string_value:
CODE:

def pipeline(
    task: str = None,
    model: Optional[Union[str, "PreTrainedModel", "TFPreTrainedModel"]] = None,
    config: Optional[Union[str, PretrainedConfig]] = None,
    tokenizer: Optional[Union[str, PreTrainedTokenizer, "PreTrainedTokenizerFast"]] = None,
    feature_extractor: Optional[Union[str, PreTrainedFeatureExtractor]] = None,
    image_processor: Optional[Union[str, BaseImageProcessor]] = None,
    processor: Optional[Union[str, ProcessorMixin]] = None,
    framework: Optional[str] = None,
    revision: Optional[str] = None,
    use_fast: bool = True,
    token: Optional[Union[str, bool]] = None,
    device: Optional[Union[int, str, "torch.device"]] = None,
    device_map=None,
    torch_dtype=None,
    trust_remote_code: Optional[bool] = None,
    model_kwargs: Dict[str, Any] = None,
    pipeline_class: Optional[Any] = None,
    **kwargs,
) -> Pipeline:
    """
    Utility factory method to build a [`Pipeline`].

    A pipeline consists of:

        - One or more components for pre-processing model inputs, such as a [tokenizer](tokenizer),
        [image_processor](image_processor), [feature_extractor](feature_extractor), or [processor](processors).
        - A [model](model) that generates predictions from the inputs.
        - Optional post-processing steps to refine the model's output, which can also be handled by processors.

    <Tip>
    While there are such optional arguments as `tokenizer`, `feature_extractor`, `image_processor`, and `processor`,
    they shouldn't be specified all at once. If these components are not provided, `pipeline` will try to load
    required ones automatically. In case you want to provide these components explicitly, please refer to a
    specific pipeline in order to get more details regarding what components are required.
    </Tip>

    Args:
        task (`str`):
            The task defining which pipeline will be returned. Currently accepted tasks are:

            - `"audio-classification"`: will return a [`AudioClassificationPipeline`].
            - `"automatic-speech-recognition"`: will return a [`AutomaticSpeechRecognitionPipeline`].
            - `"depth-estimation"`: will return a [`DepthEstimationPipeline`].
            - `"document-question-answering"`: will return a [`DocumentQuestionAnsweringPipeline`].
            - `"feature-extraction"`: will return a [`FeatureExtractionPipeline`].
            - `"fill-mask"`: will return a [`FillMaskPipeline`]:.
            - `"image-classification"`: will return a [`ImageClassificationPipeline`].
            - `"image-feature-extraction"`: will return an [`ImageFeatureExtractionPipeline`].
            - `"image-segmentation"`: will return a [`ImageSegmentationPipeline`].
            - `"image-text-to-text"`: will return a [`ImageTextToTextPipeline`].
            - `"image-to-image"`: will return a [`ImageToImagePipeline`].
            - `"image-to-text"`: will return a [`ImageToTextPipeline`].
            - `"mask-generation"`: will return a [`MaskGenerationPipeline`].
            - `"object-detection"`: will return a [`ObjectDetectionPipeline`].
            - `"question-answering"`: will return a [`QuestionAnsweringPipeline`].
            - `"summarization"`: will return a [`SummarizationPipeline`].
            - `"table-question-answering"`: will return a [`TableQuestionAnsweringPipeline`].
            - `"text2text-generation"`: will return a [`Text2TextGenerationPipeline`].
            - `"text-classification"` (alias `"sentiment-analysis"` available): will return a
              [`TextClassificationPipeline`].
            - `"text-generation"`: will return a [`TextGenerationPipeline`]:.
            - `"text-to-audio"` (alias `"text-to-speech"` available): will return a [`TextToAudioPipeline`]:.
            - `"token-classification"` (alias `"ner"` available): will return a [`TokenClassificationPipeline`].
            - `"translation"`: will return a [`TranslationPipeline`].
            - `"translation_xx_to_yy"`: will return a [`TranslationPipeline`].
            - `"video-classification"`: will return a [`VideoClassificationPipeline`].
            - `"visual-question-answering"`: will return a [`VisualQuestionAnsweringPipeline`].
            - `"zero-shot-classification"`: will return a [`ZeroShotClassificationPipeline`].
            - `"zero-shot-image-classification"`: will return a [`ZeroShotImageClassificationPipeline`].
            - `"zero-shot-audio-classification"`: will return a [`ZeroShotAudioClassificationPipeline`].
            - `"zero-shot-object-detection"`: will return a [`ZeroShotObjectDetectionPipeline`].

        model (`str` or [`PreTrainedModel`] or [`TFPreTrainedModel`], *optional*):
            The model that will be used by the pipeline to make predictions. This can be a model identifier or an
            actual instance of a pretrained model inheriting from [`PreTrainedModel`] (for PyTorch) or
            [`TFPreTrainedModel`] (for TensorFlow).

            If not provided, the default for the `task` will be loaded.
        config (`str` or [`PretrainedConfig`], *optional*):
            The configuration that will be used by the pipeline to instantiate the model. This can be a model
            identifier or an actual pretrained model configuration inheriting from [`PretrainedConfig`].

            If not provided, the default configuration file for the requested model will be used. That means that if
            `model` is given, its default configuration will be used. However, if `model` is not supplied, this
            `task`'s default model's config is used instead.
        tokenizer (`str` or [`PreTrainedTokenizer`], *optional*):
            The tokenizer that will be used by the pipeline to encode data for the model. This can be a model
            identifier or an actual pretrained tokenizer inheriting from [`PreTrainedTokenizer`].

            If not provided, the default tokenizer for the given `model` will be loaded (if it is a string). If `model`
            is not specified or not a string, then the default tokenizer for `config` is loaded (if it is a string).
            However, if `config` is also not given or not a string, then the default tokenizer for the given `task`
            will be loaded.
        feature_extractor (`str` or [`PreTrainedFeatureExtractor`], *optional*):
            The feature extractor that will be used by the pipeline to encode data for the model. This can be a model
            identifier or an actual pretrained feature extractor inheriting from [`PreTrainedFeatureExtractor`].

            Feature extractors are used for non-NLP models, such as Speech or Vision models as well as multi-modal
            models. Multi-modal models will also require a tokenizer to be passed.

            If not provided, the default feature extractor for the given `model` will be loaded (if it is a string). If
            `model` is not specified or not a string, then the default feature extractor for `config` is loaded (if it
            is a string). However, if `config` is also not given or not a string, then the default feature extractor
            for the given `task` will be loaded.
        image_processor (`str` or [`BaseImageProcessor`], *optional*):
            The image processor that will be used by the pipeline to preprocess images for the model. This can be a
            model identifier or an actual image processor inheriting from [`BaseImageProcessor`].

            Image processors are used for Vision models and multi-modal models that require image inputs. Multi-modal
            models will also require a tokenizer to be passed.

            If not provided, the default image processor for the given `model` will be loaded (if it is a string). If
            `model` is not specified or not a string, then the default image processor for `config` is loaded (if it is
            a string).
        processor (`str` or [`ProcessorMixin`], *optional*):
            The processor that will be used by the pipeline to preprocess data for the model. This can be a model
            identifier or an actual processor inheriting from [`ProcessorMixin`].

            Processors are used for multi-modal models that require multi-modal inputs, for example, a model that
            requires both text and image inputs.

            If not provided, the default processor for the given `model` will be loaded (if it is a string). If `model`
            is not specified or not a string, then the default processor for `config` is loaded (if it is a string).
        framework (`str`, *optional*):
            The framework to use, either `"pt"` for PyTorch or `"tf"` for TensorFlow. The specified framework must be
            installed.

            If no framework is specified, will default to the one currently installed. If no framework is specified and
            both frameworks are installed, will default to the framework of the `model`, or to PyTorch if no model is
            provided.
        revision (`str`, *optional*, defaults to `"main"`):
            When passing a task name or a string model identifier: 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.
        use_fast (`bool`, *optional*, defaults to `True`):
            Whether or not to use a Fast tokenizer if possible (a [`PreTrainedTokenizerFast`]).
        use_auth_token (`str` or *bool*, *optional*):
            The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
            when running `huggingface-cli login` (stored in `~/.huggingface`).
        device (`int` or `str` or `torch.device`):
            Defines the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank like `1`) on which this
            pipeline will be allocated.
        device_map (`str` or `Dict[str, Union[int, str, torch.device]`, *optional*):
            Sent directly as `model_kwargs` (just a simpler shortcut). When `accelerate` library is present, set
            `device_map="auto"` to compute the most optimized `device_map` automatically (see
            [here](https://huggingface.co/docs/accelerate/main/en/package_reference/big_modeling#accelerate.cpu_offload)
            for more information).

            <Tip warning={true}>

            Do not use `device_map` AND `device` at the same time as they will conflict

            </Tip>

        torch_dtype (`str` or `torch.dtype`, *optional*):
            Sent directly as `model_kwargs` (just a simpler shortcut) to use the available precision for this model
            (`torch.float16`, `torch.bfloat16`, ... or `"auto"`).
        trust_remote_code (`bool`, *optional*, defaults to `False`):
            Whether or not to allow for custom code defined on the Hub in their own modeling, configuration,
            tokenization or even pipeline files. This option should only be set to `True` for repositories you trust
            and in which you have read the code, as it will execute code present on the Hub on your local machine.
        model_kwargs (`Dict[str, Any]`, *optional*):
            Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,
            **model_kwargs)` function.
        kwargs (`Dict[str, Any]`, *optional*):
            Additional keyword arguments passed along to the specific pipeline init (see the documentation for the
            corresponding pipeline class for possible values).

    Returns:
        [`Pipeline`]: A suitable pipeline for the task.

    Examples:

    ```python
    >>> from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer

    >>> # Sentiment analysis pipeline
    >>> analyzer = pipeline("sentiment-analysis")

    >>> # Question answering pipeline, specifying the checkpoint identifier
    >>> oracle = pipeline(
    ...     "question-answering", model="distilbert/distilbert-base-cased-distilled-squad", tokenizer="google-bert/bert-base-cased"
    ... )

    >>> # Named entity recognition pipeline, passing in a specific model and tokenizer
    >>> model = AutoModelForTokenClassification.from_pretrained("dbmdz/bert-large-cased-finetuned-conll03-english")
    >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")
    >>> recognizer = pipeline("ner", model=model, tokenizer=tokenizer)
    ```"""
    if model_kwargs is None:
        model_kwargs = {}
    # Make sure we only pass use_auth_token once as a kwarg (it used to be possible to pass it in model_kwargs,
    # this is to keep BC).
    use_auth_token = model_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

    code_revision = kwargs.pop("code_revision", None)
    commit_hash = kwargs.pop("_commit_hash", None)

    hub_kwargs = {
        "revision": revision,
        "token": token,
        "trust_remote_code": trust_remote_code,
        "_commit_hash": commit_hash,
    }

    if task is None and model is None:
        raise RuntimeError(
            "Impossible to instantiate a pipeline without either a task or a model "
            "being specified. "
            "Please provide a task class or a model"
        )

    if model is None and tokenizer is not None:
        raise RuntimeError(
            "Impossible to instantiate a pipeline with tokenizer specified but not the model as the provided tokenizer"
            " may not be compatible with the default model. Please provide a PreTrainedModel class or a"
            " path/identifier to a pretrained model when providing tokenizer."
        )
    if model is None and feature_extractor is not None:
        raise RuntimeError(
            "Impossible to instantiate a pipeline with feature_extractor specified but not the model as the provided"
            " feature_extractor may not be compatible with the default model. Please provide a PreTrainedModel class"
            " or a path/identifier to a pretrained model when providing feature_extractor."
        )
    if isinstance(model, Path):
        model = str(model)

    if commit_hash is None:
        pretrained_model_name_or_path = None
        if isinstance(config, str):
            pretrained_model_name_or_path = config
        elif config is None and isinstance(model, str):
            pretrained_model_name_or_path = model

        if not isinstance(config, PretrainedConfig) and pretrained_model_name_or_path is not None:
            # We make a call to the config file first (which may be absent) to get the commit hash as soon as possible
            resolved_config_file = cached_file(
                pretrained_model_name_or_path,
                CONFIG_NAME,
                _raise_exceptions_for_gated_repo=False,
                _raise_exceptions_for_missing_entries=False,
                _raise_exceptions_for_connection_errors=False,
                cache_dir=model_kwargs.get("cache_dir"),
                **hub_kwargs,
            )
            hub_kwargs["_commit_hash"] = extract_commit_hash(resolved_config_file, commit_hash)
        else:
            hub_kwargs["_commit_hash"] = getattr(config, "_commit_hash", None)

    # Config is the primordial information item.
    # Instantiate config if needed
    if isinstance(config, str):
        config = AutoConfig.from_pretrained(
            config, _from_pipeline=task, code_revision=code_revision, **hub_kwargs, **model_kwargs
        )
        hub_kwargs["_commit_hash"] = config._commit_hash
    elif config is None and isinstance(model, str):
        # Check for an adapter file in the model path if PEFT is available
        if is_peft_available():
            # `find_adapter_config_file` doesn't accept `trust_remote_code`
            _hub_kwargs = {k: v for k, v in hub_kwargs.items() if k != "trust_remote_code"}
            maybe_adapter_path = find_adapter_config_file(
                model,
                token=hub_kwargs["token"],
                revision=hub_kwargs["revision"],
                _commit_hash=hub_kwargs["_commit_hash"],
            )

            if maybe_adapter_path is not None:
                with open(maybe_adapter_path, "r", encoding="utf-8") as f:
                    adapter_config = json.load(f)
                    model = adapter_config["base_model_name_or_path"]

        config = AutoConfig.from_pretrained(
            model, _from_pipeline=task, code_revision=code_revision, **hub_kwargs, **model_kwargs
        )
        hub_kwargs["_commit_hash"] = config._commit_hash

    custom_tasks = {}
    if config is not None and len(getattr(config, "custom_pipelines", {})) > 0:
        custom_tasks = config.custom_pipelines
        if task is None and trust_remote_code is not False:
            if len(custom_tasks) == 1:
                task = list(custom_tasks.keys())[0]
            else:
                raise RuntimeError(
                    "We can't infer the task automatically for this model as there are multiple tasks available. Pick "
                    f"one in {', '.join(custom_tasks.keys())}"
                )

    if task is None and model is not None:
        if not isinstance(model, str):
            raise RuntimeError(
                "Inferring the task automatically requires to check the hub with a model_id defined as a `str`. "
                f"{model} is not a valid model_id."
            )
        task = get_task(model, token)

    # Retrieve the task
    if task in custom_tasks:
        normalized_task = task
        targeted_task, task_options = clean_custom_task(custom_tasks[task])
        if pipeline_class is None:
            if not trust_remote_code:
                raise ValueError(
                    "Loading this pipeline requires you to execute the code in the pipeline file in that"
                    " repo on your local machine. Make sure you have read the code there to avoid malicious use, then"
                    " set the option `trust_remote_code=True` to remove this error."
                )
            class_ref = targeted_task["impl"]
            pipeline_class = get_class_from_dynamic_module(
                class_ref,
                model,
                code_revision=code_revision,
                **hub_kwargs,
            )
    else:
        normalized_task, targeted_task, task_options = check_task(task)
        if pipeline_class is None:
            pipeline_class = targeted_task["impl"]

    # Use default model/config/tokenizer for the task if no model is provided
    if model is None:
        # At that point framework might still be undetermined
        model, default_revision = get_default_model_and_revision(targeted_task, framework, task_options)
        revision = revision if revision is not None else default_revision
        logger.warning(
            f"No model was supplied, defaulted to {model} and revision"
            f" {revision} ({HUGGINGFACE_CO_RESOLVE_ENDPOINT}/{model}).\n"
            "Using a pipeline without specifying a model name and revision in production is not recommended."
        )
        hub_kwargs["revision"] = revision
        if config is None and isinstance(model, str):
            config = AutoConfig.from_pretrained(model, _from_pipeline=task, **hub_kwargs, **model_kwargs)
            hub_kwargs["_commit_hash"] = config._commit_hash

    if device_map is not None:
        if "device_map" in model_kwargs:
            raise ValueError(
                'You cannot use both `pipeline(... device_map=..., model_kwargs={"device_map":...})` as those'
                " arguments might conflict, use only one.)"
            )
        if device is not None:
            logger.warning(
                "Both `device` and `device_map` are specified. `device` will override `device_map`. You"
                " will most likely encounter unexpected behavior. Please remove `device` and keep `device_map`."
            )
        model_kwargs["device_map"] = device_map
    if torch_dtype is not None:
        if "torch_dtype" in model_kwargs:
            raise ValueError(
                'You cannot use both `pipeline(... torch_dtype=..., model_kwargs={"torch_dtype":...})` as those'
                " arguments might conflict, use only one.)"
            )
        if isinstance(torch_dtype, str) and hasattr(torch, torch_dtype):
            torch_dtype = getattr(torch, torch_dtype)
        model_kwargs["torch_dtype"] = torch_dtype

    model_name = model if isinstance(model, str) else None

    # Load the correct model if possible
    # Infer the framework from the model if not already defined
    if isinstance(model, str) or framework is None:
        model_classes = {"tf": targeted_task["tf"], "pt": targeted_task["pt"]}
        framework, model = infer_framework_load_model(
            model,
            model_classes=model_classes,
            config=config,
            framework=framework,
            task=task,
            **hub_kwargs,
            **model_kwargs,
        )

    model_config = model.config
    hub_kwargs["_commit_hash"] = model.config._commit_hash

    load_tokenizer = type(model_config) in TOKENIZER_MAPPING or model_config.tokenizer_class is not None
    load_feature_extractor = type(model_config) in FEATURE_EXTRACTOR_MAPPING or feature_extractor is not None
    load_image_processor = type(model_config) in IMAGE_PROCESSOR_MAPPING or image_processor is not None
    load_processor = type(model_config) in PROCESSOR_MAPPING or processor is not None

    # Check that pipeline class required loading
    load_tokenizer = load_tokenizer and pipeline_class._load_tokenizer
    load_feature_extractor = load_feature_extractor and pipeline_class._load_feature_extractor
    load_image_processor = load_image_processor and pipeline_class._load_image_processor
    load_processor = load_processor and pipeline_class._load_processor

    # If `model` (instance of `PretrainedModel` instead of `str`) is passed (and/or same for config), while
    # `image_processor` or `feature_extractor` is `None`, the loading will fail. This happens particularly for some
    # vision tasks when calling `pipeline()` with `model` and only one of the `image_processor` and `feature_extractor`.
    # TODO: we need to make `NO_IMAGE_PROCESSOR_TASKS` and `NO_FEATURE_EXTRACTOR_TASKS` more robust to avoid such issue.
    # This block is only temporarily to make CI green.
    if load_image_processor and load_feature_extractor:
        load_feature_extractor = False

    if (
        tokenizer is None
        and not load_tokenizer
        and normalized_task not in NO_TOKENIZER_TASKS
        # Using class name to avoid importing the real class.
        and (
            model_config.__class__.__name__ in MULTI_MODEL_AUDIO_CONFIGS
            or model_config.__class__.__name__ in MULTI_MODEL_VISION_CONFIGS
        )
    ):
        # This is a special category of models, that are fusions of multiple models
        # so the model_config might not define a tokenizer, but it seems to be
        # necessary for the task, so we're force-trying to load it.
        load_tokenizer = True
    if (
        image_processor is None
        and not load_image_processor
        and normalized_task not in NO_IMAGE_PROCESSOR_TASKS
        # Using class name to avoid importing the real class.
        and model_config.__class__.__name__ in MULTI_MODEL_VISION_CONFIGS
    ):
        # This is a special category of models, that are fusions of multiple models
        # so the model_config might not define a tokenizer, but it seems to be
        # necessary for the task, so we're force-trying to load it.
        load_image_processor = True
    if (
        feature_extractor is None
        and not load_feature_extractor
        and normalized_task not in NO_FEATURE_EXTRACTOR_TASKS
        # Using class name to avoid importing the real class.
        and model_config.__class__.__name__ in MULTI_MODEL_AUDIO_CONFIGS
    ):
        # This is a special category of models, that are fusions of multiple models
        # so the model_config might not define a tokenizer, but it seems to be
        # necessary for the task, so we're force-trying to load it.
        load_feature_extractor = True

    if task in NO_TOKENIZER_TASKS:
        # These will never require a tokenizer.
        # the model on the other hand might have a tokenizer, but
        # the files could be missing from the hub, instead of failing
        # on such repos, we just force to not load it.
        load_tokenizer = False

    if task in NO_FEATURE_EXTRACTOR_TASKS:
        load_feature_extractor = False
    if task in NO_IMAGE_PROCESSOR_TASKS:
        load_image_processor = False

    if load_tokenizer:
        # Try to infer tokenizer from model or config name (if provided as str)
        if tokenizer is None:
            if isinstance(model_name, str):
                tokenizer = model_name
            elif isinstance(config, str):
                tokenizer = config
            else:
                # Impossible to guess what is the right tokenizer here
                raise Exception(
                    "Impossible to guess which tokenizer to use. "
                    "Please provide a PreTrainedTokenizer class or a path/identifier to a pretrained tokenizer."
                )

        # Instantiate tokenizer if needed
        if isinstance(tokenizer, (str, tuple)):
            if isinstance(tokenizer, tuple):
                # For tuple we have (tokenizer name, {kwargs})
                use_fast = tokenizer[1].pop("use_fast", use_fast)
                tokenizer_identifier = tokenizer[0]
                tokenizer_kwargs = tokenizer[1]
            else:
                tokenizer_identifier = tokenizer
                tokenizer_kwargs = model_kwargs.copy()
                tokenizer_kwargs.pop("torch_dtype", None)

            tokenizer = AutoTokenizer.from_pretrained(
                tokenizer_identifier, use_fast=use_fast, _from_pipeline=task, **hub_kwargs, **tokenizer_kwargs
            )

    if load_image_processor:
        # Try to infer image processor from model or config name (if provided as str)
        if image_processor is None:
            if isinstance(model_name, str):
                image_processor = model_name
            elif isinstance(config, str):
                image_processor = config
            # Backward compatibility, as `feature_extractor` used to be the name
            # for `ImageProcessor`.
            elif feature_extractor is not None and isinstance(feature_extractor, BaseImageProcessor):
                image_processor = feature_extractor
            else:
                # Impossible to guess what is the right image_processor here
                raise Exception(
                    "Impossible to guess which image processor to use. "
                    "Please provide a PreTrainedImageProcessor class or a path/identifier "
                    "to a pretrained image processor."
                )

        # Instantiate image_processor if needed
        if isinstance(image_processor, (str, tuple)):
            image_processor = AutoImageProcessor.from_pretrained(
                image_processor, _from_pipeline=task, **hub_kwargs, **model_kwargs
            )

    if load_feature_extractor:
        # Try to infer feature extractor from model or config name (if provided as str)
        if feature_extractor is None:
            if isinstance(model_name, str):
                feature_extractor = model_name
            elif isinstance(config, str):
                feature_extractor = config
            else:
                # Impossible to guess what is the right feature_extractor here
                raise Exception(
                    "Impossible to guess which feature extractor to use. "
                    "Please provide a PreTrainedFeatureExtractor class or a path/identifier "
                    "to a pretrained feature extractor."
                )

        # Instantiate feature_extractor if needed
        if isinstance(feature_extractor, (str, tuple)):
            feature_extractor = AutoFeatureExtractor.from_pretrained(
                feature_extractor, _from_pipeline=task, **hub_kwargs, **model_kwargs
            )

            if (
                feature_extractor._processor_class
                and feature_extractor._processor_class.endswith("WithLM")
                and isinstance(model_name, str)
            ):
                try:
                    import kenlm  # to trigger `ImportError` if not installed
                    from pyctcdecode import BeamSearchDecoderCTC

                    if os.path.isdir(model_name) or os.path.isfile(model_name):
                        decoder = BeamSearchDecoderCTC.load_from_dir(model_name)
                    else:
                        language_model_glob = os.path.join(
                            BeamSearchDecoderCTC._LANGUAGE_MODEL_SERIALIZED_DIRECTORY, "*"
                        )
                        alphabet_filename = BeamSearchDecoderCTC._ALPHABET_SERIALIZED_FILENAME
                        allow_patterns = [language_model_glob, alphabet_filename]
                        decoder = BeamSearchDecoderCTC.load_from_hf_hub(model_name, allow_patterns=allow_patterns)

                    kwargs["decoder"] = decoder
                except ImportError as e:
                    logger.warning(f"Could not load the `decoder` for {model_name}. Defaulting to raw CTC. Error: {e}")
                    if not is_kenlm_available():
                        logger.warning("Try to install `kenlm`: `pip install kenlm")

                    if not is_pyctcdecode_available():
                        logger.warning("Try to install `pyctcdecode`: `pip install pyctcdecode")

    if load_processor:
        # Try to infer processor from model or config name (if provided as str)
        if processor is None:
            if isinstance(model_name, str):
                processor = model_name
            elif isinstance(config, str):
                processor = config
            else:
                # Impossible to guess what is the right processor here
                raise Exception(
                    "Impossible to guess which processor to use. "
                    "Please provide a processor instance or a path/identifier "
                    "to a processor."
                )

        # Instantiate processor if needed
        if isinstance(processor, (str, tuple)):
            processor = AutoProcessor.from_pretrained(processor, _from_pipeline=task, **hub_kwargs, **model_kwargs)
            if not isinstance(processor, ProcessorMixin):
                raise TypeError(
                    "Processor was loaded, but it is not an instance of `ProcessorMixin`. "
                    f"Got type `{type(processor)}` instead. Please check that you specified "
                    "correct pipeline task for the model and model has processor implemented and saved."
                )

    if task == "translation" and model.config.task_specific_params:
        for key in model.config.task_specific_params:
            if key.startswith("translation"):
                task = key
                warnings.warn(
                    f'"translation" task was used, instead of "translation_XX_to_YY", defaulting to "{task}"',
                    UserWarning,
                )
                break

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

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

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

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

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

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



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

**Overview**

This code is part of a Natural Language Processing (NLP) model loading and creation process. It defines a function `load_model` that creates an instance of the specified `pipeline_class` by parsing the task-specific configuration, loading any necessary preprocessing layers, and setting up various model parameters.

**Key Components**

1. **Preprocessing Layers**: The code loads and instantiates various preprocessing layers such as tokenizer, image processor, and feature extractor based on task specifications.
2. **Model Parameters**: It sets various model parameters including device, torch_dtype, and processor settings according to the task requirements.
3. **Return Value**: The function returns an instance of `pipeline_class` with the specified model, framework, task, and configuration.

**Notes**

* This code appears to be based on Hugging Face's Transformers library for NLP models.
* The code uses various pre-trained models and libraries such as AutoTokenizer, AutoFeatureExtractor, and BeamSearchDecoderCTC.


key: chroma:document
string_value:
CODE:

def read_video_pyav(container, indices):
    frames = []
    container.seek(0)
    start_index = indices[0]
    end_index = indices[-1]
    for i, frame in enumerate(container.decode(video=0)):
        if i > end_index:
            break
        if i >= start_index and i in indices:
            frames.append(frame)


SUMMARY:
**Function Summary:**

`read_video_pyav(container, indices)` reads video frames from PyAV container.

**Purpose:** Extract specific frames from a video stream based on provided indices and store them as RGB values.

**Key Implementation Details:**

* Uses `container.decode()` to iterate over the video frame indices.
* Filters frames within specified range (`start_index` to `end_index`) and includes only those present in `indices`.
* Converts each frame to RGB format using `to_ndarray(format="rgb24")`.

**Dependencies:**

* PyAV
* NumPy

Note: No external dependencies or interfaces are used. The summary focuses on the key functionality and algorithms used in the function.


key: chroma:document
string_value:
CODE:

def no_collate_fn(items):
    if len(items) != 1:
        raise ValueError("This collate_fn is meant to be used with batch_size=1")


SUMMARY:
**Function Summary**

`no_collate_fn`: A custom collation function that returns the only item in an iterable, or raises an error if it's not.

**Key Implementation Details**

* Raises `ValueError` if input iterable has more than one element
* Returns single-element iterable directly (i.e., No Collation)

**Notable Dependencies/Requirements**

* None referenced; uses standard Python `raise` and `return` statements


key: chroma:document
string_value:
CODE:

def _pad(items, key, padding_value, padding_side):
    batch_size = len(items)
    if isinstance(items[0][key], torch.Tensor):
        # Others include `attention_mask` etc...
        shape = items[0][key].shape
        dim = len(shape)
        if dim == 1:
            # We have a list of 1-dim torch tensors, which can be stacked without padding
            return torch.cat([item[key] for item in items], dim=0)
        if key in ["pixel_values", "image"]:
            # This is probable image so padding shouldn't be necessary
            # B, C, H, W
            return torch.cat([item[key] for item in items], dim=0)
        elif dim == 4 and key == "input_features":
            # this is probably a mel spectrogram batched
            return torch.cat([item[key] for item in items], dim=0)
        max_length = max(item[key].shape[1] for item in items)
        min_length = min(item[key].shape[1] for item in items)
        dtype = items[0][key].dtype

        if dim == 2:
            if max_length == min_length:
                # Bypass for `ImageGPT` which doesn't provide a padding value, yet
                # we can consistently pad since the size should be matching
                return torch.cat([item[key] for item in items], dim=0)
            tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value
        elif dim == 3:
            tensor = torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype) + padding_value
        elif dim == 4:
            tensor = torch.zeros((batch_size, max_length, shape[-2], shape[-1]), dtype=dtype) + padding_value

        for i, item in enumerate(items):
            if dim == 2:
                if padding_side == "left":
                    tensor[i, -len(item[key][0]) :] = item[key][0].clone()
                else:
                    tensor[i, : len(item[key][0])] = item[key][0].clone()
            elif dim == 3:
                if padding_side == "left":
                    tensor[i, -len(item[key][0]) :, :] = item[key][0].clone()
                else:
                    tensor[i, : len(item[key][0]), :] = item[key][0].clone()
            elif dim == 4:
                if padding_side == "left":
                    tensor[i, -len(item[key][0]) :, :, :] = item[key][0].clone()
                else:
                    tensor[i, : len(item[key][0]), :, :] = item[key][0].clone()

        return tensor
    else:


SUMMARY:
**Function Summary: `_pad(items, key, padding_value, padding_side)`**

**Main Purpose and Functionality:**
This function pads or truncates the last dimension of tensors along the `key` axis to achieve a uniform length (`max_length`) in batches. It handles different tensor dimensions (1D, 2D, 3D, 4D) and applies padding according to the specified `padding_side`.

**Key Implementation Details:**

* It checks if the first item in the batch is a torch Tensor using `isinstance(items[0][key], torch.Tensor)`.
* For different tensor dimensions, it appends or prepends padding values to maintain consistency.
* It uses `torch.cat` and slicing operations (`tensor[i, ...]`) to construct the padded tensors.

**Notable Dependencies and Requirements:**

* `torch`
* `torchvision`

No external dependencies are listed as they appear to be part of the PyTorch library.


key: chroma:document
string_value:
CODE:

def pad_collate_fn(tokenizer, feature_extractor):
    # Tokenizer
    t_padding_side = None
    # Feature extractor
    f_padding_side = None
    if tokenizer is None and feature_extractor is None:
        raise ValueError("Pipeline without tokenizer or feature_extractor cannot do batching")
    if tokenizer is not None:
        if tokenizer.pad_token_id is None:
            raise ValueError(
                "Pipeline with tokenizer without pad_token cannot do batching. You can try to set it with "
                "`pipe.tokenizer.pad_token_id = model.config.eos_token_id`."
            )
        else:
            t_padding_value = tokenizer.pad_token_id
            t_padding_side = tokenizer.padding_side
    if feature_extractor is not None:
        # Feature extractor can be images, where no padding is expected
        f_padding_value = getattr(feature_extractor, "padding_value", None)
        f_padding_side = getattr(feature_extractor, "padding_side", None)

    if t_padding_side is not None and f_padding_side is not None and t_padding_side != f_padding_side:
        raise ValueError(
            f"The feature extractor, and tokenizer don't agree on padding side {t_padding_side} != {f_padding_side}"
        )
    padding_side = "right"
    if t_padding_side is not None:
        padding_side = t_padding_side
    if f_padding_side is not None:
        padding_side = f_padding_side

    def inner(items):
        keys = set(items[0].keys())
        for item in items:
            if set(item.keys()) != keys:
                raise ValueError(
                    f"The elements of the batch contain different keys. Cannot batch them ({set(item.keys())} !="
                    f" {keys})"
                )
        # input_values, input_pixels, input_ids, ...
        padded = {}
        for key in keys:
            if key in {"input_ids"}:
                # ImageGPT uses a feature extractor
                if tokenizer is None and feature_extractor is not None:
                    _padding_value = f_padding_value
                else:
                    _padding_value = t_padding_value
            elif key in {"input_values", "pixel_values", "input_features"}:
                _padding_value = f_padding_value
            elif key in {"p_mask", "special_tokens_mask"}:
                _padding_value = 1
            elif key in {"attention_mask", "token_type_ids"}:
                _padding_value = 0
            else:
                # This is likely another random key maybe even user provided
                _padding_value = 0
            padded[key] = _pad(items, key, _padding_value, padding_side)
        return padded



SUMMARY:
**Function Summary: `pad_collate_fn`**

**Purpose and Functionality:**
The `pad_collate_fn` function is a batching utility for preparing input data for model inference. It aligns input features from multiple sources (tokenizer & feature extractor) by padding values appropriately.

**Key Implementation Details:**

*   Verifies tokenizer and feature extractor configurations.
*   Checks for consistency in padding sides between tokenizer and feature extractor.
*   Applies padding to consistent side of inputs based on `padding_side` configuration.

**Notable Dependencies and Requirements:**

*   `tokenizer`: used for text input formatting
*   `feature_extractor`: used for image input formatting (optional)
    *   `pad_token_id`
    *   `padding_value`

**External Interface:*

The function returns a transformation object (`inner`) that accepts batches of input data.

*   Example usage: `padded = inner([dataset_example])`, where `dataset_example` is an example batch containing input data for the model.


key: chroma:document
string_value:
CODE:

def infer_framework_load_model(
    model,
    config: AutoConfig,
    model_classes: Optional[Dict[str, Tuple[type]]] = None,
    task: Optional[str] = None,
    framework: Optional[str] = None,
    **model_kwargs,
):
    """
    Select framework (TensorFlow or PyTorch) to use from the `model` passed. Returns a tuple (framework, model).

    If `model` is instantiated, this function will just infer the framework from the model class. Otherwise `model` is
    actually a checkpoint name and this method will try to instantiate it using `model_classes`. Since we don't want to
    instantiate the model twice, this model is returned for use by the pipeline.

    If both frameworks are installed and available for `model`, PyTorch is selected.

    Args:
        model (`str`, [`PreTrainedModel`] or [`TFPreTrainedModel]`):
            The model to infer the framework from. If `str`, a checkpoint name. The model to infer the framewrok from.
        config ([`AutoConfig`]):
            The config associated with the model to help using the correct class
        model_classes (dictionary `str` to `type`, *optional*):
            A mapping framework to class.
        task (`str`):
            The task defining which pipeline will be returned.
        model_kwargs:
            Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,
            **model_kwargs)` function.

    Returns:
        `Tuple`: A tuple framework, model.
    """
    if not is_tf_available() and not is_torch_available():
        raise RuntimeError(
            "At least one of TensorFlow 2.0 or PyTorch should be installed. "
            "To install TensorFlow 2.0, read the instructions at https://www.tensorflow.org/install/ "
            "To install PyTorch, read the instructions at https://pytorch.org/."
        )
    if isinstance(model, str):
        model_kwargs["_from_pipeline"] = task
        class_tuple = ()
        look_pt = is_torch_available() and framework in {"pt", None}
        look_tf = is_tf_available() and framework in {"tf", None}
        if model_classes:
            if look_pt:
                class_tuple = class_tuple + model_classes.get("pt", (AutoModel,))
            if look_tf:
                class_tuple = class_tuple + model_classes.get("tf", (TFAutoModel,))
        if config.architectures:
            classes = []
            for architecture in config.architectures:
                transformers_module = importlib.import_module("transformers")
                if look_pt:
                    _class = getattr(transformers_module, architecture, None)
                    if _class is not None:
                        classes.append(_class)
                if look_tf:
                    _class = getattr(transformers_module, f"TF{architecture}", None)
                    if _class is not None:
                        classes.append(_class)
            class_tuple = class_tuple + tuple(classes)

        if len(class_tuple) == 0:
            raise ValueError(f"Pipeline cannot infer suitable model classes from {model}")

        all_traceback = {}
        for model_class in class_tuple:
            kwargs = model_kwargs.copy()
            if framework == "pt" and model.endswith(".h5"):
                kwargs["from_tf"] = True
                logger.warning(
                    "Model might be a TensorFlow model (ending with `.h5`) but TensorFlow is not available. "
                    "Trying to load the model with PyTorch."
                )
            elif framework == "tf" and model.endswith(".bin"):
                kwargs["from_pt"] = True
                logger.warning(
                    "Model might be a PyTorch model (ending with `.bin`) but PyTorch is not available. "
                    "Trying to load the model with Tensorflow."
                )

            try:
                model = model_class.from_pretrained(model, **kwargs)
                if hasattr(model, "eval"):
                    model = model.eval()
                # Stop loading on the first successful load.
                break
            except (OSError, ValueError):
                all_traceback[model_class.__name__] = traceback.format_exc()
                continue

        if isinstance(model, str):
            error = ""
            for class_name, trace in all_traceback.items():
                error += f"while loading with {class_name}, an error is thrown:\n{trace}\n"
            raise ValueError(
                f"Could not load model {model} with any of the following classes: {class_tuple}. See the original errors:\n\n{error}\n"
            )

    if framework is None:
        framework = infer_framework(model.__class__)


SUMMARY:
**Summary**

Infer Framework Load Model Function
=====================================

### Purpose and Functionality

This function inferences the framework (TensorFlow or PyTorch) from a given model. It returns a tuple containing the framework and loaded model.

### Key Implementation Details

*   Inferences framework based on TensorFlow availability and framework name.
*   Creates a dictionary of models using `model_classes` mapping.
*   Installs necessary modules and checks their availability before loading models.
*   Tries to load multiple classes in sequence, with error handling for each failed attempt.

### Notable Dependencies or Requirements

*   Torch library (`transformers`)
*   TensorFlow library
*   `AutoConfig`
*   `TFPreTrainedModel`

### Return Value

A tuple containing the inferred framework and loaded model.
```python
(framework: str, model)
```


key: chroma:document
string_value:
CODE:

def infer_framework_from_model(
    model,
    model_classes: Optional[Dict[str, Tuple[type]]] = None,
    task: Optional[str] = None,
    framework: Optional[str] = None,
    **model_kwargs,
):
    """
    Select framework (TensorFlow or PyTorch) to use from the `model` passed. Returns a tuple (framework, model).

    If `model` is instantiated, this function will just infer the framework from the model class. Otherwise `model` is
    actually a checkpoint name and this method will try to instantiate it using `model_classes`. Since we don't want to
    instantiate the model twice, this model is returned for use by the pipeline.

    If both frameworks are installed and available for `model`, PyTorch is selected.

    Args:
        model (`str`, [`PreTrainedModel`] or [`TFPreTrainedModel]`):
            The model to infer the framework from. If `str`, a checkpoint name. The model to infer the framewrok from.
        model_classes (dictionary `str` to `type`, *optional*):
            A mapping framework to class.
        task (`str`):
            The task defining which pipeline will be returned.
        model_kwargs:
            Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,
            **model_kwargs)` function.

    Returns:
        `Tuple`: A tuple framework, model.
    """
    if isinstance(model, str):
        config = AutoConfig.from_pretrained(model, _from_pipeline=task, **model_kwargs)
    else:
        config = model.config
    return infer_framework_load_model(
        model, config, model_classes=model_classes, _from_pipeline=task, task=task, framework=framework, **model_kwargs


SUMMARY:
**Function Summary: `infer_framework_from_model`**

*Purpose*: Selects the framework (TensorFlow or PyTorch) to use with a given model based on its class and configuration.

**Key Implementation Details:**

* Infers framework from model class if model is an instance
* Tries to instantiate model using `model_classes` dictionary if model is a string checkpoint name
* Priority given to PyTorch over TensorFlow when both are installed

**Method Signature:** `infer_framework_from_model(model, model_classes, task, framework, **model_kwargs) -> Tuple[framework, model]`

**Notable Dependencies and Requirements:**

* `AutoConfig.from_pretrained` for loading configuration
* `model.classes` for accessing model configuration
* `model_classes` dictionary for mapping frameworks to classes
* `TensorFlow` and `PyTorch` installations determined at runtime


key: chroma:document
string_value:
CODE:

def get_framework(model, revision: Optional[str] = None):
    """
    Select framework (TensorFlow or PyTorch) to use.

    Args:
        model (`str`, [`PreTrainedModel`] or [`TFPreTrainedModel]`):
            If both frameworks are installed, picks the one corresponding to the model passed (either a model class or
            the model name). If no specific model is provided, defaults to using PyTorch.
    """
    warnings.warn(
        "`get_framework` is deprecated and will be removed in v5, use `infer_framework_from_model` instead.",
        FutureWarning,
    )
    if not is_tf_available() and not is_torch_available():
        raise RuntimeError(
            "At least one of TensorFlow 2.0 or PyTorch should be installed. "
            "To install TensorFlow 2.0, read the instructions at https://www.tensorflow.org/install/ "
            "To install PyTorch, read the instructions at https://pytorch.org/."
        )
    if isinstance(model, str):
        if is_torch_available() and not is_tf_available():
            model = AutoModel.from_pretrained(model, revision=revision)
        elif is_tf_available() and not is_torch_available():
            model = TFAutoModel.from_pretrained(model, revision=revision)
        else:
            try:
                model = AutoModel.from_pretrained(model, revision=revision)
            except OSError:
                model = TFAutoModel.from_pretrained(model, revision=revision)

    framework = infer_framework(model.__class__)


SUMMARY:
**Function Summary: `get_framework`**

### Purpose and Functionality

The `get_framework` function determines the underlying framework (TensorFlow or PyTorch) to use for a given model. It takes a `model` argument, which can be either a string representing the model name or an instance of `PreTrainedModel`/`TFPreTrainedModel`.

### Key Implementation Details

* Uses the `infer_framework_from_model` function to determine the framework based on the model's class.
* Handles different cases for framework availability and assigns the corresponding `AutoModel`/`TFAutoModel` instance.

### Notable Dependencies or Requirements

* TensorFlow 2.0 or PyTorch must be installed.
* In cases where a specific model is not provided, the function defaults to using PyTorch.

### Example Parameters

* `model`: string (model name) or `PreTrainedModel`/`TFPreTrainedModel` instance.

### Return Value
The returned framework instance corresponds to the underlying framework used for the given model.


key: chroma:document
string_value:
CODE:

def get_default_model_and_revision(
    targeted_task: Dict, framework: Optional[str], task_options: Optional[Any]
) -> Union[str, Tuple[str, str]]:
    """
    Select a default model to use for a given task. Defaults to pytorch if ambiguous.

    Args:
        targeted_task (`Dict`):
           Dictionary representing the given task, that should contain default models

        framework (`str`, None)
           "pt", "tf" or None, representing a specific framework if it was specified, or None if we don't know yet.

        task_options (`Any`, None)
           Any further value required by the task to get fully specified, for instance (SRC, TGT) languages for
           translation task.

    Returns

        `str` The model string representing the default model for this pipeline
    """
    if is_torch_available() and not is_tf_available():
        framework = "pt"
    elif is_tf_available() and not is_torch_available():
        framework = "tf"

    defaults = targeted_task["default"]
    if task_options:
        if task_options not in defaults:
            raise ValueError(f"The task does not provide any default models for options {task_options}")
        default_models = defaults[task_options]["model"]
    elif "model" in defaults:
        default_models = targeted_task["default"]["model"]
    else:
        # XXX This error message needs to be updated to be more generic if more tasks are going to become
        # parametrized
        raise ValueError('The task defaults can\'t be correctly selected. You probably meant "translation_XX_to_YY"')

    if framework is None:
        framework = "pt"



SUMMARY:
**Function Summary: `get_default_model_and_revision`**

**Purpose and Functionality:**
This function selects a default model for a given task based on the provided framework and task options. It returns the model string representing the default model for this pipeline.

*   If no explicit framework is specified, it defaults to PyTorch.
*   It checks if specific task options are provided and uses the corresponding default models.
*   It raises an error if the framework or task options are not supported.

**Key Implementation Details:**

*   Uses conditional statements to handle different frameworks (PyTorch/TensorFlow) based on availability.
*   Utilizes a dictionary (`defaults`) to store default models for specific task options or translations.

**Notable Dependencies and Requirements:**

*   The function uses the following external dependencies:
    *   `is_torch_available()` and `is_tf_available()` functions
    *   `Optional` and `Union` types from the standard library
    *   `Any` type as an argument for some function parameters

There are no inherited or implemented interfaces mentioned in this code chunk.


key: chroma:document
string_value:
CODE:

def load_assistant_model(
    model: "PreTrainedModel",
    assistant_model: Optional[Union[str, "PreTrainedModel"]],
    assistant_tokenizer: Optional[PreTrainedTokenizer],
) -> Tuple[Optional["PreTrainedModel"], Optional[PreTrainedTokenizer]]:
    """
    Prepares the assistant model and the assistant tokenizer for a pipeline whose model that can call `generate`.

    Args:
        model ([`PreTrainedModel`]):
            The main model that will be used by the pipeline to make predictions.
        assistant_model (`str` or [`PreTrainedModel`], *optional*):
            The assistant model that will be used by the pipeline to make predictions.
        assistant_tokenizer ([`PreTrainedTokenizer`], *optional*):
            The assistant tokenizer that will be used by the pipeline to encode data for the model.

    Returns:
        Tuple: The loaded assistant model and (optionally) the loaded tokenizer.
    """
    if not model.can_generate() or assistant_model is None:
        return None, None

    if not isinstance(model, PreTrainedModel):
        raise ValueError(
            "Assisted generation, triggered by the `assistant_model` argument, is only available for "
            "`PreTrainedModel` model instances. For instance, TF or JAX models are not supported."
        )

    # If the model is passed as a string, load the model and the corresponding tokenizer
    if isinstance(assistant_model, str):
        assistant_config = AutoConfig.from_pretrained(assistant_model)
        _, loaded_assistant_model = infer_framework_load_model(assistant_model, config=assistant_config)
        loaded_assistant_model = loaded_assistant_model.to(device=model.device, dtype=model.dtype)
        loaded_assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_model)
    else:
        loaded_assistant_model = assistant_model
        loaded_assistant_tokenizer = assistant_tokenizer

    # Finally, let's check the tokenizers: if the two models have different tokenizers, we need to keep the assistant
    # tokenizer
    same_vocab_size = model.config.vocab_size == loaded_assistant_model.config.vocab_size
    same_special_tokens = all(
        getattr(model.config, token) == getattr(loaded_assistant_model.config, token)
        for token in ("eos_token_id", "pad_token_id", "bos_token_id")
    )
    if same_vocab_size and same_special_tokens:
        loaded_assistant_tokenizer = None
    elif loaded_assistant_tokenizer is None:
        raise ValueError(
            "The assistant model has a different tokenizer than the main model. You should pass the assistant "
            "tokenizer."
        )



SUMMARY:
**Function Summary: `load_assistant_model`**

Prepares an assistant model and tokenizer for a pipeline that can perform generation.

```markdown
### Purpose and Functionality
The main purpose of this function is to load the assistant model (`assistant_model`) along with its tokenizer (`assistant_tokenizer`) if applicable. It is essential for pipeline operations that involve assisted generation, which requires both model instances to have identical vocab sizes and token configurations.
### Key Implementation Details

- The function first checks whether the `assistant_model` can generate and if it is a valid `PreTrainedModel`.
- If the `assistant_model` argument is passed as a string, it loads the corresponding model and tokenizer instances using pre-trained AutoConfigs from Hugging Face's library.

### Notable Dependencies and Requirements
The function relies on the following:
- `AutoConfig`
- `InferFrameworkLoadModel`
- `PreTrainedTokenizer` (`assistant_tokenizer`)
- `PreTrainedModel` (`model`, `assistant_model`)
- The user-provided or inferred device type for model operations.

### Return Value
A tuple containing the loaded `assistant_model` and optionally, the loaded `assistant_tokenizer`. If both are the same or if no tokenization is required due to matching vocab size, the `assistant_tokenizer` will be set to None.


key: chroma:document
string_value:
CODE:

class PipelineException(Exception):
    """
    Raised by a [`Pipeline`] when handling __call__.

    Args:
        task (`str`): The task of the pipeline.
        model (`str`): The model used by the pipeline.
        reason (`str`): The error message to display.
    """

    def __init__(self, task: str, model: str, reason: str):
        super().__init__(reason)

        self.task = task


SUMMARY:
**PipelineException Class Summary**

The `PipelineException` class is a custom exception raised by the Pipeline when handling the `__call__` method. It provides information about the task, model used, and error reason.

**Key Implementation Details:**

- Custom exception type.
- Initializes with `task`, `model`, and `reason`.
  - `super().__init__(reason)` sets the error message.
  
**Notable Dependencies or Requirements:**

- [`Pipeline` class] (related class).
- No external dependencies.

**Primary Purpose:**
Raise a custom exception with relevant task, model, and error information when handling the Pipeline's `__call__` method.


key: chroma:document
string_value:
CODE:

class ArgumentHandler(ABC):
    """
    Base interface for handling arguments for each [`~pipelines.Pipeline`].
    """

    @abstractmethod
    def __call__(self, *args, **kwargs):


SUMMARY:
**ArgumentHandler Class Summary**

The `ArgumentHandler` class serves as a base interface for handling arguments specific to each pipeline in a workflow. 

**Key Functionality:**
- Provides a single, abstract constructor (`__call__`) that expects variable number of positional and keyword arguments.

**Notable Implementation Details:**
- Enforces the use of a concrete implementation through the `NotImplementedError` exception, indicating that subclasses must provide an actual implementation for handling arguments.
 
**Dependencies and Requirements:**
- No explicit dependencies are mentioned. However, it is implicitly dependent on the `pipelines.Pipeline` class, referencing its scope.


key: chroma:document
string_value:
CODE:

class PipelineDataFormat:
    """
    Base class for all the pipeline supported data format both for reading and writing. Supported data formats
    currently includes:

    - JSON
    - CSV
    - stdin/stdout (pipe)

    `PipelineDataFormat` also includes some utilities to work with multi-columns like mapping from datasets columns to
    pipelines keyword arguments through the `dataset_kwarg_1=dataset_column_1` format.

    Args:
        output_path (`str`): Where to save the outgoing data.
        input_path (`str`): Where to look for the input data.
        column (`str`): The column to read.
        overwrite (`bool`, *optional*, defaults to `False`):
            Whether or not to overwrite the `output_path`.
    """

    SUPPORTED_FORMATS = ["json", "csv", "pipe"]

    def __init__(
        self,
        output_path: Optional[str],
        input_path: Optional[str],
        column: Optional[str],
        overwrite: bool = False,
    ):
        self.output_path = output_path
        self.input_path = input_path
        self.column = column.split(",") if column is not None else [""]
        self.is_multi_columns = len(self.column) > 1

        if self.is_multi_columns:
            self.column = [tuple(c.split("=")) if "=" in c else (c, c) for c in self.column]

        if output_path is not None and not overwrite:
            if exists(abspath(self.output_path)):
                raise OSError(f"{self.output_path} already exists on disk")

        if input_path is not None:
            if not exists(abspath(self.input_path)):
                raise OSError(f"{self.input_path} doesnt exist on disk")

    @abstractmethod
    def __iter__(self):
        raise NotImplementedError()

    @abstractmethod
    def save(self, data: Union[dict, List[dict]]):
        """
        Save the provided data object with the representation for the current [`~pipelines.PipelineDataFormat`].

        Args:
            data (`dict` or list of `dict`): The data to store.
        """
        raise NotImplementedError()

    def save_binary(self, data: Union[dict, List[dict]]) -> str:
        """
        Save the provided data object as a pickle-formatted binary data on the disk.

        Args:
            data (`dict` or list of `dict`): The data to store.

        Returns:
            `str`: Path where the data has been saved.
        """
        path, _ = os.path.splitext(self.output_path)
        binary_path = os.path.extsep.join((path, "pickle"))

        with open(binary_path, "wb+") as f_output:
            pickle.dump(data, f_output)

        return binary_path

    @staticmethod
    def from_str(
        format: str,
        output_path: Optional[str],
        input_path: Optional[str],
        column: Optional[str],
        overwrite=False,
    ) -> "PipelineDataFormat":
        """
        Creates an instance of the right subclass of [`~pipelines.PipelineDataFormat`] depending on `format`.

        Args:
            format (`str`):
                The format of the desired pipeline. Acceptable values are `"json"`, `"csv"` or `"pipe"`.
            output_path (`str`, *optional*):
                Where to save the outgoing data.
            input_path (`str`, *optional*):
                Where to look for the input data.
            column (`str`, *optional*):
                The column to read.
            overwrite (`bool`, *optional*, defaults to `False`):
                Whether or not to overwrite the `output_path`.

        Returns:
            [`~pipelines.PipelineDataFormat`]: The proper data format.
        """
        if format == "json":
            return JsonPipelineDataFormat(output_path, input_path, column, overwrite=overwrite)
        elif format == "csv":
            return CsvPipelineDataFormat(output_path, input_path, column, overwrite=overwrite)
        elif format == "pipe":
            return PipedPipelineDataFormat(output_path, input_path, column, overwrite=overwrite)
        else:


SUMMARY:
**Summary of `PipelineDataFormat` Class**

The `PipelineDataFormat` class is a base class for various pipeline data formats, including JSON, CSV, and pipe. It provides utilities for working with multi-columns and supports reading/writing data to disk.

### Main Purpose and Functionality

- **Read/Write Data**: Supports reading from and writing to various data formats.
- **Multi-column Handling**: Provides functionality for mapping dataset columns to pipeline keyword arguments.

### Key Implementation Details

*   **`__init__` method**:
    *   Initializes the class with output, input paths, column names, and overwrite flag.
    *   Checks if the output path exists and raises an error if it does not exist or needs to be overwritten.
*   **`save` interface**: Defines a method for saving data to disk, which is implemented by subclasses (JSON or CSV format).
*   **`save_binary` method**: Saves data in binary pickle format on disk.
*   **`from_str` static method**: Creates an instance of the correct subclass based on the specified `format`.

### Notable Dependencies and Requirements

-   No external dependencies are listed; however, it likely requires `os`, `pickle`, and `json` libraries.

Note that there's no inheritance or implemented interfaces explicitly mentioned in this class summary. However, each formatted pipeline (`PipelineDataFormat` subclass) inherits from the base format.


key: chroma:document
string_value:
CODE:

class CsvPipelineDataFormat(PipelineDataFormat):
    """
    Support for pipelines using CSV data format.

    Args:
        output_path (`str`): Where to save the outgoing data.
        input_path (`str`): Where to look for the input data.
        column (`str`): The column to read.
        overwrite (`bool`, *optional*, defaults to `False`):
            Whether or not to overwrite the `output_path`.
    """

    def __init__(
        self,
        output_path: Optional[str],
        input_path: Optional[str],
        column: Optional[str],
        overwrite=False,
    ):
        super().__init__(output_path, input_path, column, overwrite=overwrite)

    def __iter__(self):
        with open(self.input_path, "r") as f:
            reader = csv.DictReader(f)
            for row in reader:
                if self.is_multi_columns:
                    yield {k: row[c] for k, c in self.column}
                else:
                    yield row[self.column[0]]

    def save(self, data: List[dict]):
        """
        Save the provided data object with the representation for the current [`~pipelines.PipelineDataFormat`].

        Args:
            data (`List[dict]`): The data to store.
        """
        with open(self.output_path, "w") as f:
            if len(data) > 0:
                writer = csv.DictWriter(f, list(data[0].keys()))
                writer.writeheader()


SUMMARY:
**CsvPipelineDataFormat Class Summary**

The `CsvPipelineDataFormat` class provides support for pipelines using CSV data format. Its main purpose is to read and write CSV data.

**Primary Purpose:**
To load and save CSV data from `input_path` to `output_path` based on a specified column.

**Key Implementation Details:**

* The `__iter__` method iterates over the CSV rows, yielding a dictionary representation (dict) of each row.
* The `save` method writes data to a CSV file at `output_path`.

**Method Signatures and Interfaces:**

* `__init__(output_path: Optional[str], input_path: Optional[str], column: Optional[str], overwrite=False)`
* `__iter__()`
* `save(data: List[dict])`

**Notable Dependencies or Requirements:**

* `csv.DictReader` and `csv.DictWriter` for reading and writing CSV data, respectively.
* The `pipelines.PipelineDataFormat` interface is inherited.

**External Dependencies:**

* The `csv` module (part of the Python Standard Library)
* Assumed dependency on other pipelines classes in the `pipelines` package.


key: chroma:document
string_value:
CODE:

class JsonPipelineDataFormat(PipelineDataFormat):
    """
    Support for pipelines using JSON file format.

    Args:
        output_path (`str`): Where to save the outgoing data.
        input_path (`str`): Where to look for the input data.
        column (`str`): The column to read.
        overwrite (`bool`, *optional*, defaults to `False`):
            Whether or not to overwrite the `output_path`.
    """

    def __init__(
        self,
        output_path: Optional[str],
        input_path: Optional[str],
        column: Optional[str],
        overwrite=False,
    ):
        super().__init__(output_path, input_path, column, overwrite=overwrite)

        with open(input_path, "r") as f:
            self._entries = json.load(f)

    def __iter__(self):
        for entry in self._entries:
            if self.is_multi_columns:
                yield {k: entry[c] for k, c in self.column}
            else:
                yield entry[self.column[0]]

    def save(self, data: dict):
        """
        Save the provided data object in a json file.

        Args:
            data (`dict`): The data to store.
        """
        with open(self.output_path, "w") as f:


SUMMARY:
**JsonPipelineDataFormat Class Summary**

### Purpose and Functionality

The `JsonPipelineDataFormat` class is designed for Pipelines using JSON file format. It supports reading input data from a JSON file, processing the data through various columns (if applicable), and writing processed data to an output JSON file.

### Key Implementation Details

*   Reads input data from a specified JSON file and stores it in `_entries`.
*   Implements a custom iterator (`__iter__`) that yields either individual entries or columns depending on `is_multi_columns` setting.
*   Saves the provided `data` object in a formatted JSON file using `json.dump`.

### Notable Dependencies and Requirements

*   **Python Standard Library:** `json`, `io`
*   **Optional Packages/(Classes):** `PipelineDataFormat`

### Key Method Signatures and Interfaces

*   `__init__(output_path: str, input_path: str, column: str, overwrite: bool = False)`
*   `save(data: dict)`
*   `is_multi_columns.setter` (implies `is_multi_columns` is a property)


key: chroma:document
string_value:
CODE:

class PipedPipelineDataFormat(PipelineDataFormat):
    """
    Read data from piped input to the python process. For multi columns data, columns should separated by \t

    If columns are provided, then the output will be a dictionary with {column_x: value_x}

    Args:
        output_path (`str`): Where to save the outgoing data.
        input_path (`str`): Where to look for the input data.
        column (`str`): The column to read.
        overwrite (`bool`, *optional*, defaults to `False`):
            Whether or not to overwrite the `output_path`.
    """

    def __iter__(self):
        for line in sys.stdin:
            # Split for multi-columns
            if "\t" in line:
                line = line.split("\t")
                if self.column:
                    # Dictionary to map arguments
                    yield {kwargs: l for (kwargs, _), l in zip(self.column, line)}
                else:
                    yield tuple(line)

            # No dictionary to map arguments
            else:
                yield line

    def save(self, data: dict):
        """
        Print the data.

        Args:
            data (`dict`): The data to store.
        """
        print(data)

    def save_binary(self, data: Union[dict, List[dict]]) -> str:
        if self.output_path is None:
            raise KeyError(
                "When using piped input on pipeline outputting large object requires an output file path. "
                "Please provide such output path through --output argument."
            )



SUMMARY:
**PipedPipelineDataFormat Class Summary**

### Purpose and Functionality
Implements data pipeline for piped input, producing dictionaries or tuples based on specified columns.

### Implementation Details

*   **`__iter__` method**: Iterates over system stdin, splitting lines into multiple columns (if separated by `\t`) and yielding dictionaries with column names as keys.
*   **`save` method**: Prints data as a dictionary.
*   **`save_binary` method**: Saves binary data to output path (requires `output_path` argument). Calls superclass method if not provided.

### Dependencies and Requirements

*   Uses `sys.stdin` for piped input
*   Requires `PipelineDataFormat` superclass (not shown in snippet)
*   Assumes Python's built-in `print` function and dictionary data type.
*   No external dependencies noted.


key: chroma:document
string_value:
CODE:

class _ScikitCompat(ABC):
    """
    Interface layer for the Scikit and Keras compatibility.
    """

    @abstractmethod
    def transform(self, X):
        raise NotImplementedError()

    @abstractmethod
    def predict(self, X):


SUMMARY:
**Summary:**
```python
class _ScikitCompat(ABC):
    """
    Interface layer for Scikit/Keras compatibility.
    
    Provides abstract methods for transform and predict operations.
    """

    @abstractmethod
    def transform(self, X):
        # Not implemented; raised on call
    
    @abstractmethod
    def predict(self, X):
        # Not implemented; raised on call
```

**Analysis:**

* The `_ScikitCompat` class serves as an interface for compatibility between Scikit and Keras libraries.
* It has two abstract methods:
	+ `transform(X)`: not implemented, raises `NotImplementedError`.
	+ `predict(X)`: not implemented, raises `NotImplementedError`.

**Key Points:**

* Defines the interface layer for compatibility between Scikit and Keras libraries.
* Abstracts transform and predict operations with empty implementations.
* Notable dependency: None (builtin Python/standard library).

Note: The summary is concise and focused on the main purpose, key implementation details, and notable dependencies or requirements.


key: chroma:document
string_value:
CODE:

def build_pipeline_init_args(
    has_tokenizer: bool = False,
    has_feature_extractor: bool = False,
    has_image_processor: bool = False,
    has_processor: bool = False,
    supports_binary_output: bool = True,
) -> str:
    docstring = r"""
    Arguments:
        model ([`PreTrainedModel`] or [`TFPreTrainedModel`]):
            The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from
            [`PreTrainedModel`] for PyTorch and [`TFPreTrainedModel`] for TensorFlow."""
    if has_tokenizer:
        docstring += r"""
        tokenizer ([`PreTrainedTokenizer`]):
            The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from
            [`PreTrainedTokenizer`]."""
    if has_feature_extractor:
        docstring += r"""
        feature_extractor ([`SequenceFeatureExtractor`]):
            The feature extractor that will be used by the pipeline to encode data for the model. This object inherits from
            [`SequenceFeatureExtractor`]."""
    if has_image_processor:
        docstring += r"""
        image_processor ([`BaseImageProcessor`]):
            The image processor that will be used by the pipeline to encode data for the model. This object inherits from
            [`BaseImageProcessor`]."""
    if has_processor:
        docstring += r"""
        processor ([`ProcessorMixin`]):
            The processor that will be used by the pipeline to encode data for the model. This object inherits from
            [`ProcessorMixin`]. Processor is a composite object that might contain `tokenizer`, `feature_extractor`, and
            `image_processor`."""
    docstring += r"""
        modelcard (`str` or [`ModelCard`], *optional*):
            Model card attributed to the model for this pipeline.
        framework (`str`, *optional*):
            The framework to use, either `"pt"` for PyTorch or `"tf"` for TensorFlow. The specified framework must be
            installed.

            If no framework is specified, will default to the one currently installed. If no framework is specified and
            both frameworks are installed, will default to the framework of the `model`, or to PyTorch if no model is
            provided.
        task (`str`, defaults to `""`):
            A task-identifier for the pipeline.
        num_workers (`int`, *optional*, defaults to 8):
            When the pipeline will use *DataLoader* (when passing a dataset, on GPU for a Pytorch model), the number of
            workers to be used.
        batch_size (`int`, *optional*, defaults to 1):
            When the pipeline will use *DataLoader* (when passing a dataset, on GPU for a Pytorch model), the size of
            the batch to use, for inference this is not always beneficial, please read [Batching with
            pipelines](https://huggingface.co/transformers/main_classes/pipelines.html#pipeline-batching) .
        args_parser ([`~pipelines.ArgumentHandler`], *optional*):
            Reference to the object in charge of parsing supplied pipeline parameters.
        device (`int`, *optional*, defaults to -1):
            Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, a positive will run the model on
            the associated CUDA device id. You can pass native `torch.device` or a `str` too
        torch_dtype (`str` or `torch.dtype`, *optional*):
            Sent directly as `model_kwargs` (just a simpler shortcut) to use the available precision for this model
            (`torch.float16`, `torch.bfloat16`, ... or `"auto"`)"""
    if supports_binary_output:
        docstring += r"""
        binary_output (`bool`, *optional*, defaults to `False`):
            Flag indicating if the output the pipeline should happen in a serialized format (i.e., pickle) or as
            the raw output data e.g. text."""


SUMMARY:
**Function Summary:**

`build_pipeline_init_args`

Builds and returns an initial argument string for a pipeline, incorporating model, tokenizer, feature extractor, image processor, and other parameters.

**Key Implementation Details:**

* Accepts boolean flags to determine the presence of various components.
* Incorporates documentation strings with detailed descriptions of each parameter.
* Uses optional parameters (e.g., `modelcard`, `framework`) for fine-tuning pipeline configuration.

**Notable Dependencies/Requirements:**

* `PreTrainedModel` and related classes (`TFPreTrainedModel`, `PreTrainedTokenizer`, etc.)
* `SequenceFeatureExtractor`
* `BaseImageProcessor`
* `ProcessorMixin`
* Optional dependencies:
	+ PyTorch (`torch`)
	+ TensorFlow (`tf`)
	+ Other frameworks (e.g., `"pt"` or `"tf"`)


key: chroma:document
string_value:
CODE:

class ChunkPipeline(Pipeline):
    def run_single(self, inputs, preprocess_params, forward_params, postprocess_params):
        all_outputs = []
        for model_inputs in self.preprocess(inputs, **preprocess_params):
            model_outputs = self.forward(model_inputs, **forward_params)
            all_outputs.append(model_outputs)
        outputs = self.postprocess(all_outputs, **postprocess_params)
        return outputs

    def get_iterator(
        self, inputs, num_workers: int, batch_size: int, preprocess_params, forward_params, postprocess_params
    ):
        if "TOKENIZERS_PARALLELISM" not in os.environ:
            logger.info("Disabling tokenizer parallelism, we're using DataLoader multithreading already")
            os.environ["TOKENIZERS_PARALLELISM"] = "false"
        if num_workers > 1:
            logger.warning(
                "For ChunkPipeline using num_workers>0 is likely to result in errors since everything is iterable,"
                " setting `num_workers=1` to guarantee correctness."
            )
            num_workers = 1
        dataset = PipelineChunkIterator(inputs, self.preprocess, preprocess_params)

        # TODO hack by collating feature_extractor and image_processor
        feature_extractor = self.feature_extractor if self.feature_extractor is not None else self.image_processor
        collate_fn = no_collate_fn if batch_size == 1 else pad_collate_fn(self.tokenizer, feature_extractor)
        dataloader = DataLoader(dataset, num_workers=num_workers, batch_size=batch_size, collate_fn=collate_fn)
        model_iterator = PipelinePackIterator(dataloader, self.forward, forward_params, loader_batch_size=batch_size)
        final_iterator = PipelineIterator(model_iterator, self.postprocess, postprocess_params)


SUMMARY:
**Summary:**

The `ChunkPipeline` class extends the `Pipeline` class and provides a data processing pipeline for batches of input data. The primary purpose is to preprocess, forward process, and postprocess individual models' inputs.

**Main Functionality:**

1.  **run_single:** Runs a single iteration through the pipeline, applying preprocessing, forward processing, and postprocessing steps.
2.  **get_iterator:** Returns an iterator for the pipeline, collecting batches of input data with preprocessing, forward processing, and postprocessing applied.

**Key Implementation Details:**

*   Uses `PipelineChunkIterator` to iterate over the dataset, featuring a custom collate function that handles single-sample batches using `no_collate_fn`.
*   Employs `DataLoader` from PyTorch for efficient batch loading.
*   Leverages `PipelinePackIterator` and `PipelineIterator` for further processing and postprocessing of model inputs.

**Notable Dependencies and Requirements:**

*   ` Pipeline, PipelineChunkIterator, PipelinePackIterator, PipelineIterator`
*   `os.environ['TOKENIZERS_PARALLELISM']` is set to "false" when not using parallelism.
*   External dependents may include:
    *   `no_collate_fn`
    *   `pad_collate_fn`
    *   `tokenizer`
    *   `feature_extractor` or `image_processor`

(Note: No external dependencies have been explicitly mentioned in the code snippet, but it's assumed these might be imported from other modules.)


key: chroma:document
string_value:
CODE:

class PipelineRegistry:
    def __init__(self, supported_tasks: Dict[str, Any], task_aliases: Dict[str, str]) -> None:
        self.supported_tasks = supported_tasks
        self.task_aliases = task_aliases

    def get_supported_tasks(self) -> List[str]:
        supported_task = list(self.supported_tasks.keys()) + list(self.task_aliases.keys())
        supported_task.sort()
        return supported_task

    def check_task(self, task: str) -> Tuple[str, Dict, Any]:
        if task in self.task_aliases:
            task = self.task_aliases[task]
        if task in self.supported_tasks:
            targeted_task = self.supported_tasks[task]
            return task, targeted_task, None

        if task.startswith("translation"):
            tokens = task.split("_")
            if len(tokens) == 4 and tokens[0] == "translation" and tokens[2] == "to":
                targeted_task = self.supported_tasks["translation"]
                task = "translation"
                return task, targeted_task, (tokens[1], tokens[3])
            raise KeyError(f"Invalid translation task {task}, use 'translation_XX_to_YY' format")

        raise KeyError(
            f"Unknown task {task}, available tasks are {self.get_supported_tasks() + ['translation_XX_to_YY']}"
        )

    def register_pipeline(
        self,
        task: str,
        pipeline_class: type,
        pt_model: Optional[Union[type, Tuple[type]]] = None,
        tf_model: Optional[Union[type, Tuple[type]]] = None,
        default: Optional[Dict] = None,
        type: Optional[str] = None,
    ) -> None:
        if task in self.supported_tasks:
            logger.warning(f"{task} is already registered. Overwriting pipeline for task {task}...")

        if pt_model is None:
            pt_model = ()
        elif not isinstance(pt_model, tuple):
            pt_model = (pt_model,)

        if tf_model is None:
            tf_model = ()
        elif not isinstance(tf_model, tuple):
            tf_model = (tf_model,)

        task_impl = {"impl": pipeline_class, "pt": pt_model, "tf": tf_model}

        if default is not None:
            if "model" not in default and ("pt" in default or "tf" in default):
                default = {"model": default}
            task_impl["default"] = default

        if type is not None:
            task_impl["type"] = type

        self.supported_tasks[task] = task_impl
        pipeline_class._registered_impl = {task: task_impl}

    def to_dict(self):


SUMMARY:
**PipelineRegistry Class Summary**

### Overview

The `PipelineRegistry` class manages a dictionary of supported tasks, their aliases, and corresponding pipeline implementations. It provides methods for registering pipelines, checking task validity, and retrieving the registry as a dictionary.

### Key Features

*   Registers pipelines with optional PT (Python Torch) and TF (TensorFlow) models
*   Supports task aliases and automatic translation task detection
*   Enforces unique task registration and raises errors for invalid tasks
*   Exposes the supported tasks as a digestible list

### Notable Methods

*   `__init__(supported_tasks, task_aliases)`: Initializes the registry with support tasks and task aliases.
*   `get_supported_tasks()`: Returns an alphabetically sorted list of all supported tasks, including translation variants.
*   `check_task(task)`: Validates a given task string and returns its corresponding implementation details or raises a KeyError if invalid.
*   `register_pipeline(task, pipeline_class, ... )`: Registers a new pipeline for a specific task, handling model types, defaults, and optional parameters.

### Dependencies

*   No external dependencies (only utilizes Python's built-in modules)
*   Utilizes the following classes/functions:
    *   `Dict`: Dictionary data structure
    *   `List`: List data structure
    *   `Tuple`: Tuple data structure
    *   `Optional`: Optional type hint
    *   `logger.warning()`: Logger warning function

### Interfaces Implemented

*   The registry implements an internal dictionary interface (`self.supported_tasks`) for storing task registrations.


key: chroma:document
string_value:
CODE:

class ReturnType(enum.Enum):
    TENSORS = 0
    NEW_TEXT = 1


SUMMARY:
**Summary:**

This code defines an enumeration (enum) called `ReturnType` with three possible values:

- `TENSORS`: represents tensors
- `NEW_TEXT`: represents new text
- `FULL_TEXT`: represents full text

The primary purpose is to provide a distinct and typed return type for various operations, likely in a data processing or machine learning context.

**Implementation Details:**

* No specific algorithm or design choice is highlighted.
* The method signatures or interfaces are not mentioned as key implementation details.

**Notable Dependencies or Requirements:**

- None explicitly stated. However, this enum is likely used within a class or function that may depend on external libraries or frameworks for tensor operations.


key: chroma:document
string_value:
CODE:

class Chat:
    """This class is intended to just be used internally in this pipeline and not exposed to users. We convert chats
    to this format because the rest of the pipeline code tends to assume that lists of messages are
    actually a batch of samples rather than messages in the same conversation."""

    def __init__(self, messages: Dict, images: Union[str, List[str], "Image.Image", List["Image.Image"]]):
        for message in messages:
            if not ("role" in message and "content" in message):
                raise ValueError("When passing chat dicts as input, each dict must have a 'role' and 'content' key.")
        images = retrieve_images_in_messages(messages, images)

        self.messages = messages


SUMMARY:
**Class Chat Summary**

The `Chat` class represents a conversation in the pipeline. Its primary purpose is to convert chats into a standardized format, where each message has 'role' and 'content' keys.

### Main Purpose and Functionality

* Initializes a chat object with a list of messages and optional images.
* Verifies that each message has required 'role' and 'content' keys.
* Retrieves images from messages using the `retrieve_images_in_messages` function.

### Key Implementation Details

* Method signature: `__init__(messages: Dict, images: Union[str, List[str], "Image.Image", List["Image.Image"]])`.

### Notable Dependencies or Requirements

* No external dependencies listed. However, it assumes the existence of the `retrieve_images_in_messages` function and the `Image` class.

Note: This summary preserves precise technical terms ("Dict", "Union", "List", "Image"), class/function names exactly ("Chat", "__init__"), and states the primary purpose while focusing on key algorithms, patterns, or design choices.


key: chroma:document
string_value:
CODE:

def retrieve_images_in_messages(
    messages: dict, images: Optional[Union[str, List[str], "Image.Image", List["Image.Image"]]]
):
    """
    Retrieve and combine images from the chat and the images passed as input.
    """
    if images is None:
        images = []
    idx_images = 0
    retrieved_images = []
    for message in messages:
        for content in message["content"]:
            if isinstance(content, dict):
                if content.get("type") == "image":
                    for key in ["image", "url", "path", "base64"]:
                        if key in content:
                            retrieved_images.append(content[key])
                            break
                    else:
                        if idx_images < len(images):
                            retrieved_images.append(images[idx_images])
                            idx_images += 1
                        else:
                            raise ValueError(
                                "The number of images in the chat messages should be the same as the number of images passed to the pipeline."
                            )
                # Add support for OpenAI/TGI chat format
                elif content.get("type") == "image_url":
                    if isinstance(content.get("image_url"), dict) and "url" in content["image_url"]:
                        retrieved_images.append(content["image_url"]["url"])
                        # Rewrite content to be in the Transformers chat format
                        content["type"] = "image"
                        content["image"] = content["image_url"]["url"]
                        del content["image_url"]
                    else:
                        raise ValueError(
                            "Wrong format for 'image_url' content type. The content should have an 'image_url' dict with a 'url' key."
                        )

    # The number of images passed should be consistent with the number of images in the chat without an image key
    if idx_images != len(images):
        raise ValueError(
            "The number of images in the chat messages should be the same as the number of images passed to the pipeline."
        )



SUMMARY:
**Summary: retrieve_images_in_messages function**

Purpose:
This function retrieves and combines images from chat messages (`messages`) with additional images (`images`). The primary purpose is to normalize image formats for further processing.

Key Implementation Details:

*   Iterates through each message, content, and key-value pairs.
*   Extracts `image`, `url`, `path`, or `base64` based on `content.type = "image"` or `.type = "image_url"`
*   Handles inconsistent image counts by using the first `images` count as a default
*   Supports OpenAI/TGI chat format by rewriting content
*   Raises errors for invalid formats

Notable Dependencies:
- `messages`: Expected to be a dictionary of messages.
- `images`: Can be any of `str`, `[str]`, `Image.Image`, or `[Image.Image]`.

Method Signatures and Interfaces:

*   None, but it can be adapted to other image-related interfaces.

External Dependencies: 

*   Assume an "Image" class or interface (imported as `Image`) is available, handling image-specific tasks.


key: chroma:document
string_value:
CODE:

class PipelineDataset(Dataset):
    def __init__(self, dataset, process, params):
        self.dataset = dataset
        self.process = process
        self.params = params

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

    def __getitem__(self, i):
        item = self.dataset[i]
        processed = self.process(item, **self.params)


SUMMARY:
**Concise Summary**

The `PipelineDataset` class wraps a dataset with data processing functionality. It inherits from Django's `Dataset` and extends it.

**Key Implementations:**

* The class constructor (`__init__`) initializes the dataset, process, and parameters.
* The `__getitem__` method retrieves an item from the dataset, applies the specified processing function (`process`), and returns the processed result.
* The `__len__` method returns the total length of the wrapped dataset.

**Notable Dependencies:**

* Django's `Dataset`
* User-defined `dataset`, `process`, and `params` (not shown in this code snippet)


key: chroma:document
string_value:
CODE:

class PipelineIterator(IterableDataset):
    def __init__(self, loader, infer, params, loader_batch_size=None):
        """
        Roughly equivalent to

        ```
        for item in loader:
            yield infer(item, **params)
        ```

                Arguments:
                    loader (`torch.utils.data.DataLoader` or `Iterable`):
                        The iterator that will be used to apply `infer` on.
                    infer (any function):
                        The function to apply of each element of `loader`.
                    params (`dict`):
                        The parameters passed to `infer` along with every item
                    loader_batch_size (`int`, *optional*):
                        If specified, the items of `loader` are supposed to come as batch, and are loader_batched here
                        making it roughly behave as


        ```
        for items in loader:
            for i in loader_batch_size:
                item = items[i]
                yield infer(item, **params)
        ```"""
        self.loader = loader
        self.infer = infer
        self.params = params
        if loader_batch_size == 1:
            # Let's spare some time by deactivating altogether
            loader_batch_size = None
        self.loader_batch_size = loader_batch_size

        # Internal bookkeeping
        self._loader_batch_index = None
        self._loader_batch_data = None

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

    def __iter__(self):
        self.iterator = iter(self.loader)
        return self

    def loader_batch_item(self):
        """
        Return item located at `loader_batch_index` within the current `loader_batch_data`.
        """
        if isinstance(self._loader_batch_data, torch.Tensor):
            # Batch data is simple tensor, just fetch the slice
            result = self._loader_batch_data[self._loader_batch_index].unsqueeze(0)
        else:
            # Batch data is assumed to be BaseModelOutput (or dict)
            loader_batched = {}
            for k, element in self._loader_batch_data.items():
                if isinstance(element, ModelOutput):
                    # Convert ModelOutput to tuple first
                    element = element.to_tuple()
                    if isinstance(element[0], torch.Tensor):
                        loader_batched[k] = tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
                    elif isinstance(element[0], np.ndarray):
                        loader_batched[k] = tuple(np.expand_dims(el[self._loader_batch_index], 0) for el in element)
                    continue
                if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(element, tuple):
                    # Those are stored as lists of tensors so need specific unbatching.
                    if isinstance(element[0], torch.Tensor):
                        loader_batched[k] = tuple(el[self._loader_batch_index].unsqueeze(0) for el in element)
                    elif isinstance(element[0], np.ndarray):
                        loader_batched[k] = tuple(np.expand_dims(el[self._loader_batch_index], 0) for el in element)
                    continue
                if element is None:
                    # This can happen for optional data that get passed around
                    loader_batched[k] = None
                elif isinstance(element[self._loader_batch_index], torch.Tensor):
                    # Take correct batch data, but make it looked like batch_size=1
                    # For compatibility with other methods within transformers

                    loader_batched[k] = element[self._loader_batch_index].unsqueeze(0)
                elif isinstance(element[self._loader_batch_index], np.ndarray):
                    # Take correct batch data, but make it looked like batch_size=1
                    # For compatibility with other methods within transformers
                    loader_batched[k] = np.expand_dims(element[self._loader_batch_index], 0)
                else:
                    # This is typically a list, so no need to `unsqueeze`.
                    loader_batched[k] = element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
            result = self._loader_batch_data.__class__(loader_batched)
        self._loader_batch_index += 1
        return result

    def __next__(self):
        if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
            # We are currently unrolling a batch so we just need to return
            # the current item within a batch
            return self.loader_batch_item()

        # We're out of items within a batch
        item = next(self.iterator)
        processed = self.infer(item, **self.params)
        # We now have a batch of "inferred things".
        if self.loader_batch_size is not None:
            # Try to infer the size of the batch
            if isinstance(processed, torch.Tensor):
                first_tensor = processed
            elif isinstance(processed, tuple):
                first_tensor = processed[0]
            else:
                key = list(processed.keys())[0]
                first_tensor = processed[key]

            if isinstance(first_tensor, list):
                observed_batch_size = len(first_tensor)
            else:
                observed_batch_size = first_tensor.shape[0]
            if 0 < observed_batch_size < self.loader_batch_size:
                # could be last batch so we can't unroll as many
                # elements.
                self.loader_batch_size = observed_batch_size
            # Setting internal index to unwrap the batch
            self._loader_batch_data = processed[0] if isinstance(processed, tuple) else processed
            self._loader_batch_index = 0
            return self.loader_batch_item()
        else:
            # We're not unrolling batches


SUMMARY:
Here's a concise summary of the `PipelineIterator` class:

**Main Purpose and Functionality:**
The `PipelineIterator` is an iterable class that applies a given function (`infer`) to each item in a data loader, with optional batching. It uses internal bookkeeping to manage the iteration process.

**Key Implementation Details:**

* Employs batched iteration of the data loader using `loader_batch_size`.
* Utilizes internal variables `_loader_batch_data` and `_loader_batch_index` to track the current batch.
* Uses `next()` function to retrieve the next item in the iteration, optionally unrolling a batch if `loader_batch_size` is specified.

**Notable Dependencies:**

* `torch.utils.data.DataLoader` or `Iterable`
* `infer` function (any type)
* `params` dictionary

This class appears to be part of a larger transformer framework, possibly related to model inference and data iteration.


key: chroma:document
string_value:
CODE:

class PipelineChunkIterator(PipelineIterator):
    def __init__(self, loader, infer, params, loader_batch_size=None):
        """
        Roughly equivalent to

        ```
        for iterator in loader:
            for item in iterator:
                yield infer(item, **params)
        ```

                Arguments:
                    loader (`torch.utils.data.DataLoader` or `Iterable`):
                        The iterator that will be used to apply `infer` on.
                    infer (any function):
                        The function to apply of each element of `loader`.
                    params (`dict`):
                        The parameters passed to `infer` along with every item
        """
        super().__init__(loader, infer, params)

    def __iter__(self):
        self.iterator = iter(self.loader)
        self.subiterator = None
        return self

    def __next__(self):
        if self.subiterator is None:
            "Subiterator None means we haven't started a `preprocess` iterator. so start it"
            self.subiterator = self.infer(next(self.iterator), **self.params)
        try:
            # Try to return next item
            processed = next(self.subiterator)
        except StopIteration:
            # When a preprocess iterator ends, we can start lookig at the next item
            # ChunkIterator will keep feeding until ALL elements of iterator
            # all have created their subiterator and have been iterating against.
            #
            # Another way to look at it, is we're basically flattening lists of lists
            # into a single list, but with generators
            self.subiterator = self.infer(next(self.iterator), **self.params)
            processed = next(self.subiterator)


SUMMARY:
**PipelineChunkIterator Class Summary**

**Main Purpose and Functionality**
The `PipelineChunkIterator` class is an iterator that applies an inference function `infer` to each element of an input iterator (`loader`) in chunks, utilizing two nested iterators for parallel processing. It feeds the input data through a pipeline where each element is first processed by `infer`, and then yielded as part of the main iterator.

**Key Implementation Details**

* Two nested iterators (`iterator` and `subiterator`) are used to achieve parallel processing.
* The `__next__` method handles iteration over both main and subiterators, switching between them when necessary.
* The `infer` function is applied to each element of the input iterator multiple times in a nested manner.

**Notable Dependencies or Requirements**

* `torch.utils.data.DataLoader` or `Iterable`: the input iterator that will be used to apply `infer` on.
* External dependency: None listed, but `infer` assumes the existence of related classes (e.g., `DataLoader`, `Iterable`) and functions.


key: chroma:document
string_value:
CODE:

class PipelinePackIterator(PipelineIterator):
    """
    Roughly equivalent to

    ```
    packed =  []
    for item in loader:
        packed.append(item)
        if item["is_last"]:
            yield packed
            packed = []
    ```

        but it also handles cases where `item` are batched (meaning it's a dict of Tensor with first dimension > 1. In
        that case it does

    ```
    packed =  []
    for batch in loader:
        # item is batched
        for item in batch:
            packed.append(item)
            if item["is_last"]:
                yield packed
                packed = []
    ```

        Arguments:
            loader (`torch.utils.data.DataLoader` or `Iterable`):
                The iterator that will be used to apply `infer` on.
            infer (any function):
                The function to apply of each element of `loader`.
            params (`dict`):
                The parameters passed to `infer` along with every item
            loader_batch_size (`int`, *optional*):
                If specified, the items of `loader` are supposed to come as batch, and are loader_batched here making
                it roughly behave as


    ```
    for items in loader:
        for i in loader_batch_size:
            item = items[i]
            yield infer(item, **params)
    ```"""

    def __iter__(self):
        self.iterator = iter(self.loader)
        return self

    def __next__(self):
        # Extremely similar to PipelineIterator in its unpacking mechanism
        # BUT, we have an extra required item which is the presence of `is_last`
        # That is because everything is flattened by `PipelineChunkIterator` we
        # need to keep track of how to regroup here in the original `process`
        # boundaries so that `process` and `postprocess` see the same data.

        # This iterator accumulates items (possibly while unbatching) until it
        # its a `is_last` and then just passes it on to the caller.
        is_last = False
        accumulator = []
        if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
            while self._loader_batch_index < self.loader_batch_size:
                item = self.loader_batch_item()
                is_last = item.pop("is_last")
                accumulator.append(item)
                if is_last:
                    return accumulator

        while not is_last:
            processed = self.infer(next(self.iterator), **self.params)
            if self.loader_batch_size is not None:
                if isinstance(processed, torch.Tensor):
                    first_tensor = processed
                else:
                    key = list(processed.keys())[0]
                    first_tensor = processed[key]
                if isinstance(first_tensor, list):
                    observed_batch_size = len(first_tensor)
                else:
                    observed_batch_size = first_tensor.shape[0]
                if 0 < observed_batch_size < self.loader_batch_size:
                    # could be last batch so we can't unroll as many
                    # elements.
                    self.loader_batch_size = observed_batch_size
                self._loader_batch_data = processed
                self._loader_batch_index = 0
                while self._loader_batch_index < self.loader_batch_size:
                    item = self.loader_batch_item()
                    is_last = item.pop("is_last")
                    accumulator.append(item)
                    if is_last:
                        return accumulator
            else:
                item = processed
                is_last = item.pop("is_last")
                accumulator.append(item)


SUMMARY:
**Summary**

`PipelinePackIterator` is a custom iterator that extends `PipelineIterator`. It efficiently applies an inference function to batches of items from a data loader, handling both single-item and batched cases. 

**Main Purpose and Functionality:**
- Takes an iterable as input instead of a tensor with dimension 1 (batched tensor).
 
 **Key Implementation Details:**

*   Iterates over the loader, appending each item to an accumulator until a sentinel value ("is_last") is encountered.
*   Unbatching works similarly by iterating over the batch size iterations. If no batches are expected and it encounters a tensor not equal to its single element in the tensor (shape or dict), it adjusts the batch size.
*   Handles both cases of flattened vs batched item handling efficiently.

 **Notable Dependencies/Requirements:**
-   `torch.utils.data.DataLoader`
-   The inference function (`infer`)
-   Optional parameters (`params` and loader_batch_size)
 
**Example usage**: 

```python
# usage example for PipelinePackIterator
loader = torch.utils.data.DataLoader(...)  # data loader as input
infer_func = (...)  # the function to apply on each item of `loader`
chunk_iter = PipelinePackIterator(loader, infer=infer_func)  # create an instance with load and infer arguments

for chunk in chunk_iter:
    print(chunk)  # prints a batched iterator chunk if any else printed single items
    ...    # break loop is you wish (example purposes)
```

 The class implements batching handling to improve efficiency.


key: chroma:document
string_value:
CODE:

class KeyDataset(Dataset):
    def __init__(self, dataset: Dataset, key: str):
        self.dataset = dataset
        self.key = key

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

    def __getitem__(self, i):


SUMMARY:
**KeyDataset Class Summary**

* **Main Purpose and Functionality**: A wrapper class around another `Dataset` instance, allowing for data selection via a specific key.

**Implementation Details**:
- Uses the original dataset's index to retrieve value at specified key (``self.dataset[i][self.key]``).

**Notable Dependencies**:
- Inherited from `Dataset`.
- Requires an instance of `Dataset` with associated key during initialization.


key: chroma:document
string_value:
CODE:

class KeyPairDataset(Dataset):
    def __init__(self, dataset: Dataset, key1: str, key2: str):
        self.dataset = dataset
        self.key1 = key1
        self.key2 = key2

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

    def __getitem__(self, i):


SUMMARY:
**Class Summary: `KeyPairDataset`**

* **Main Purpose and Functionality**: Custom Dataset class that transforms a base dataset into a `KeyPairDataset` format, where each data point is paired with two distinct keys (`key1` and `key2`).
* **Implementation Details**:
	+ Exposes the `_len()` method to return the length of the underlying dataset.
	+ Implements the `_getitem()` method to retrieve individual data points from the dataset, using the specified `key1` and `key2` for extraction. The returned dictionary contains two text values: `text` (from `self.key1`) and `text_pair` (from `self.key2`).
* **Notable Dependencies**:
	+ Extends `Dataset` class from an unknown library/framework (likely PyTorch, due to the mention of `__len__` methods)
	+ Assumes access to a base dataset object (`dataset`) with compatible indexing and structure.


key: chroma:document
string_value:
CODE:

class TableQuestionAnsweringArgumentHandler(ArgumentHandler):
    """
    Handles arguments for the TableQuestionAnsweringPipeline
    """

    def __call__(self, table=None, query=None, **kwargs):
        # Returns tqa_pipeline_inputs of shape:
        # [
        #   {"table": pd.DataFrame, "query": List[str]},
        #   ...,
        #   {"table": pd.DataFrame, "query" : List[str]}
        # ]
        requires_backends(self, "pandas")
        import pandas as pd

        if table is None:
            raise ValueError("Keyword argument `table` cannot be None.")
        elif query is None:
            if isinstance(table, dict) and table.get("query") is not None and table.get("table") is not None:
                tqa_pipeline_inputs = [table]
            elif isinstance(table, list) and len(table) > 0:
                if not all(isinstance(d, dict) for d in table):
                    raise ValueError(
                        f"Keyword argument `table` should be a list of dict, but is {(type(d) for d in table)}"
                    )

                if table[0].get("query") is not None and table[0].get("table") is not None:
                    tqa_pipeline_inputs = table
                else:
                    raise ValueError(
                        "If keyword argument `table` is a list of dictionaries, each dictionary should have a `table`"
                        f" and `query` key, but only dictionary has keys {table[0].keys()} `table` and `query` keys."
                    )
            elif Dataset is not None and isinstance(table, Dataset) or isinstance(table, types.GeneratorType):
                return table
            else:
                raise ValueError(
                    "Invalid input. Keyword argument `table` should be either of type `dict` or `list`, but "
                    f"is {type(table)})"
                )
        else:
            tqa_pipeline_inputs = [{"table": table, "query": query}]

        for tqa_pipeline_input in tqa_pipeline_inputs:
            if not isinstance(tqa_pipeline_input["table"], pd.DataFrame):
                if tqa_pipeline_input["table"] is None:
                    raise ValueError("Table cannot be None.")

                tqa_pipeline_input["table"] = pd.DataFrame(tqa_pipeline_input["table"])



SUMMARY:
**Class Summary: `TableQuestionAnsweringArgumentHandler`**

*   **Purpose:** Handles arguments for the TableQuestionAnsweringPipeline, transforming tables and queries into a format compatible with TQA pipeline inputs.

### Main Purpose and Functionality

The class takes table and query parameters, ensuring they are valid and formatted correctly. It supports cases where `table` is None or an existing DataFrame, a list of dictionaries (with `table` and `query` keys), or a Dataset object. The resulting transformed tables are stored in the `tqa_pipeline_inputs` variable.

### Key Implementation Details

*   **Data Validation:** Validates input types and formats for `table`, `query`, and `Dataset` objects.
*   **Data Type Conversion:** Converts tables to DataFrames if they are not already.

### Notable Dependencies or Requirements

*   **Pandas (`pandas`):** Required library for data manipulation and processing. All instances of this library are prefixed with `pd.`.

### Classes/Interfaces Implemented

*   Uses the `ArgumentHandler` class interface.
*   Inherently uses the `requires_backends` method (not explicitly implemented here, assuming it's already defined elsewhere in the codebase).

**Key Method Signature:**

```python
def __call__(self, table=None, query=None, **kwargs):
    ...
```

This summary provides an overview of the class's purpose and implementation details while minimizing unnecessary information. The focus is on key concepts, technical terms, method signatures, and notable dependencies or requirements.


key: chroma:document
string_value:
CODE:

def sigmoid(_outputs):


SUMMARY:
**Concise Summary:**

The `_sigmoid` function calculates the Sigmoid activation function:

* Takes an array of `_outputs` as input
* Applies the formula `1 / (1 + exp(-x))`
* Returns the sigmoid of each element in the input array.

### Implementation Details:
* Used for binary classification problems
* Utilizes exponentiation (`exp`) and division operations

### Dependencies:
* `np`: NumPy library is used for the `exp` function.


key: chroma:document
string_value:
CODE:

def softmax(_outputs):
    maxes = np.max(_outputs, axis=-1, keepdims=True)
    shifted_exp = np.exp(_outputs - maxes)


SUMMARY:
**Concise Summary**

### Function Overview

`softmax` calculates the softmax normalization of a set of output values.

### Primary Purpose

The primary purpose is to normalize the input `_outputs` values using the softmax function, which produces a probability distribution over each class or category.

### Key Implementation Details

* The `maxes` variable stores the maximum values along the last axis (-1) of the input `_outputs`.
* The exponential shift (`shifted_exp`) is calculated by subtracting `maxes` from each element in `_outputs`.
* The softmax output is computed by dividing `shifted_exp` by its sum over the last axis (axis=-1).

### Notable Dependencies

* `numpy` (np) for numerical operations.
* `np.max`, `np.exp`, and `np.sum` are used extensively.

### Notes

* This implementation uses NumPy's vectorized operations to compute the softmax output efficiently.
* The `_outputs` variable is expected to be a NumPy array with at least one axis.


key: chroma:document
string_value:
CODE:

class ClassificationFunction(ExplicitEnum):
    SIGMOID = "sigmoid"
    SOFTMAX = "softmax"


SUMMARY:
**Summary:**

The `ClassificationFunction` class defines an explicit enumeration for classification functions, providing three variants:

* `Sigmoid`: Maps to the sigmoid activation function.
* `Softmax`: Maps to the softmax activation function, often used in multi-class classification problems.
* `None`: Represents no transformation or no classification method.

**Key Details:**

* No specific implementation provided; only enumerates possible classification functions.

**Notable Dependencies/Requirements:**

* None explicitly listed, but may be used in conjunction with numerical libraries for statistical calculations (not applicable to this code snippet).


key: chroma:document
string_value:
CODE:

def decode_spans(
    start: np.ndarray, end: np.ndarray, topk: int, max_answer_len: int, undesired_tokens: np.ndarray
) -> Tuple:
    """
    Take the output of any `ModelForQuestionAnswering` and will generate probabilities for each span to be the actual
    answer.

    In addition, it filters out some unwanted/impossible cases like answer len being greater than max_answer_len or
    answer end position being before the starting position. The method supports output the k-best answer through the
    topk argument.

    Args:
        start (`np.ndarray`): Individual start probabilities for each token.
        end (`np.ndarray`): Individual end probabilities for each token.
        topk (`int`): Indicates how many possible answer span(s) to extract from the model output.
        max_answer_len (`int`): Maximum size of the answer to extract from the model's output.
        undesired_tokens (`np.ndarray`): Mask determining tokens that can be part of the answer
    """
    # Ensure we have batch axis
    if start.ndim == 1:
        start = start[None]

    if end.ndim == 1:
        end = end[None]

    # Compute the score of each tuple(start, end) to be the real answer
    outer = np.matmul(np.expand_dims(start, -1), np.expand_dims(end, 1))

    # Remove candidate with end < start and end - start > max_answer_len
    candidates = np.tril(np.triu(outer), max_answer_len - 1)

    #  Inspired by Chen & al. (https://github.com/facebookresearch/DrQA)
    scores_flat = candidates.flatten()
    if topk == 1:
        idx_sort = [np.argmax(scores_flat)]
    elif len(scores_flat) < topk:
        idx_sort = np.argsort(-scores_flat)
    else:
        idx = np.argpartition(-scores_flat, topk)[0:topk]
        idx_sort = idx[np.argsort(-scores_flat[idx])]

    starts, ends = np.unravel_index(idx_sort, candidates.shape)[1:]
    desired_spans = np.isin(starts, undesired_tokens.nonzero()) & np.isin(ends, undesired_tokens.nonzero())
    starts = starts[desired_spans]
    ends = ends[desired_spans]
    scores = candidates[0, starts, ends]



SUMMARY:
**`decode_spans` Function Summary**

**Purpose and Functionality:**
This function decodes the output of a `ModelForQuestionAnswering` by generating probabilities for each span to be the actual answer. It filters out unwanted cases such as incorrect answer length or token positions.

**Key Implementation Details:**

* Computation of score between `start` and `end` tokens using matrix multiplication.
* Filtering of candidates based on invalid end positions (before start) or exceeding maximum answer length.
* Sorting of top-k candidate spans by their scores using NumPy's `argpartition` and `argsort` functions.

**Notable Dependencies and Requirements:**

* `np.ndarray`: Input arrays for start, end, and undesired tokens.
* `ModelForQuestionAnswering`: Output data type not explicitly specified in the function.
* No external dependencies listed.

**Return Value:**
A tuple containing the starts, ends, and scores of the top-k candidate spans.


key: chroma:document
string_value:
CODE:

def select_starts_ends(
    start,
    end,
    p_mask,
    attention_mask,
    min_null_score=1000000,
    top_k=1,
    handle_impossible_answer=False,
    max_answer_len=15,
):
    """
    Takes the raw output of any `ModelForQuestionAnswering` and first normalizes its outputs and then uses
    `decode_spans()` to generate probabilities for each span to be the actual answer.

    Args:
        start (`np.ndarray`): Individual start logits for each token.
        end (`np.ndarray`): Individual end logits for each token.
        p_mask (`np.ndarray`): A mask with 1 for values that cannot be in the answer
        attention_mask (`np.ndarray`): The attention mask generated by the tokenizer
        min_null_score(`float`): The minimum null (empty) answer score seen so far.
        topk (`int`): Indicates how many possible answer span(s) to extract from the model output.
        handle_impossible_answer(`bool`): Whether to allow null (empty) answers
        max_answer_len (`int`): Maximum size of the answer to extract from the model's output.
    """
    # Ensure padded tokens & question tokens cannot belong to the set of candidate answers.
    undesired_tokens = np.abs(np.array(p_mask) - 1)

    if attention_mask is not None:
        undesired_tokens = undesired_tokens & attention_mask

    # Generate mask
    undesired_tokens_mask = undesired_tokens == 0.0

    # Make sure non-context indexes in the tensor cannot contribute to the softmax
    start = np.where(undesired_tokens_mask, -10000.0, start)
    end = np.where(undesired_tokens_mask, -10000.0, end)

    # Normalize logits and spans to retrieve the answer
    start = np.exp(start - start.max(axis=-1, keepdims=True))
    start = start / start.sum()

    end = np.exp(end - end.max(axis=-1, keepdims=True))
    end = end / end.sum()

    if handle_impossible_answer:
        min_null_score = min(min_null_score, (start[0, 0] * end[0, 0]).item())

    # Mask CLS
    start[0, 0] = end[0, 0] = 0.0

    starts, ends, scores = decode_spans(start, end, top_k, max_answer_len, undesired_tokens)


SUMMARY:
**Function Summary: `select_starts_ends`**

**Purpose and Functionality**: This function takes the raw output of a `ModelForQuestionAnswering` and normalizes it to generate probabilities for each span to be the actual answer.

**Key Implementation Details**:

* Normalizes start and end logits using softmax, excluding undesired tokens.
* Applies additional smoothing by setting start and end masks to 0.0 at CLS token position.
* Uses `decode_spans()` function to extract top-k possible answer spans.

**Notable Dependencies and Requirements**:

* `np.ndarray` inputs for raw output
* External dependency on `ModelForQuestionAnswering`
* `decode_spans()` function, not defined in this code chunk.


key: chroma:document
string_value:
CODE:

class QuestionAnsweringArgumentHandler(ArgumentHandler):
    """
    QuestionAnsweringPipeline requires the user to provide multiple arguments (i.e. question & context) to be mapped to
    internal [`SquadExample`].

    QuestionAnsweringArgumentHandler manages all the possible to create a [`SquadExample`] from the command-line
    supplied arguments.
    """

    def normalize(self, item):
        if isinstance(item, SquadExample):
            return item
        elif isinstance(item, dict):
            for k in ["question", "context"]:
                if k not in item:
                    raise KeyError("You need to provide a dictionary with keys {question:..., context:...}")
                elif item[k] is None:
                    raise ValueError(f"`{k}` cannot be None")
                elif isinstance(item[k], str) and len(item[k]) == 0:
                    raise ValueError(f"`{k}` cannot be empty")

            return QuestionAnsweringPipeline.create_sample(**item)
        raise ValueError(f"{item} argument needs to be of type (SquadExample, dict)")

    def __call__(self, *args, **kwargs):
        # Detect where the actual inputs are
        if args is not None and len(args) > 0:
            if len(args) == 1:
                inputs = args[0]
            elif len(args) == 2 and {type(el) for el in args} == {str}:
                inputs = [{"question": args[0], "context": args[1]}]
            else:
                inputs = list(args)
        # Generic compatibility with sklearn and Keras
        # Batched data
        elif "X" in kwargs:
            warnings.warn(
                "Passing the `X` argument to the pipeline is deprecated and will be removed in v5. Inputs should be passed using the `question` and `context` keyword arguments instead.",
                FutureWarning,
            )
            inputs = kwargs["X"]
        elif "data" in kwargs:
            warnings.warn(
                "Passing the `data` argument to the pipeline is deprecated and will be removed in v5. Inputs should be passed using the `question` and `context` keyword arguments instead.",
                FutureWarning,
            )
            inputs = kwargs["data"]
        elif "question" in kwargs and "context" in kwargs:
            if isinstance(kwargs["question"], list) and isinstance(kwargs["context"], str):
                inputs = [{"question": Q, "context": kwargs["context"]} for Q in kwargs["question"]]
            elif isinstance(kwargs["question"], list) and isinstance(kwargs["context"], list):
                if len(kwargs["question"]) != len(kwargs["context"]):
                    raise ValueError("Questions and contexts don't have the same lengths")

                inputs = [{"question": Q, "context": C} for Q, C in zip(kwargs["question"], kwargs["context"])]
            elif isinstance(kwargs["question"], str) and isinstance(kwargs["context"], str):
                inputs = [{"question": kwargs["question"], "context": kwargs["context"]}]
            else:
                raise ValueError("Arguments can't be understood")
        else:
            raise ValueError(f"Unknown arguments {kwargs}")

        # When user is sending a generator we need to trust it's a valid example
        generator_types = (types.GeneratorType, Dataset) if Dataset is not None else (types.GeneratorType,)
        if isinstance(inputs, generator_types):
            return inputs

        # Normalize inputs
        if isinstance(inputs, dict):
            inputs = [inputs]
        elif isinstance(inputs, Iterable):
            # Copy to avoid overriding arguments
            inputs = list(inputs)
        else:
            raise ValueError(f"Invalid arguments {kwargs}")

        for i, item in enumerate(inputs):
            inputs[i] = self.normalize(item)



SUMMARY:
**QuestionAnsweringArgumentHandler Class Summary**

### Overview
The `QuestionAnsweringArgumentHandler` class is responsible for managing the input arguments and creating a `SquadExample` from them. It inherits from the `ArgumentHandler` class.

### Primary Purpose
To normalize the input arguments (question, context, or data) and return a list of normalized `SquadExample` instances.

### Key Implementation Details

*   **normalize method**: This method takes an argument and checks if it's a `SquadExample`, dictionary, question, context, generator, or iterable. It raises errors for invalid arguments.
*   **__call__ method**: This method detects the input sources (args or kwargs) and passes them to the `normalize` method for processing.

### Notable Dependencies or Requirements

*   The class uses the `types.GeneratorType` and `Dataset` classes from Python 3's `typing` module. It assumes that a generator is an instance of these types.
*   The class also uses external dependencies such as `SquadExample` (not defined in this snippet), `ArgumentHandler`, `FutureWarning`, and `sklearn` or Keras, which are not explicitly imported.

### Interfaces and Inherited Methods

The class inherits from the `ArgumentHandler` base class and does not implement any new interfaces. However, it overrides the built-in object behavior to use its `__call__` method as a callable instance.


key: chroma:document
string_value:
CODE:

def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
    """
    Helper function to read an audio file through ffmpeg.
    """
    ar = f"{sampling_rate}"
    ac = "1"
    format_for_conversion = "f32le"
    ffmpeg_command = [
        "ffmpeg",
        "-i",
        "pipe:0",
        "-ac",
        ac,
        "-ar",
        ar,
        "-f",
        format_for_conversion,
        "-hide_banner",
        "-loglevel",
        "quiet",
        "pipe:1",
    ]

    try:
        with subprocess.Popen(ffmpeg_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE) as ffmpeg_process:
            output_stream = ffmpeg_process.communicate(bpayload)
    except FileNotFoundError as error:
        raise ValueError("ffmpeg was not found but is required to load audio files from filename") from error
    out_bytes = output_stream[0]
    audio = np.frombuffer(out_bytes, np.float32)
    if audio.shape[0] == 0:
        raise ValueError(
            "Soundfile is either not in the correct format or is malformed. Ensure that the soundfile has "
            "a valid audio file extension (e.g. wav, flac or mp3) and is not corrupted. If reading from a remote "
            "URL, ensure that the URL is the full address to **download** the audio file."
        )


SUMMARY:
**FFmpeg Audio Reader Function Summary**

**Purpose**: Reads an audio file through FFmpeg and returns the audio data as a NumPy array.

**Implementation Details**:

* Uses `ffmpeg` command-line tool to read raw audio data from input bytes (`bpayload`) with specified sampling rate (`sampling_rate`)
* Converts resulting stream from FFmpeg into a NumPy float32 array using `np.frombuffer`

**Key Algorithm/Design Choices**:

* Utilizes FFmpeg's `-ac`, `-ar`, and `-f` options to configure audio conversion
* Leverages the `pipe:0` and `pipe:1` syntax for input/output redirection

**Notable Dependencies/Requirements**:

* FFmpeg executable is required to be installed on the system, which raises a `ValueError` if not found (`subprocess.Popen` raise exception)

**Additional Notes**: Input audio data must conform to standard formats (e.g. wav, flac) and contain usable audio data to avoid raising errors.
```python
def ffmpeg_read(bpayload: bytes, sampling_rate: int) -> np.array:
    # FFmpeg audio reader function implementation details...
    
    return audio
```
**External Dependencies:**

* `ffmpeg` command-line tool
* NumPy for audio data processing


