code cleanup

bert_padding.py

@@ -1,156 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-# Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
-# Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
-
-"""Helper functions for padding and unpadding batches.  """
-
-from typing import Tuple, cast
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange, repeat
-
-
-class IndexFirstAxis(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, input: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
-        """Get just the values of `input` which are at `indices`.
-
-        Arguments:
-            ctx: the autograd context object
-            input: (b, ...) 2+ dimensional tensor
-            indices: (num_idx) 1D tensor
-        """
-        ctx.save_for_backward(indices)
-        assert input.ndim >= 2
-        ctx.first_axis_dim, other_shape = (
-            input.shape[0],
-            input.shape[1:],
-        )  # type: ignore
-        second_dim = other_shape.numel()  # product of sizes of all but first dimension
-        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
-        return torch.gather(
-            rearrange(input, 'b ... -> b (...)'),  # (b, ...) -> (b, second_dim)
-            0,
-            repeat(
-                indices, 'z -> z d', d=second_dim
-            ),  # (indices,) -> (indices, second_dim)
-        ).reshape(
-            -1, *other_shape
-        )  # (num_idx, ...)
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
-        (indices,) = ctx.saved_tensors
-        assert grad_output.ndim >= 2
-        other_shape = grad_output.shape[1:]
-        grad_output = rearrange(grad_output, 'b ... -> b (...)')
-        grad_input = torch.zeros(
-            [ctx.first_axis_dim, grad_output.shape[1]],
-            device=grad_output.device,
-            dtype=grad_output.dtype,
-        )
-        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
-        # grad_input[indices] = grad_output
-        grad_input.scatter_(
-            0, repeat(indices, 'z -> z d', d=grad_output.shape[1]), grad_output
-        )
-        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
-
-
-index_first_axis = IndexFirstAxis.apply
-
-
-class IndexPutFirstAxis(torch.autograd.Function):
-    @staticmethod
-    def forward(
-        ctx, values: torch.Tensor, indices: torch.Tensor, first_axis_dim
-    ) -> torch.Tensor:
-        ctx.save_for_backward(indices)
-        assert indices.ndim == 1
-        assert values.ndim >= 2
-        output = torch.zeros(
-            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
-        )
-        output[indices] = values
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
-        (indices,) = ctx.saved_tensors
-        grad_values = grad_output[indices]
-        return grad_values, None, None
-
-
-index_put_first_axis = IndexPutFirstAxis.apply
-
-
-def unpad_input(
-    hidden_states: torch.Tensor,
-    attention_mask: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    """Remove padding from input sequences.
-
-    Arguments:
-        hidden_states: (batch, seqlen, ...)
-        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
-
-    Returns:
-        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
-        indices: (total_nnz)
-        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
-        max_seqlen_in_batch: int ()
-    """
-    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    max_seqlen_in_batch = int(seqlens_in_batch.max().item())
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
-    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
-    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
-    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
-    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
-    # so we write custom forward and backward to make it a bit faster.
-    hidden_states = cast(
-        torch.Tensor,
-        index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'), indices),
-    )
-    return hidden_states, indices, cu_seqlens, max_seqlen_in_batch
-
-
-def unpad_input_only(
-    hidden_states: torch.Tensor,
-    attention_mask: torch.Tensor,
-) -> torch.Tensor:
-    """Like unpad_input, but only return the unpadded first tensor.
-
-    Save a small amount of overhead.
-
-    Arguments:
-        hidden_states: (batch, seqlen, ...)
-        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
-
-    Returns:
-        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
-    """
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    return index_first_axis(rearrange(hidden_states, 'b s ... -> (b s) ...'), indices)
-
-
-def pad_input(
-    hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int
-) -> torch.Tensor:
-    """Add padding to sequences.
-
-    Arguments:
-        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
-        indices: (total_nnz)
-        batch: int batch_size
-        seqlen: int max sequence length
-
-    Returns:
-        hidden_states: (batch, seqlen, ...)
-    """
-    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
-    return rearrange(output, '(b s) ... -> b s ...', b=batch)

config.json

@@ -1,15 +1,15 @@
 {
-  "_name_or_path": "jinaai/jina-bert-b-en-v1",
+  "_name_or_path": "jinaai/jina-embedding-v2",
   "model_max_length": 8192,
   "architectures": [
-    "JBertForMaskedLM"
+    "JinaBertForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.0,
   "auto_map": {
-    "AutoConfig": "configuration_jbert.JBertConfig",
-    "AutoModelForMaskedLM": "modeling_jbert.JBertForMaskedLM",
-    "AutoModel": "modeling_jbert.JBertModel",
-    "AutoModelForSequenceClassification": "modeling_jbert.JBertForSequenceClassification"
+    "AutoConfig": "jinaai/jina-embedding-v2--configuration_bert.JinaBertConfig",
+    "AutoModelForMaskedLM": "jinaai/jina-embedding-v2--modeling_bert.JinaBertForMaskedLM",
+    "AutoModel": "jinaai/jina-embedding-v2--modeling_bert.JinaBertModel",
+    "AutoModelForSequenceClassification": "jinaai/jina-embedding-v2--modeling_bert.JinaBertForSequenceClassification"
   },
   "classifier_dropout": null,
   "gradient_checkpointing": false,

configuration_jbert.py

@@ -1,26 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-from transformers import BertConfig as TransformersBertConfig
-
-
-class JBertConfig(TransformersBertConfig):
-    def __init__(
-        self,
-        model_max_length: int = 8192,
-        attention_probs_dropout_prob: float = 0.0,
-        **kwargs,
-    ):
-        """Configuration class for MosaicBert.
-
-        Args:
-            model_max_length (int): Use `model_max_length` to determine how large of an alibi tensor to
-                create when initializing the model. You should be able to ignore this parameter in most cases.
-                Defaults to 8192.
-            attention_probs_dropout_prob (float): By default, turn off attention dropout in Mosaic BERT
-                (otherwise, Flash Attention will be off by default). Defaults to 0.0.
-        """
-        super().__init__(
-            attention_probs_dropout_prob=attention_probs_dropout_prob, **kwargs
-        )
-        self.model_max_length = model_max_length

modeling_jbert.py

@@ -1,908 +0,0 @@
-# Copyright 2022 MosaicML Examples authors
-# SPDX-License-Identifier: Apache-2.0
-
-# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
-# Copyright (c) 2018-2021, NVIDIA CORPORATION.  All rights reserved.
-# Copyright (c) 2022, Tri Dao.
-
-import copy
-import logging
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from einops import rearrange
-from transformers.activations import ACT2FN
-from transformers.modeling_outputs import (
-    MaskedLMOutput,
-    SequenceClassifierOutput,
-    BaseModelOutputWithPastAndCrossAttentions,
-    BaseModelOutputWithPoolingAndCrossAttentions,
-)
-from transformers.models.bert.modeling_bert import BertPreTrainedModel
-
-from .bert_padding import (index_first_axis, index_put_first_axis, pad_input,
-                           unpad_input, unpad_input_only)
-from .configuration_jbert import JBertConfig
-
-logger = logging.getLogger(__name__)
-
-
-class JBertEmbeddings(nn.Module):
-    """Construct the embeddings for words, ignoring position.
-
-    There are no positional embeddings since we use ALiBi and token_type
-    embeddings.
-
-    This module is modeled after the Hugging Face BERT's
-    :class:`~transformers.model.bert.modeling_bert.BertEmbeddings`, but is
-    modified to implement ALiBi. The key change is
-    that position embeddings are removed. Position information instead comes
-    from attention biases that scale linearly with the position distance
-    between query and key tokens.
-
-    This module ignores the `position_ids` input to the `forward` method.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        self.word_embeddings = nn.Embedding(
-            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
-        )
-        # ALiBi doesn't use position embeddings
-        self.token_type_embeddings = nn.Embedding(
-            config.type_vocab_size, config.hidden_size
-        )
-
-        # self.LayerNorm is not snake-cased to stick with TensorFlow model
-        # variable name and be able to load any TensorFlow checkpoint file
-        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.register_buffer(
-            "token_type_ids", torch.zeros((1, config.model_max_length), dtype=torch.long), persistent=False
-        )
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        past_key_values_length: int = 0,
-    ) -> torch.Tensor:
-        if (input_ids is not None) == (inputs_embeds is not None):
-            raise ValueError('Must specify either input_ids or input_embeds!')
-        if input_ids is not None:
-            input_shape = input_ids.size()
-        else:
-            assert inputs_embeds is not None  # just for type checking
-            input_shape = inputs_embeds.size()[:-1]
-
-        seq_length = input_shape[1]
-
-        if position_ids is not None:
-            warnings.warn('position_ids is not used in JBertEmbeddings as it does not have position embeddings.')
-
-        # Setting the token_type_ids to the registered buffer in constructor
-        # where it is all zeros, which usually occurs when it's auto-generated;
-        # registered buffer helps users when tracing the model without passing
-        # token_type_ids, solves issue #5664
-        if token_type_ids is None:
-            if hasattr(self, 'token_type_ids'):
-                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
-                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(
-                    input_shape[0], seq_length
-                )
-                token_type_ids = buffered_token_type_ids_expanded  # type: ignore
-            else:
-                token_type_ids = torch.zeros(
-                    input_shape,  # type: ignore
-                    dtype=torch.long,
-                    device=self.word_embeddings.device,
-                )  # type: ignore  # yapf: disable
-
-        if inputs_embeds is None:
-            inputs_embeds = self.word_embeddings(input_ids)
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
-        embeddings = inputs_embeds + token_type_embeddings
-        embeddings = self.LayerNorm(embeddings)
-        embeddings = self.dropout(embeddings)
-        return embeddings
-
-
-class BertUnpadSelfAttention(nn.Module):
-    """Performs multi-headed self attention on a batch of unpadded sequences.
-
-    If Triton is installed, this module uses Flash Attention to greatly improve throughput.
-    The Flash Attention implementation used in Mosaic BERT supports arbitrary attention biases (which
-    we use to implement ALiBi), but does not support attention dropout. If either Triton is not installed
-    or `config.attention_probs_dropout_prob > 0`, the implementation will default to a
-    math-equivalent pytorch version, which is much slower.
-
-    See `forward` method for additional detail.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
-            config, 'embedding_size'
-        ):
-            raise ValueError(
-                f'The hidden size ({config.hidden_size}) is not a multiple of the number of attention '
-                f'heads ({config.num_attention_heads})'
-            )
-
-        self.num_attention_heads = config.num_attention_heads
-        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
-        # TODO: self.all_head_size == config.hidden_size? Why not just use config.hidden_size?
-        self.all_head_size = self.num_attention_heads * self.attention_head_size
-
-        self.Wqkv = nn.Linear(self.all_head_size, 3 * config.hidden_size)
-
-        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        max_seqlen_in_batch: int,
-        indices: torch.Tensor,
-        attn_mask: torch.Tensor,
-        bias: torch.Tensor,
-    ) -> torch.Tensor:
-        """Perform self-attention.
-
-        If dropout is zero, then we can use the Triton kernel, so we do that. However, if not, we send through a standard PyTorch
-        implementation of self-attention.
-
-        The arguments are unpadded, and our implementations of attention require padded arguments,
-        so we first call `pad_input`. Once we compute attention, we re-unpad our outputs for the other layers.
-        The pad/unpad operations add overhead, but not sending pad tokens through ffs saves compute.
-        It is possible to write an unpadded implementation of attention (in Triton and PyTorch), which we will eventually do.
-
-        Args:
-            hidden_states: (total_nnz, dim)
-            cu_seqlens: (batch + 1,)
-            max_seqlen_in_batch: int
-            indices: (total_nnz,)
-            attn_mask: (batch, max_seqlen_in_batch)
-            bias: (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
-
-        Returns:
-            attention: (total_nnz, dim)
-        """
-        qkv = self.Wqkv(hidden_states)
-        qkv = pad_input(
-            qkv, indices, cu_seqlens.shape[0] - 1, max_seqlen_in_batch
-        )  # batch, max_seqlen_in_batch, thd
-        qkv = rearrange(
-            qkv, 'b s (t h d) -> b s t h d', t=3, h=self.num_attention_heads
-        )
-        # if we have nonzero attention dropout (e.g. during fine-tuning) or no Triton, compute attention in PyTorch
-        q = qkv[:, :, 0, :, :].permute(0, 2, 1, 3)  # b h s d
-        k = qkv[:, :, 1, :, :].permute(0, 2, 3, 1)  # b h d s
-        v = qkv[:, :, 2, :, :].permute(0, 2, 1, 3)  # b h s d
-        attention_scores = torch.matmul(q, k) / math.sqrt(self.attention_head_size)
-        attention_scores = attention_scores + bias
-        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
-        attention_probs = self.dropout(attention_probs)
-        attention_probs = attention_probs.to(dtype=v.dtype)
-        attention = torch.matmul(attention_probs, v).permute(0, 2, 1, 3)  # b s h
-
-        # attn_mask is 1 for attend and 0 for don't
-        attention = unpad_input_only(attention, torch.squeeze(attn_mask) == 1)
-        return rearrange(attention, 'nnz h d -> nnz (h d)')
-
-
-# Copy of transformer's library BertSelfOutput that will not be caught by surgery methods looking for HF BERT modules.
-class BertSelfOutput(nn.Module):
-    """Computes the output of the attention layer.
-
-    This module is modeled after the Hugging Face BERT's
-    :class:`~transformers.model.bert.modeling_bert.BertSelfOutput`.
-    The implementation is identical. Rather than use the original module
-    directly, we re-implement it here so that Mosaic BERT's modules will not
-    be affected by any Composer surgery algorithm that modifies Hugging Face
-    BERT modules.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-
-    def forward(
-        self, hidden_states: torch.Tensor, input_tensor: torch.Tensor
-    ) -> torch.Tensor:
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.dropout(hidden_states)
-        hidden_states = self.LayerNorm(hidden_states + input_tensor)
-        return hidden_states
-
-
-class BertUnpadAttention(nn.Module):
-    """Chains attention, Dropout, and LayerNorm for Mosaic BERT."""
-
-    def __init__(self, config):
-        super().__init__()
-        self.self = BertUnpadSelfAttention(config)
-        self.output = BertSelfOutput(config)
-
-    def forward(
-        self,
-        input_tensor: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        max_s: int,
-        subset_idx: Optional[torch.Tensor] = None,
-        indices: Optional[torch.Tensor] = None,
-        attn_mask: Optional[torch.Tensor] = None,
-        bias: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Forward pass for scaled self-attention without padding.
-
-        Arguments:
-            input_tensor: (total_nnz, dim)
-            cu_seqlens: (batch + 1,)
-            max_s: int
-            subset_idx: () set of indices whose values we care about at the end of the layer
-                        (e.g., the masked tokens, if this is the final layer).
-            indices: None or (total_nnz,)
-            attn_mask: None or (batch, max_seqlen_in_batch)
-            bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
-        """
-        self_output = self.self(
-            input_tensor, cu_seqlens, max_s, indices, attn_mask, bias
-        )
-        if subset_idx is not None:
-            return self.output(
-                index_first_axis(self_output, subset_idx),
-                index_first_axis(input_tensor, subset_idx),
-            )
-        else:
-            return self.output(self_output, input_tensor)
-
-
-class BertGatedLinearUnitMLP(nn.Module):
-    """Applies the FFN at the end of each Mosaic BERT layer.
-
-    Compared to the default BERT architecture, this block replaces :class:`~transformers.model.bert.modeling_bert.BertIntermediate`
-    and :class:`~transformers.model.bert.modeling_bert.SelfOutput` with a single module that has similar functionality, but
-    introduces Gated Linear Units.
-
-    Note: Mosaic BERT adds parameters in order to implement Gated Linear Units. To keep parameter count consistent with that of a
-    standard Hugging Face BERT, scale down `config.intermediate_size` by 2/3. For example, a Mosaic BERT constructed with
-    `config.intermediate_size=2048` will have the same parameter footprint as its Hugging Face BERT counterpart constructed
-    with the `config.intermediate_size=3072`.
-    However, in most cases it will not be necessary to adjust `config.intermediate_size` since, despite the increased
-    parameter size, Mosaic BERT typically offers a net higher throughput than a Hugging Face BERT built from the same `config`.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.gated_layers = nn.Linear(
-            config.hidden_size, config.intermediate_size * 2, bias=False
-        )
-        self.act = nn.GELU(approximate='none')
-        self.wo = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        """Compute new hidden states from current hidden states.
-
-        Args:
-            hidden_states (torch.Tensor): The (unpadded) hidden states from
-                the attention layer [nnz, dim].
-        """
-        residual_connection = hidden_states
-        # compute the activation
-        hidden_states = self.gated_layers(hidden_states)
-        gated = hidden_states[:, : self.config.intermediate_size]
-        non_gated = hidden_states[:, self.config.intermediate_size :]
-        hidden_states = self.act(gated) * non_gated
-        hidden_states = self.dropout(hidden_states)
-        # multiply by the second matrix
-        hidden_states = self.wo(hidden_states)
-        # add the residual connection and post-LN
-        hidden_states = self.layernorm(hidden_states + residual_connection)
-        return hidden_states
-
-
-class BertLayer(nn.Module):
-    """Composes the Mosaic BERT attention and FFN blocks into a single layer."""
-
-    def __init__(self, config: JBertConfig):
-        super().__init__()
-        self.attention = BertUnpadAttention(config)
-        self.mlp = BertGatedLinearUnitMLP(config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        cu_seqlens: torch.Tensor,
-        seqlen: int,
-        subset_idx: Optional[torch.Tensor] = None,
-        indices: Optional[torch.Tensor] = None,
-        attn_mask: Optional[torch.Tensor] = None,
-        bias: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Forward pass for a BERT layer, including both attention and MLP.
-
-        Args:
-            hidden_states: (total_nnz, dim)
-            cu_seqlens: (batch + 1,)
-            seqlen: int
-            subset_idx: () set of indices whose values we care about at the end of the layer
-                        (e.g., the masked tokens, if this is the final layer).
-            indices: None or (total_nnz,)
-            attn_mask: None or (batch, max_seqlen_in_batch)
-            bias: None or (batch, heads, max_seqlen_in_batch, max_seqlen_in_batch)
-        """
-        attention_output = self.attention(
-            hidden_states, cu_seqlens, seqlen, subset_idx, indices, attn_mask, bias
-        )
-        layer_output = self.mlp(attention_output)
-        return layer_output
-
-
-class JBertEncoder(nn.Module):
-    """A stack of BERT layers providing the backbone.
-
-    This module is modeled after the Hugging Face BERT's :class:`~transformers.model.bert.modeling_bert.BertEncoder`,
-    but with substantial modifications to implement unpadding and ALiBi.
-
-    Compared to the analogous Hugging Face BERT module, this module handles unpadding to reduce unnecessary computation
-    at padded tokens, and pre-computes attention biases to implement ALiBi.
-    """
-
-    def __init__(self, config: JBertConfig):
-        super().__init__()
-        self.layer = nn.ModuleList(
-            [BertLayer(config) for _ in range(config.num_hidden_layers)]
-        )
-
-        self.num_attention_heads = config.num_attention_heads
-
-        # The alibi mask will be dynamically expanded if it is too small for
-        # the input the model receives. But it generally helps to initialize it
-        # to a reasonably large size to help pre-allocate CUDA memory.
-        # The default `model_max_length` is 8192.
-        self._current_alibi_size = int(config.model_max_length)
-        self.alibi = torch.zeros(
-            (
-                1,
-                self.num_attention_heads,
-                self._current_alibi_size,
-                self._current_alibi_size,
-            )
-        )
-        self.rebuild_alibi_tensor(size=config.model_max_length)
-
-    def rebuild_alibi_tensor(
-        self, size: int, device: Optional[Union[torch.device, str]] = None
-    ):
-        # Alibi
-        # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
-        # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation
-        # of the logits, which makes the math work out *after* applying causal masking. If no causal masking
-        # will be applied, it is necessary to construct the diagonal mask.
-        n_heads = self.num_attention_heads
-
-        def _get_alibi_head_slopes(n_heads: int) -> List[float]:
-            def get_slopes_power_of_2(n_heads: int) -> List[float]:
-                start = 2 ** (-(2 ** -(math.log2(n_heads) - 3)))
-                ratio = start
-                return [start * ratio**i for i in range(n_heads)]
-
-            # In the paper, they only train models that have 2^a heads for some a. This function
-            # has some good properties that only occur when the input is a power of 2. To
-            # maintain that even when the number of heads is not a power of 2, we use a
-            # workaround.
-            if math.log2(n_heads).is_integer():
-                return get_slopes_power_of_2(n_heads)
-
-            closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
-            slopes_a = get_slopes_power_of_2(closest_power_of_2)
-            slopes_b = _get_alibi_head_slopes(2 * closest_power_of_2)
-            slopes_b = slopes_b[0::2][: n_heads - closest_power_of_2]
-            return slopes_a + slopes_b
-
-        context_position = torch.arange(size, device=device)[:, None]
-        memory_position = torch.arange(size, device=device)[None, :]
-        relative_position = torch.abs(memory_position - context_position)
-        # [n_heads, max_token_length, max_token_length]
-        relative_position = relative_position.unsqueeze(0).expand(n_heads, -1, -1)
-        slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device)
-        alibi = slopes.unsqueeze(1).unsqueeze(1) * -relative_position
-        # [1, n_heads, max_token_length, max_token_length]
-        alibi = alibi.unsqueeze(0)
-        assert alibi.shape == torch.Size([1, n_heads, size, size])
-
-        self._current_alibi_size = size
-        self.alibi = alibi
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        encoder_hidden_states: Optional[torch.FloatTensor] = None,
-        encoder_attention_mask: Optional[torch.FloatTensor] = None,
-        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = False,
-        output_hidden_states: Optional[bool] = False,
-        return_dict: Optional[bool] = True,
-    ) -> List[torch.Tensor]:
-        all_hidden_states = [] if output_hidden_states else None
-
-        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
-        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
-
-        attention_mask_bool = attention_mask.bool()
-        batch, seqlen = hidden_states.shape[:2]
-        # Unpad inputs and mask. It will remove tokens that are padded.
-        # Assume ntokens is total number of tokens (padded and non-padded)
-        # and ntokens_unpad is total number of non-padded tokens.
-        # Then unpadding performs the following compression of the inputs:
-        # hidden_states[ntokens,hidden] -> hidden_states[ntokens_unpad,hidden]
-        hidden_states, indices, cu_seqlens, _ = unpad_input(
-            hidden_states, attention_mask_bool
-        )
-
-        # Add alibi matrix to extended_attention_mask
-        if self._current_alibi_size < seqlen:
-            # Rebuild the alibi tensor when needed
-            warnings.warn(
-                f'Increasing alibi size from {self._current_alibi_size} to {seqlen}'
-            )
-            self.rebuild_alibi_tensor(size=seqlen, device=hidden_states.device)
-        elif self.alibi.device != hidden_states.device:
-            # Device catch-up
-            self.alibi = self.alibi.to(hidden_states.device)
-        alibi_bias = self.alibi[:, :, :seqlen, :seqlen]
-        attn_bias = extended_attention_mask[:, :, :seqlen, :seqlen]
-        alibi_attn_mask = attn_bias + alibi_bias
-
-        for layer_module in self.layer:
-            if output_hidden_states:
-                all_hidden_states.append(rearrange(hidden_states, '(b n) d -> b n d', b=batch))
-            hidden_states = layer_module(
-                hidden_states,
-                cu_seqlens,
-                seqlen,
-                None,
-                indices,
-                attn_mask=attention_mask,
-                bias=alibi_attn_mask,
-            )
-        # Pad inputs and mask. It will insert back zero-padded tokens.
-        # Assume ntokens is total number of tokens (padded and non-padded)
-        # and ntokens_unpad is total number of non-padded tokens.
-        # Then padding performs the following de-compression:
-        #     hidden_states[ntokens_unpad,hidden] -> hidden_states[ntokens,hidden]
-        hidden_states = pad_input(hidden_states, indices, batch, seqlen)
-
-        if output_hidden_states:
-            all_hidden_states.append(hidden_states)
-        
-        if not return_dict:
-            return tuple(
-                v for v in [hidden_states, all_hidden_states] if v is not None
-            )
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=None,
-            hidden_states=all_hidden_states,
-            attentions=None,
-            cross_attentions=None,
-        )
-
-
-class JBertPooler(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.activation = nn.Tanh()
-
-    def forward(
-        self, hidden_states: torch.Tensor, pool: Optional[bool] = True
-    ) -> torch.Tensor:
-        # We "pool" the model by simply taking the hidden state corresponding
-        # to the first token.
-        first_token_tensor = hidden_states[:, 0] if pool else hidden_states
-        pooled_output = self.dense(first_token_tensor)
-        pooled_output = self.activation(pooled_output)
-        return pooled_output
-
-
-class BertPredictionHeadTransform(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        if isinstance(config.hidden_act, str):
-            self.transform_act_fn = ACT2FN[config.hidden_act]
-        else:
-            self.transform_act_fn = config.hidden_act
-        self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=1e-12)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.transform_act_fn(hidden_states)
-        hidden_states = self.LayerNorm(hidden_states)
-        return hidden_states
-
-
-class JBertModel(BertPreTrainedModel):
-    """Overall BERT model.
-
-    Args:
-        config: a JBertConfig class instance with the configuration to build a new model
-
-    Inputs:
-        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
-            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
-            `extract_features.py`, `run_classifier.py` and `run_squad.py`)
-        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
-            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
-            a `sentence B` token (see BERT paper for more details).
-        `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
-            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
-            input sequence length in the current batch. It's the mask that we typically use for attention when
-            a batch has varying length sentences.
-        `output_all_encoded_layers`: boolean which controls the content of the `encoded_layers` output as described below. Default: `True`.
-
-    Outputs: Tuple of (encoded_layers, pooled_output)
-        `encoded_layers`: controlled by `output_all_encoded_layers` argument:
-            - `output_all_encoded_layers=True`: outputs a list of the full sequences of encoded-hidden-states at the end
-                of each attention block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each
-                encoded-hidden-state is a torch.FloatTensor of size [batch_size, sequence_length, hidden_size],
-            - `output_all_encoded_layers=False`: outputs only the full sequence of hidden-states corresponding
-                to the last attention block of shape [batch_size, sequence_length, hidden_size],
-        `pooled_output`: a torch.FloatTensor of size [batch_size, hidden_size] which is the output of a
-            classifier pretrained on top of the hidden state associated to the first character of the
-            input (`CLS`) to train on the Next-Sentence task (see BERT's paper).
-
-    Example usage:
-    ```python
-    # Already been converted into WordPiece token ids
-    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
-    input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
-    token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
-    config = modeling.JBertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
-        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
-    model = JBertModel(config=config)
-    all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
-    ```
-    """
-
-    config_class = JBertConfig
-
-    def __init__(self, config, add_pooling_layer=True):
-        super().__init__(config)
-        self.embeddings = JBertEmbeddings(config)
-        self.encoder = JBertEncoder(config)
-        self.pooler = JBertPooler(config) if add_pooling_layer else None
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embeddings.word_embeddings
-
-    def set_input_embeddings(self, value):
-        self.embeddings.word_embeddings = value
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = False,
-        output_hidden_states: Optional[bool] = False,
-        return_dict: Optional[bool] = True,
-    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
-        if attention_mask is None:
-            attention_mask = torch.ones_like(input_ids)
-        if token_type_ids is None:
-            token_type_ids = torch.zeros_like(input_ids)
-
-        embedding_output = self.embeddings(input_ids, token_type_ids, position_ids)
-
-        encoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.encoder(
-            hidden_states=embedding_output,
-            attention_mask=attention_mask,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        sequence_output = encoder_outputs[0]
-        pooled_output = (
-            self.pooler(sequence_output) if self.pooler is not None else None
-        )
-
-        if not return_dict:
-            return (sequence_output, pooled_output) + encoder_outputs[1:]
-
-        #return encoder_outputs, None
-        return BaseModelOutputWithPoolingAndCrossAttentions(
-            last_hidden_state=sequence_output,
-            pooler_output=pooled_output,
-            past_key_values=encoder_outputs.past_key_values,
-            hidden_states=encoder_outputs.hidden_states,
-            attentions=encoder_outputs.attentions,
-            cross_attentions=encoder_outputs.cross_attentions,
-        )
-
-
-###################
-# Bert Heads
-###################
-class BertLMPredictionHead(nn.Module):
-    def __init__(self, config, bert_model_embedding_weights):
-        super().__init__()
-        self.transform = BertPredictionHeadTransform(config)
-        # The output weights are the same as the input embeddings, but there is
-        # an output-only bias for each token.
-        self.decoder = nn.Linear(
-            bert_model_embedding_weights.size(1), bert_model_embedding_weights.size(0)
-        )
-        self.decoder.weight = bert_model_embedding_weights
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.transform(hidden_states)
-        hidden_states = self.decoder(hidden_states)
-        return hidden_states
-
-
-class BertOnlyMLMHead(nn.Module):
-    def __init__(self, config, bert_model_embedding_weights):
-        super().__init__()
-        self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)
-
-    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
-        prediction_scores = self.predictions(sequence_output)
-        return prediction_scores
-
-
-class BertOnlyNSPHead(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.seq_relationship = nn.Linear(config.hidden_size, 2)
-
-    def forward(self, pooled_output: torch.Tensor) -> torch.Tensor:
-        seq_relationship_score = self.seq_relationship(pooled_output)
-        return seq_relationship_score
-
-
-#####################
-# Various Bert models
-#####################
-class JBertForMaskedLM(BertPreTrainedModel):
-    config_class = JBertConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        if config.is_decoder:
-            warnings.warn(
-                'If you want to use `JBertForMaskedLM` make sure `config.is_decoder=False` for '
-                'bi-directional self-attention.'
-            )
-
-        self.bert = JBertModel(config, add_pooling_layer=False)
-        self.cls = BertOnlyMLMHead(config, self.bert.embeddings.word_embeddings.weight)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_output_embeddings(self):
-        return self.cls.predictions.decoder
-
-    def set_output_embeddings(self, new_embeddings):
-        self.cls.predictions.decoder = new_embeddings
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        encoder_attention_mask: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
-        # labels should be a `torch.LongTensor` of shape
-        # `(batch_size, sequence_length)`. These are used for computing the
-        #  masked language modeling loss.
-        #
-        # Indices should be in `[-100, 0, ..., config.vocab_size]` (see
-        # `input_ids` docstring) Tokens with indices set to `-100` are ignored
-        # (masked), the loss is only computed for the tokens with labels in `[0,
-        # ..., config.vocab_size]`
-        #
-        # Prediction scores are only computed for masked tokens and the (bs,
-        # seqlen) dimensions are flattened
-        if (input_ids is not None) == (inputs_embeds is not None):
-            raise ValueError('Must specify either input_ids or input_embeds!')
-
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-
-        outputs = self.bert(
-            input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            encoder_hidden_states=encoder_hidden_states,
-            encoder_attention_mask=encoder_attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        sequence_output = outputs[0]
-        prediction_scores = self.cls(sequence_output)
-
-        loss = None
-        if labels is not None:
-            # Compute loss
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
-
-        if not return_dict:
-            output = (prediction_scores,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return MaskedLMOutput(
-            loss=loss,
-            logits=prediction_scores,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(
-        self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **model_kwargs
-    ):
-        input_shape = input_ids.shape
-        effective_batch_size = input_shape[0]
-
-        #  add a dummy token
-        if self.config.pad_token_id is None:
-            raise ValueError('The PAD token should be defined for generation')
-
-        attention_mask = torch.cat(
-            [attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))],
-            dim=-1,
-        )
-        dummy_token = torch.full(
-            (effective_batch_size, 1),
-            self.config.pad_token_id,
-            dtype=torch.long,
-            device=input_ids.device,
-        )
-        input_ids = torch.cat([input_ids, dummy_token], dim=1)
-
-        return {'input_ids': input_ids, 'attention_mask': attention_mask}
-
-
-
-class JBertForSequenceClassification(BertPreTrainedModel):
-    """Bert Model transformer with a sequence classification/regression head.
-
-    This head is just a linear layer on top of the pooled output. Used for,
-    e.g., GLUE tasks.
-    """
-
-    config_class = JBertConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
-        self.config = config
-
-        self.bert = JBertModel(config)
-        classifier_dropout = (
-            config.classifier_dropout
-            if config.classifier_dropout is not None
-            else config.hidden_dropout_prob
-        )
-        self.dropout = nn.Dropout(classifier_dropout)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        token_type_ids: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
-        # labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-        # Labels for computing the sequence classification/regression loss.
-        # Indices should be in `[0, ..., config.num_labels - 1]`.
-        # If `config.num_labels == 1` a regression loss is computed
-        # (mean-square loss). If `config.num_labels > 1` a classification loss
-        # is computed (cross-entropy).
-
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-
-        outputs = self.bert(
-            input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        pooled_output = outputs[1]
-
-        pooled_output = self.dropout(pooled_output)
-        logits = self.classifier(pooled_output)
-
-        loss = None
-        if labels is not None:
-            # Compute loss
-            if self.config.problem_type is None:
-                if self.num_labels == 1:
-                    self.config.problem_type = 'regression'
-                elif self.num_labels > 1 and (
-                    labels.dtype == torch.long or labels.dtype == torch.int
-                ):
-                    self.config.problem_type = 'single_label_classification'
-                else:
-                    self.config.problem_type = 'multi_label_classification'
-
-            if self.config.problem_type == 'regression':
-                loss_fct = nn.MSELoss()
-                if self.num_labels == 1:
-                    loss = loss_fct(logits.squeeze(), labels.squeeze())
-                else:
-                    loss = loss_fct(logits, labels)
-            elif self.config.problem_type == 'single_label_classification':
-                loss_fct = nn.CrossEntropyLoss()
-                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-            elif self.config.problem_type == 'multi_label_classification':
-                loss_fct = nn.BCEWithLogitsLoss()
-                loss = loss_fct(logits, labels)
-
-        if not return_dict:
-            output = (logits,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return SequenceClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=None,
-            attentions=None,
-        )