Upload processor

preprocessor_config.json

@@ -0,0 +1,48 @@
+{
+  "_valid_processor_keys": [
+    "images",
+    "do_resize",
+    "size",
+    "resample",
+    "do_center_crop",
+    "crop_size",
+    "do_rescale",
+    "rescale_factor",
+    "do_normalize",
+    "image_mean",
+    "image_std",
+    "do_convert_rgb",
+    "return_tensors",
+    "data_format",
+    "input_data_format"
+  ],
+  "auto_map": {
+    "AutoProcessor": "processing_jat.JatProcessor"
+  },
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "processor_class": "JatProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 224
+  }
+}

processing_jat.py

@@ -0,0 +1,404 @@
+import copy
+import warnings
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+import torchvision.transforms.functional as F
+from transformers import BatchEncoding
+from transformers.processing_utils import ProcessorMixin
+
+
+def to_tensor(x):
+    """
+    Convert a nested structure of numpy arrays or tensors (including lists and tuples of them)
+    into a tensor. Assumes that all nested structures can be converted into a tensor directly.
+
+    :param x: Nested structure containing numpy arrays, tensors, lists, or tuples
+    :return: torch.Tensor
+    """
+    with warnings.catch_warnings():
+        # Convert specific warning to an error
+        warnings.filterwarnings(
+            "error",
+            category=UserWarning,
+            message=".*Creating a tensor from a list of numpy.ndarrays is extremely slow.*",
+        )
+        try:
+            return torch.Tensor(x)
+        except Exception:
+            if isinstance(x, list):
+                return torch.stack([to_tensor(item) for item in x])
+            else:
+                raise TypeError("Unsupported type for conversion to tensor")
+
+
+def truncate(
+    encoding: Dict[str, List[List[Any]]], max_length: int, truncation_side: str = "right", preserve: bool = False
+) -> Dict[str, List[List[Any]]]:
+    """
+    Truncate the sequences in the encoding to the specified maximum length.
+
+    This function is designed to process batch of sequences represented in the encoding dictionary.
+    Depending on the chosen strategy, sequences are either truncated with loss of residual data or with preservation
+    and incorporation of residual data into the batch.
+
+    Args:
+        encoding (`Mapping`):
+            A dictionary where each key-value pair consists of a feature name and its corresponding batch of sequences.
+            The sequences are expected to be lists.
+        max_length (`int`):
+            The maximum allowable length for the sequences.
+        truncation_side (`str`, **optional**):
+            The strategy to use for truncation. Can be `"left"` or `"right"`. Defaults to `"right"`.
+        preserve (`bool`, **optional**):
+            Whether to preserve the residual data by adding them as new sequences in the batch. Defaults to `False`.
+
+    Returns:
+        `Dict[str, List[List[Any]]]`:
+            A dictionary with the same keys as the input `encoding`, containing the truncated batch of sequences.
+            If `preserve` is set to `True`, the batch size may increase due to the addition of new sequences formed
+            from the residual data.
+
+    Example:
+
+        >>> encoding = {'feature1': [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]}
+        >>> truncate(encoding, 3, preserve=False)
+        {'feature1': [[1, 2, 3], [6, 7, 8]]}
+
+        >>> truncate(encoding, 3, preserve=True)
+        {'feature1': [[1, 2, 3], [4, 5], [6, 7, 8], [9, 10]]}
+    """
+    truncated_encoding = {}
+
+    for key, sequences in encoding.items():
+        if not all(isinstance(seq, list) for seq in sequences):
+            raise TypeError(f"All sequences under key {key} should be of type list.")
+
+        truncated_sequences = []
+
+        for seq in sequences:
+            if len(seq) <= max_length:
+                truncated_sequences.append(seq)
+                continue
+
+            if preserve:  # truncate and append the residual as new sequences
+                if truncation_side == "right":
+                    truncated_sequences.extend([seq[i : i + max_length] for i in range(0, len(seq), max_length)])
+                elif truncation_side == "left":
+                    n = len(seq) // max_length + int(len(seq) % max_length > 0)
+                    low, high = len(seq) - n * max_length, len(seq)
+                    truncated_sequences.extend(
+                        [seq[max(0, i - max_length) : i] for i in range(high, low, -max_length)]
+                    )
+                else:
+                    raise ValueError(f"Invalid truncation_side: {truncation_side}")
+            else:  # simply truncate the sequence
+                if truncation_side == "right":
+                    truncated_sequences.append(seq[:max_length])
+                elif truncation_side == "left":
+                    truncated_sequences.append(seq[-max_length:])
+
+        truncated_encoding[key] = truncated_sequences
+
+    return truncated_encoding
+
+
+def pad(encoding: Dict[str, List[List[Any]]], target_length: int) -> Dict[str, List[List[Any]]]:
+    """
+    Pad the sequences in the encoding to the specified maximum length.
+
+    This function is designed to process batch of sequences represented in the encoding dictionary.
+    The padding value is set to be the first element in the sequence.
+
+    Args:
+        encoding (`Mapping`):
+            A dictionary where each key-value pair consists of a feature name and its corresponding batch of sequences.
+            The sequences are expected to be lists.
+        target_length (`int`):
+            The desired length for the sequences.
+
+    Returns:
+        `Dict[str, List[List[Any]]]`:
+            A dictionary with the same keys as the input `encoding`, containing the padded batch of sequences.
+            An additional key `attention_mask` is added to the dictionary to indicate the positions of the non-padding
+            elements with 1s and the padding elements with 0s. If the input `encoding` already contains an
+            `attention_mask` key, the corresponding mask will be updated such that the original masking is preserved,
+            and the newly added padding elements will be masked with 0s. In other words, the resulting
+            `attention_mask` is a logical "AND" between the provided mask and the mask created due to padding, ensuring
+            that any element masked originally remains masked.
+
+    Example:
+
+        >>> encoding = {'feature1': [[1, 2], [3, 4, 5]]}
+        >>> pad(encoding, 4)
+        {'feature1': [[1, 2, 1, 1], [3, 4, 5, 3]], 'attention_mask': [[1, 1, 0, 0], [1, 1, 1, 0]]}
+
+        >>> encoding = {'feature1': [[1, 2], [3, 4, 5]], "attention_mask": [[1, 0], [0, 1, 1]]}
+        >>> pad(encoding, 4)
+        {'feature1': [[1, 2, 1, 1], [3, 4, 5, 3]], 'attention_mask': [[1, 0, 0, 0], [0, 1, 1, 0]]}
+    """
+    padded_encoding = {}
+
+    for key, sequences in encoding.items():
+        if not all(isinstance(seq, (list, torch.Tensor)) for seq in sequences):
+            raise TypeError(f"All sequences under key {key} should be of type list or tensor.")
+        if key == "attention_mask":  # attention_mask is handled separately
+            continue
+
+        padded_sequences = []
+        pad_mask = []
+
+        for seq in sequences:
+            pad_len = target_length - len(seq)
+            padded_seq = list(seq) + [seq[0]] * max(0, pad_len)
+            mask = [1] * len(seq) + [0] * max(0, pad_len)
+
+            padded_sequences.append(padded_seq)
+            pad_mask.append(mask)
+
+        padded_encoding[key] = padded_sequences
+
+    if "attention_mask" in encoding:
+        padded_encoding["attention_mask"] = [
+            [a * (b[i] if i < len(b) else 0) for i, a in enumerate(row)]
+            for row, b in zip(pad_mask, encoding["attention_mask"])
+        ]
+    else:
+        padded_encoding["attention_mask"] = pad_mask
+
+    return padded_encoding
+
+
+class JatProcessor(ProcessorMixin):
+    r"""
+    JAT processor which wraps a CLIP image processor and a BERT tokenizer into a single processor.
+
+    [`JatProcessor`] offers all the functionalities of [`CLIPImageProcessor`] and [`BertTokenizerFast`]. See the
+    [`~JatProcessor.__call__`] and [`~JatProcessor.decode`] for more information.
+
+    Args:
+        image_processor ([`AutoImageProcessor`]):
+            The image processor is a required input.
+        tokenizer ([`AutoTokenizer`]):
+            The tokenizer is a required input.
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    DONT_TRUNCATE_OR_PAD = {"pixel_values"}  # Or, a better name for this would be
+
+    def __init__(self, image_processor, tokenizer):
+        super().__init__(image_processor, tokenizer)
+        self.current_processor = self.image_processor
+
+    def _truncate_and_pad(
+        self,
+        encoding: dict,
+        padding: Union[bool, str],
+        truncation: Union[bool, str],
+        truncation_side: str = "right",
+        max_length: Optional[int] = None,
+    ) -> dict:
+        # If max_length is not provided, use the maximum length accepted by the model.
+        if max_length is None:
+            max_length = self.tokenizer.model_max_length
+
+        # Exclude keys that we don't want to truncate or pad.
+        excluded = {key: value for key, value in encoding.items() if key in self.DONT_TRUNCATE_OR_PAD}
+        encoding = {key: value for key, value in encoding.items() if key not in self.DONT_TRUNCATE_OR_PAD}
+
+        # Apply Truncation
+        if truncation in [True, "lossy"]:
+            encoding = truncate(encoding, max_length, truncation_side, preserve=False)
+        elif truncation == "preserve":
+            encoding = truncate(encoding, max_length, truncation_side, preserve=True)
+        elif truncation in [False, "do_not_truncate"]:
+            pass
+        else:
+            raise ValueError("Invalid truncation strategy:" + str(truncation))
+
+        # Apply Padding
+        if padding in [True, "longest"]:
+            target_length = max(len(seq) for sequences in encoding.values() for seq in sequences)
+            encoding = pad(encoding, target_length)
+        elif padding == "max_length":
+            encoding = pad(encoding, max_length)
+        elif padding in [False, "do_not_pad"]:
+            pass
+        else:
+            raise ValueError("Invalid padding strategy:" + str(padding))
+
+        # Add back the excluded keys.
+        encoding.update(excluded)
+
+        # Particular case, we handle the conversion to tensor of image_observations, as the format used
+        # (list of tensors) is not properly handled by the BatchEncoding class:
+        if "image_observations" in encoding:
+            encoding["image_observations"] = to_tensor(encoding["image_observations"])
+
+        return encoding
+
+    def __call__(
+        self,
+        text=None,
+        images=None,
+        continuous_observations=None,
+        discrete_observations=None,
+        text_observations=None,
+        image_observations=None,
+        continuous_actions=None,
+        discrete_actions=None,
+        rewards=None,
+        return_tensors=None,
+        **kwargs,
+    ):
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to BertTokenizerFast's [`~BertTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
+        CLIPImageProcessor's [`~CLIPImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
+        of the above two methods for more information.
+
+        Args:
+            text (`str`, `List[str]`, `List[List[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`,
+                    `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
+                number of channels, H and W are image height and width.
+            continuous_observations (`List[List[List[float]]]`):
+                The continuous observations or batch of continuous observations to be encoded.
+            discrete_observations (`List[List[List[int]]]`):
+                The discrete observations or batch of discrete observations to be encoded.
+            text_observations (`List[List[str]]`):
+                The text observations or batch of text observations to be encoded.
+            image_observations (`List[List[PIL.Image.Image]]`, `List[List[np.ndarray]]`, `List[List[torch.Tensor]]`):
+                The image observations or batch of image observations to be encoded.
+            continuous_actions (`List[List[List[float]]]`):
+                The continuous actions or batch of continuous actions to be encoded.
+            discrete_actions (``List[List[int]]`):
+                The discrete actions or batch of discrete actions to be encoded.
+            rewards (``List[List[float]]`):
+                The rewards or batch of rewards to be encoded.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            [`BatchEncoding`]: A [`BatchEncoding`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+        """
+        # we truncate and pad ourselves so we need to pass padding=False and truncation=False to the tokenizer
+        padding = kwargs.pop("padding", False)
+        truncation = kwargs.pop("truncation", False)
+        truncation_side = kwargs.pop("truncation_side", "right")
+        max_length = kwargs.pop("max_length", None)
+
+        # Ensure that the input is batched
+        if text is not None and not isinstance(text, list):
+            text = [text]
+
+        encoding = {}
+        if text is not None:
+            encoding["input_ids"] = self.tokenizer(text, **kwargs)["input_ids"]
+        if images is not None:
+            encoding["pixel_values"] = self.image_processor(images, **kwargs).pixel_values
+        if continuous_observations is not None:
+            encoding["continuous_observations"] = copy.deepcopy(continuous_observations)
+        if discrete_observations is not None:
+            encoding["discrete_observations"] = copy.deepcopy(discrete_observations)
+        if text_observations is not None:
+            if "discrete_observations" not in encoding:
+                raise ValueError("discrete_observations must be provided if text_observations is provided")
+            for batch_idx, sequence in enumerate(text_observations):
+                encoded_text = self.tokenizer(sequence, max_length=64, padding="max_length")["input_ids"]
+                for timestep, text_tokens in enumerate(encoded_text):
+                    encoding["discrete_observations"][batch_idx][timestep].extend(text_tokens)
+        if image_observations is not None:
+            image_observations = [[(F.to_tensor(im) - 0.5) / 0.5 for im in ep] for ep in image_observations]
+            encoding["image_observations"] = image_observations
+        if continuous_actions is not None:
+            encoding["continuous_actions"] = copy.deepcopy(continuous_actions)
+        if discrete_actions is not None:
+            encoding["discrete_actions"] = copy.deepcopy(discrete_actions)
+
+        if rewards is not None:
+            encoding["rewards"] = [[float(r) for r in ep] for ep in rewards]
+
+        # Handle image+text case, need to reduce the max_len as the image and text will be concatenated
+        if text is not None and images is not None:
+            if max_length is None:
+                max_length = self.tokenizer.model_max_length
+            max_length -= (224 // 16) ** 2  # substract the number of image tokens
+        elif (
+            continuous_observations is not None
+            or discrete_observations is not None
+            or text_observations is not None
+            or image_observations is not None
+        ):
+            if max_length is None:
+                max_length = self.tokenizer.model_max_length
+            max_length //= 2  # observations and actions are interleaved
+
+        encoding = self._truncate_and_pad(encoding, padding, truncation, truncation_side, max_length)
+
+        return BatchEncoding(encoding, tensor_type=return_tensors)
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to BertTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to BertTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def pad(self, *args, **kwargs):
+        inputs = args[0]
+        keys = [key for key in inputs[0].keys() if inputs[0][key] is not None]
+        inputs = {key: [arg[key] for arg in inputs] for key in keys}
+        elmt = next(iter(inputs.values()))
+        if isinstance(elmt[0], torch.Tensor) and not isinstance(elmt, torch.Tensor):
+            encoding = {key: torch.stack(inputs[key]) for key in inputs.keys()}
+        else:
+            encoding = self._truncate_and_pad(
+                inputs, padding=kwargs.get("padding", False), truncation=False, max_length=kwargs.get("max_length")
+            )
+
+        return BatchEncoding(encoding, tensor_type=kwargs.get("return_tensors"))
+
+    @property
+    def model_input_names(self):
+        return [
+            "input_ids",
+            "attention_mask",
+            "pixel_values",
+            "continuous_observations",
+            "discrete_observations",
+            "image_observations",
+            "continuous_actions",
+            "discrete_actions",
+            "rewards",
+        ]
+
+
+JatProcessor.register_for_auto_class("AutoProcessor")

processor_config.json

@@ -0,0 +1,6 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_jat.JatProcessor"
+  },
+  "processor_class": "JatProcessor"
+}

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,28 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoProcessor": "processing_jat.JatProcessor"
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 40,
+  "pad_token": "<|endoftext|>",
+  "processor_class": "JatProcessor",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}