Aluode/PerceptionLabPortable

1# Copyright 2022 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15from typing import Optional, Union16 17import numpy as np18import tensorflow as tf19 20from .feature_extraction_utils import BatchFeature21from .tokenization_utils_base import BatchEncoding22from .utils import logging23 24 25logger = logging.get_logger(__name__)26 27 28def shape_list(tensor: Union[tf.Tensor, np.ndarray]) -> list[int]:29    """30    Deal with dynamic shape in tensorflow cleanly.31 32    Args:33        tensor (`tf.Tensor` or `np.ndarray`): The tensor we want the shape of.34 35    Returns:36        `list[int]`: The shape of the tensor as a list.37    """38    if isinstance(tensor, np.ndarray):39        return list(tensor.shape)40 41    dynamic = tf.shape(tensor)42 43    if tensor.shape == tf.TensorShape(None):44        return dynamic45 46    static = tensor.shape.as_list()47 48    return [dynamic[i] if s is None else s for i, s in enumerate(static)]49 50 51def stable_softmax(logits: tf.Tensor, axis: Optional[int] = None, name: Optional[str] = None) -> tf.Tensor:52    """53    Stable wrapper that returns the same output as `tf.nn.softmax`, but that works reliably with XLA on CPU. It is54    meant as a workaround for the [following issue](https://github.com/tensorflow/tensorflow/issues/55682), and will be55    removed after it gets fixed. The arguments and outputs are the same as `tf.nn.softmax`, and relies on the fact that56    `softmax(x) = softmax(x + c)` (see https://ogunlao.github.io/2020/04/26/you_dont_really_know_softmax.html).57 58    Args:59        logits (`tf.Tensor`):60            Must be one of the following types: half, float32, float64.61        axis (`int`, *optional*):62            The dimension softmax would be performed on. The default is -1 which indicates the last dimension.63        name (`str`, *optional*):64            A name for the operation.65 66    Returns:67        `tf.Tensor`:68            A Tensor. Has the same type and shape as logits.69    """70    # TODO: When the issue linked above gets sorted, add a check on TF version here and use the original function if71    # it has the fix. After we drop the support for unfixed versions, remove this function.72    return tf.nn.softmax(logits=logits + 1e-9, axis=axis, name=name)73 74 75def functional_layernorm(inputs, weight, bias, epsilon=1e-5, axis=-1):76    # This is a very simplified functional layernorm, designed to duplicate77    # the functionality of PyTorch nn.functional.layer_norm when this is needed to port78    # models in Transformers.79 80    if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(axis, int):81        raise NotImplementedError("Only 1D weight and bias tensors are supported for now, with only a single axis.")82 83    # Get mean and variance on the axis to be normalized84    mean, variance = tf.nn.moments(inputs, axes=[axis], keepdims=True)85 86    if axis != -1:87        # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions88        # on every dimension except axis89        shape = [1] * inputs.shape.rank90        shape[axis] = shape_list(inputs)[axis]91        weight = tf.reshape(weight, shape)92        bias = tf.reshape(bias, shape)93 94    # Compute layer normalization using the batch_normalization95    # function.96    outputs = tf.nn.batch_normalization(97        inputs,98        mean,99        variance,100        offset=bias,101        scale=weight,102        variance_epsilon=epsilon,103    )104    return outputs105 106 107def scaled_dot_product_attention(108    query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale: Optional[float] = None109):110    """TF equivalent for torch's nn.functional.scaled_dot_product_attention"""111    if dropout_p != 0.0:112        raise ValueError(113            "Dropout is not supported in this implementation - file an issue "114            "with Transformers and ping @Rocketknight1 if you need it for a port!"115        )116    if is_causal and attn_mask is not None:117        raise ValueError("You cannot specify an attn_mask and is_causal at the same time!")118    if is_causal:119        attn_mask = tf.ones((tf.shape(query)[-2], tf.shape(key)[-2]), dtype=tf.int32)120        attn_mask = tf.experimental.numpy.tril(attn_mask, k=0)121    if attn_mask is not None and (attn_mask.dtype.is_integer or attn_mask.dtype.is_bool):122        # Convert boolean mask to a negative logit bias123        attn_mask = tf.where(attn_mask > 0, tf.cast(0.0, query.dtype), tf.cast(-1000.0, query.dtype))124    logits = tf.einsum("...qd, ...kd -> ...qk", query, key)125    if scale is None:126        scale = tf.cast(tf.shape(key)[-1], logits.dtype) ** -0.5127    logits *= scale  # scale by 1/sqrt(key_dim)128    if attn_mask is not None:129        logits += attn_mask130    probs = tf.nn.softmax(logits)131    return probs @ value132 133 134def flatten(input, start_dim=0, end_dim=-1):135    # Replicates the behavior of torch.flatten in TF136 137    # If end_dim or start_dim is negative, count them from the end138    if end_dim < 0:139        end_dim += input.shape.rank140    if start_dim < 0:141        start_dim += input.shape.rank142 143    if start_dim == end_dim:144        return input145 146    in_shape = tf.shape(input)147    flattened_dim = tf.math.reduce_prod(in_shape[start_dim : end_dim + 1])148    out_shape = tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]], axis=0)149    return tf.reshape(input, out_shape)150 151 152def invert_attention_mask(encoder_attention_mask: tf.Tensor) -> tf.Tensor:153    """154    Invert an attention mask (e.g., switches 0. and 1.).155 156    Args:157        encoder_attention_mask (`torch.Tensor`): An attention mask.158 159    Returns:160        `tf.Tensor`: The inverted attention mask.161    """162    if not isinstance(encoder_attention_mask, tf.Tensor):163        encoder_attention_mask = tf.convert_to_tensor(encoder_attention_mask)  # Catches stray NumPy inputs164    if encoder_attention_mask.shape.rank == 3:165        encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]166    if encoder_attention_mask.shape.rank == 2:167        encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]168    # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition169    # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow170    # /transformer/transformer_layers.py#L270171    # encoder_extended_attention_mask = (encoder_extended_attention_mask ==172    # encoder_extended_attention_mask.transpose(-1, -2))173    encoder_extended_attention_mask = (174        tf.cast(1, encoder_attention_mask.dtype) - encoder_extended_attention_mask175    ) * encoder_extended_attention_mask.dtype.min176 177    return encoder_extended_attention_mask178 179 180def check_embeddings_within_bounds(tensor: tf.Tensor, embed_dim: int, tensor_name: str = "input_ids") -> None:181    """182    `tf.gather`, on which TF embedding layers are based, won't check positive out of bound indices on GPU, returning183    zeros instead. This function adds a check against that dangerous silent behavior.184 185    Args:186        tensor (`tf.Tensor`): The tensor of indices to check.187        embed_dim (`int`): The embedding dimension.188        tensor_name (`str`, *optional*): The name of the tensor to use in the error message.189    """190    tf.debugging.assert_less(191        tensor,192        tf.cast(embed_dim, dtype=tensor.dtype),193        message=(194            f"The maximum value of {tensor_name} ({tf.math.reduce_max(tensor)}) must be smaller than the embedding "195            f"layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."196        ),197    )198 199 200def save_attributes_to_hdf5_group(group, name, data):201    """Saves attributes (data) of the specified name into the HDF5 group.202 203    This method deals with an inherent problem of HDF5 file which is not able to store data larger than204    HDF5_OBJECT_HEADER_LIMIT bytes.205 206    Args:207        group: A pointer to a HDF5 group.208        name: A name of the attributes to save.209        data: Attributes data to store.210 211    Raises:212      RuntimeError: If any single attribute is too large to be saved.213 214    Copied from Keras to Transformers to avoid versioning issues.215    """216    HDF5_OBJECT_HEADER_LIMIT = 64512217    # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`218    # because in that case even chunking the array would not make the saving219    # possible.220    bad_attributes = [x for x in data if len(x) > HDF5_OBJECT_HEADER_LIMIT]221 222    # Expecting this to never be true.223    if bad_attributes:224        raise RuntimeError(225            "The following attributes cannot be saved to HDF5 file because "226            f"they are larger than {HDF5_OBJECT_HEADER_LIMIT} "227            f"bytes: {bad_attributes}"228        )229 230    data_npy = np.asarray(data)231 232    num_chunks = 1233    chunked_data = np.array_split(data_npy, num_chunks)234 235    # This will never loop forever thanks to the test above.236    while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data):237        num_chunks += 1238        chunked_data = np.array_split(data_npy, num_chunks)239 240    if num_chunks > 1:241        for chunk_id, chunk_data in enumerate(chunked_data):242            group.attrs["%s%d" % (name, chunk_id)] = chunk_data243    else:244        group.attrs[name] = data245 246 247def load_attributes_from_hdf5_group(group, name):248    """Loads attributes of the specified name from the HDF5 group.249 250    This method deals with an inherent problem of HDF5 file which is not able to store data larger than251    HDF5_OBJECT_HEADER_LIMIT bytes.252 253    Args:254        group: A pointer to a HDF5 group.255        name: A name of the attributes to load.256 257    Returns:258        data: Attributes data.259 260    Copied from Keras to Transformers to avoid versioning issues.261    """262    if name in group.attrs:263        data = [n.decode("utf8") if hasattr(n, "decode") else n for n in group.attrs[name]]264    else:265        data = []266        chunk_id = 0267        while "%s%d" % (name, chunk_id) in group.attrs:268            data.extend(269                [n.decode("utf8") if hasattr(n, "decode") else n for n in group.attrs["%s%d" % (name, chunk_id)]]270            )271            chunk_id += 1272    return data273 274 275def expand_1d(data):276    """Expands 1-dimensional `Tensor`s into 2-dimensional `Tensor`s.277    Copied from Keras to here to avoid versioning issues."""278 279    def _expand_single_1d_tensor(t):280        if isinstance(t, tf.Tensor) and t.shape.rank == 1:281            return tf.expand_dims(t, axis=-1)282        return t283 284    return tf.nest.map_structure(_expand_single_1d_tensor, data)285 286 287def convert_batch_encoding(*args, **kwargs):288    # Convert HF BatchEncoding/BatchFeature objects in the inputs to dicts that Keras understands289    if args and isinstance(args[0], (BatchEncoding, BatchFeature)):290        args = list(args)291        args[0] = dict(args[0])292    elif "x" in kwargs and isinstance(kwargs["x"], (BatchEncoding, BatchFeature)):293        kwargs["x"] = dict(kwargs["x"])294    return args, kwargs295