first commit

README.md

@@ -1,3 +0,0 @@
----
-license: apache-2.0
----

__init__.py

@@ -0,0 +1,10 @@
+from .modeling_openba import OpenBAForConditionalGeneration
+from .configuration_openba import OpenBAConfig
+from .tokenization_openba import OpenBATokenizer
+
+
+__all__ = [
+    "OpenBAForConditionalGeneration",
+    "OpenBAConfig",
+    "OpenBATokenizer",
+]

added_tokens.json

@@ -0,0 +1,46 @@
+{
+  "\t": 250204,
+  "\n": 250203,
+  "  ": 250205,
+  "   ": 250206,
+  "    ": 250207,
+  "     ": 250208,
+  "      ": 250209,
+  "       ": 250210,
+  "        ": 250211,
+  "         ": 250212,
+  "          ": 250213,
+  "           ": 250214,
+  "            ": 250215,
+  "             ": 250216,
+  "              ": 250217,
+  "               ": 250218,
+  "                ": 250219,
+  "                 ": 250220,
+  "                  ": 250221,
+  "                   ": 250222,
+  "                    ": 250223,
+  "                     ": 250224,
+  "                      ": 250225,
+  "                       ": 250226,
+  "                        ": 250227,
+  "                         ": 250228,
+  "                          ": 250229,
+  "                           ": 250230,
+  "                            ": 250231,
+  "                             ": 250232,
+  "                              ": 250233,
+  "                               ": 250234,
+  "                                ": 250235,
+  "                                 ": 250236,
+  "                                  ": 250237,
+  "                                   ": 250238,
+  "                                    ": 250239,
+  "                                     ": 250240,
+  "                                      ": 250241,
+  "                                       ": 250242,
+  "                                        ": 250243,
+  "<R>": 250200,
+  "<S>": 250201,
+  "<X>": 250202
+}

config.json

@@ -0,0 +1,35 @@
+{
+  "add_ffn_bias": false,
+  "add_lm_head_bias": true,
+  "add_qkv_bias": true,
+  "architectures": [
+    "OpenBAForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_openba.OpenBAConfig",
+    "AutoModel": "modeling_openba.OpenBAForConditionalGeneration",
+    "AutoModelForCausalLM": "modeling_openba.OpenBAForConditionalGeneration",
+    "AutoModelForSeq2SeqLM": "modeling_openba.OpenBAForConditionalGeneration"
+  },
+  "attention_dropout": 0.1,
+  "decoder_max_seq_length": 1024,
+  "decoder_start_token_id": 0,
+  "eos_token_id": 1,
+  "ffn_hidden_size": 16384,
+  "hidden_dropout": 0.1,
+  "hidden_size": 4096,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "kv_channels": 128,
+  "max_seq_length": 1024,
+  "model_type": "openba",
+  "num_decoder_layers": 36,
+  "num_heads": 40,
+  "num_layers": 12,
+  "pad_token_id": 0,
+  "tie_word_embeddings": false,
+  "tokenizer_class": "OpenBATokenizer",
+  "transformers_version": "4.31.0",
+  "use_cache": true,
+  "vocab_size": 250880
+}

configuration_openba.py

@@ -0,0 +1,66 @@
+from transformers.utils import logging
+from transformers.configuration_utils import PretrainedConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class OpenBAConfig(PretrainedConfig):
+    model_type = "openba"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "hidden_size": "hidden_size",
+        "num_attention_heads": "num_heads",
+        "num_hidden_layers": "num_layers"
+    }
+
+    def __init__(
+        self,
+        vocab_size=32128,
+        hidden_size=512,
+        kv_channels=64,
+        ffn_hidden_size=2048,
+        num_layers=12,
+        num_decoder_layers=None,
+        hidden_dropout=0.1,
+        attention_dropout=0.1,
+        num_heads=8,
+        is_encoder_decoder=True,
+        use_cache=True,
+        initializer_factor=1.0,
+        pad_token_id=0,
+        eos_token_id=1,
+        decoder_start_token_id=0,
+        add_qkv_bias=False,
+        add_ffn_bias=False,
+        add_lm_head_bias=False,
+        max_seq_length=1024,
+        decoder_max_seq_length=256,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.kv_channels = kv_channels
+        self.ffn_hidden_size = ffn_hidden_size
+        self.num_layers = num_layers
+        self.num_decoder_layers = (
+            num_decoder_layers if num_decoder_layers is not None else self.num_layers
+        )  # default = symmetry
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+        self.initializer_factor = initializer_factor
+        self.num_heads = num_heads
+        self.add_qkv_bias = add_qkv_bias
+        self.add_ffn_bias = add_ffn_bias
+        self.add_lm_head_bias = add_lm_head_bias
+        self.max_seq_length = max_seq_length
+        self.decoder_max_seq_length = decoder_max_seq_length
+        self.use_cache = use_cache
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            **kwargs,
+        )

modeling_openba.py

@@ -0,0 +1,706 @@
+from typing import Optional, Tuple, Union
+import copy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    Seq2SeqLMOutput,
+    BaseModelOutput,
+)
+from transformers.utils import logging, is_torch_fx_proxy
+
+from .configuration_openba import OpenBAConfig
+
+
+logger = logging.get_logger(__name__)
+
+# Copied from transformers.models.gptj.modeling_gptj.create_sinusoidal_positions
+def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor:
+    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
+    sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(num_pos, dtype=torch.float), inv_freq).float()
+    return torch.cat((torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)), dim=1)
+
+
+def rotate_half(x) -> torch.Tensor:
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
+    sin = torch.cat((sin, sin), dim=-1).to(tensor.device)[:, :, None, :]
+    cos = torch.cat((cos, cos), dim=-1).to(tensor.device)[:, :, None, :]
+    return (tensor * cos) + (rotate_half(tensor) * sin)
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2
+        hidden_size = config.hidden_size
+        ffn_hidden_size = int(2 * config.ffn_hidden_size / 3) 
+        ffn_hidden_size = multiple_of * ((ffn_hidden_size + multiple_of - 1) // multiple_of)
+        self.ffn_hidden_size = ffn_hidden_size
+
+        self.fc_in = nn.Linear(hidden_size, 2 * ffn_hidden_size, bias=config.add_ffn_bias)
+        self.fc_out = nn.Linear(ffn_hidden_size, hidden_size, bias=config.add_ffn_bias)
+        
+        def swiglu(x):
+            x = torch.chunk(x, 2, dim=-1)
+            return F.silu(x[0]) * x[1]
+        self.act_func = swiglu
+
+    def forward(self, hidden_states: Optional[torch.FloatTensor]) -> torch.FloatTensor:
+        hidden_states = self.fc_in(hidden_states)
+        hidden_states = self.act_func(hidden_states)
+        hidden_states = self.fc_out(hidden_states)
+        return hidden_states
+
+
+class OpenBAAttention(nn.Module):
+    def __init__(self, config, attn_type='self'):
+        super().__init__()
+        self.attn_type = attn_type
+        self.is_decoder = config.is_decoder
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.kv_channels = config.kv_channels
+        self.proj_size = self.kv_channels * self.num_heads
+        self.dropout = config.attention_dropout
+        self.scale_attn = torch.sqrt(torch.tensor(self.kv_channels, dtype=torch.float32))
+
+        if self.attn_type == 'self':
+            self.qkv = nn.Linear(self.hidden_size, 3 * self.proj_size, bias=config.add_qkv_bias)
+        else:
+            assert self.attn_type == 'cross'
+            self.q = nn.Linear(self.hidden_size, self.proj_size, bias=config.add_qkv_bias)
+            self.kv = nn.Linear(self.hidden_size, 2 * self.proj_size, bias=config.add_qkv_bias)
+
+        self.rotary_embedding = create_sinusoidal_positions(
+            num_pos=config.max_seq_length,
+            dim=self.kv_channels,
+        )
+
+        self.o = nn.Linear(self.proj_size, self.hidden_size, bias=config.add_qkv_bias)
+
+    def forward(
+        self,
+        hidden_states: Optional[torch.FloatTensor],
+        attention_mask: Optional[torch.FloatTensor] = None,
+        key_value_states: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        layer_head_mask: Optional[Tuple[torch.Tensor]] = None,
+        position_ids:Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+    ):
+        # input is (batch_size, seq_length, hidden_size)
+        batch_size, seq_length = hidden_states.shape[:2]
+        if past_key_value is not None:
+            if len(past_key_value) != 2:
+                raise ValueError(
+                    f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
+                )
+
+        if self.rotary_embedding.device != position_ids.device:
+            self.rotary_embedding = self.rotary_embedding.to(position_ids.device)
+
+        if self.attn_type == 'self':
+            mixed_qkv_states = self.qkv(hidden_states)
+            new_tensor_shape = mixed_qkv_states.size()[:-1] + (self.num_heads, 3 * self.kv_channels)
+            mixed_qkv_states = mixed_qkv_states.view(*new_tensor_shape)
+            query_states, key_states, value_states = torch.chunk(mixed_qkv_states, 3, dim=-1)
+            # rotary position embedding
+            sincos = self.rotary_embedding[position_ids]
+            sin, cos = torch.chunk(sincos, 2, dim=-1)
+            query_states = apply_rotary_pos_emb(query_states, sin, cos)
+            key_states = apply_rotary_pos_emb(key_states, sin, cos)
+            # reshape to (batch_size, num_head, seq_length, kv_channels)
+            query_states = query_states.transpose(1, 2)
+            key_states = key_states.transpose(1, 2)
+            value_states = value_states.transpose(1, 2)
+            if past_key_value is not None:
+                past_key_states, past_value_states = past_key_value
+                key_states = torch.cat([past_key_states, key_states], dim=-2)
+                value_states = torch.cat([past_value_states, value_states], dim=-2)
+        else:
+            assert self.attn_type == 'cross'
+            query_states = self.q(hidden_states)
+            new_tensor_shape = query_states.size()[:-1] + (self.num_heads, self.kv_channels)
+            query_states = query_states.view(*new_tensor_shape)
+            # reshape to (batch_size, num_head, seq_length, kv_channels)
+            query_states = query_states.transpose(1, 2)
+            if past_key_value is None:
+                mixed_kv_states = self.kv(key_value_states)
+                new_tensor_shape = mixed_kv_states.size()[:-1] + (self.num_heads, 2 * self.kv_channels)
+                mixed_kv_states = mixed_kv_states.view(*new_tensor_shape)
+                key_states, value_states = torch.chunk(mixed_kv_states, 2, dim=-1)
+                # reshape to (batch_size, num_head, seq_length, kv_channels)
+                key_states = key_states.transpose(1, 2)
+                value_states = value_states.transpose(1, 2)
+            else:
+                key_states, value_states = past_key_value
+
+        # compute attention score
+        query_states = query_states.to(torch.float32)
+        key_states = key_states.to(torch.float32)
+        attn_scores = torch.matmul(query_states, key_states.transpose(-1, -2)) / self.scale_attn
+        attn_scores = attn_scores.masked_fill_(attention_mask, -10000.0)
+        attn_weights = F.softmax(attn_scores, dim=-1).type_as(attn_scores)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+        attn_weights = attn_weights.to(value_states.dtype)
+
+        # Mask heads if we want to
+        if layer_head_mask is not None:
+            attn_weights = attn_weights * layer_head_mask
+
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.proj_size)
+        attn_output = self.o(attn_output)
+
+        present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None
+        outputs = (attn_output, present_key_value_state)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class OpenBABlock(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.is_decoder = config.is_decoder
+        self.dropout = config.hidden_dropout
+        self.input_layernorm = nn.LayerNorm(config.hidden_size)
+        self.self_attn = OpenBAAttention(config, attn_type='self')
+        self.post_attn_layernorm = nn.LayerNorm(config.hidden_size)
+        if self.is_decoder:
+            self.inter_attn = OpenBAAttention(config, attn_type='cross')
+            self.post_inter_attn_layernorm = nn.LayerNorm(config.hidden_size)
+        self.mlp = SwiGLUMLP(config)
+
+    def forward(
+        self,
+        hidden_states=None,
+        attention_mask=None,
+        position_ids=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        layer_head_mask=None,
+        cross_attn_layer_head_mask=None,
+        past_key_value=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        if past_key_value is not None:
+            if not self.is_decoder:
+                raise ValueError("`past_key_values` is passed to the encoder. Please make sure this is intended.")
+            expected_num_past_key_values = 2 if encoder_hidden_states is None else 4
+
+            if len(past_key_value) != expected_num_past_key_values:
+                raise ValueError(
+                    f"There should be {expected_num_past_key_values} past states. "
+                    f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}"
+                    f"Got {len(past_key_value)} past key / value states"
+                )
+
+            self_attn_past_key_value = past_key_value[:2]
+            cross_attn_past_key_value = past_key_value[2:]
+        else:
+            self_attn_past_key_value, cross_attn_past_key_value = None, None
+
+        # Layer norm at the beginning of the transformer layer.
+        layernorm_output = self.input_layernorm(hidden_states)
+        # Self attention.
+        attn_outputs = self.self_attn(
+            layernorm_output,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            layer_head_mask=layer_head_mask,
+            past_key_value=self_attn_past_key_value,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+        attn_output, present_key_value_state = attn_outputs[:2]
+        attn_weights = attn_outputs[2:]
+        residual = hidden_states
+        # Layer norm post the self attention.
+        attn_output = nn.functional.dropout(attn_output, p=self.dropout, training=self.training)
+        layernorm_input = residual + attn_output
+        layernorm_output = self.post_attn_layernorm(layernorm_input)
+
+        if self.is_decoder:
+            assert encoder_hidden_states is not None
+            attn_outputs = self.inter_attn(
+                layernorm_output,
+                attention_mask=encoder_attention_mask,
+                key_value_states=encoder_hidden_states,
+                position_ids=position_ids,
+                layer_head_mask=cross_attn_layer_head_mask,
+                past_key_value=cross_attn_past_key_value,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+            )
+            attn_output = attn_outputs[0]
+            attn_output = nn.functional.dropout(attn_output, p=self.dropout, training=self.training)
+            # residual connection
+            residual = layernorm_input
+            layernorm_input = residual + attn_output
+            layernorm_output = self.post_inter_attn_layernorm(layernorm_input)
+            # Combine self attn and cross attn key value states
+            if present_key_value_state is not None:
+                present_key_value_state += attn_outputs[1]
+            attn_weights += attn_outputs[2:]
+
+        # MLP.
+        mlp_output = self.mlp(layernorm_output)
+        mlp_output = nn.functional.dropout(mlp_output, p=self.dropout, training=self.training)
+        # Second residual connection.
+        residual = layernorm_input
+        output = residual + mlp_output
+        outputs = (output,)
+
+        if use_cache:
+            outputs += (present_key_value_state,) + attn_weights
+        else:
+            outputs += attn_weights
+        return outputs
+
+
+class OpenBAPreTrainedModel(PreTrainedModel):
+    config_class = OpenBAConfig
+    base_model_prefix = "transformer"
+    _no_split_modules = ["OpenBABlock"]
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (OpenBAAttention, OpenBAStack)):
+            module.gradient_checkpointing = value
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        factor = self.config.initializer_factor
+        if isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+        elif isinstance(module, OpenBAForConditionalGeneration):
+            module.shared_embedding.weight.data.normal_(mean=0.0, std=factor * 1.0)
+            if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:
+                module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0)
+        elif isinstance(module, SwiGLUMLP):
+            module.fc_in.weight.data.normal_(mean=0.0, std=factor * ((self.config.hidden_size) ** -0.5))
+            if hasattr(module.fc_in, "bias") and module.fc_in.bias is not None:
+                module.fc_in.bias.data.zero_()
+            module.fc_out.weight.data.normal_(mean=0.0, std=factor * ((module.ffn_hidden_size) ** -0.5))
+            if hasattr(module.fc_out, "bias") and module.fc_out.bias is not None:
+                module.fc_out.bias.data.zero_()
+        elif isinstance(module, OpenBAAttention):
+            hidden_size = self.config.hidden_size
+            kv_channels = self.config.kv_channels
+            n_heads = self.config.num_heads
+            if module.attn_type == 'self':
+                module.qkv.weight.data[:n_heads * kv_channels].normal_(mean=0.0, std=factor * ((hidden_size * kv_channels) ** -0.5))
+                module.qkv.weight.data[n_heads * kv_channels:].normal_(mean=0.0, std=factor * (hidden_size ** -0.5))
+            else:
+                module.q.weight.data.normal_(mean=0.0, std=factor * ((hidden_size * kv_channels) ** -0.5))
+                module.kv.weight.data.normal_(mean=0.0, std=factor * (hidden_size ** -0.5))
+            module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * kv_channels) ** -0.5))
+
+    def _shift_right(self, input_ids):
+        decoder_start_token_id = self.config.decoder_start_token_id
+        pad_token_id = self.config.pad_token_id
+
+        if decoder_start_token_id is None:
+            raise ValueError(
+                "self.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id."
+                "See T5 docs for more information."
+            )
+
+        # shift inputs to the right
+        if is_torch_fx_proxy(input_ids):
+            # Item assignment is not supported natively for proxies.
+            shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id)
+            shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)
+        else:
+            shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+            shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
+            shifted_input_ids[..., 0] = decoder_start_token_id
+
+        if pad_token_id is None:
+            raise ValueError("self.model.config.pad_token_id has to be defined.")
+        # replace possible -100 values in labels by `pad_token_id`
+        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+        return shifted_input_ids
+
+class OpenBAStack(OpenBAPreTrainedModel):
+    def __init__(self, config, embed_tokens):
+        super().__init__(config)
+        self.embed_tokens = embed_tokens
+        self.is_decoder = config.is_decoder
+        self.block = nn.ModuleList(
+            [OpenBABlock(config) for _ in range(config.num_layers)]
+        )
+        self.final_layernorm = nn.LayerNorm(config.hidden_size)
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        inputs_embeds=None,
+        head_mask=None,
+        cross_attn_head_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # get batch size and seq_length
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        batch_size, seq_length = input_shape
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # required mask seq length can be calculated via length of past
+        if past_key_values is None:
+            past_length = 0
+            past_key_values = [None] * len(self.block)
+        else:
+            past_length = past_key_values[0][0].size(-2)
+        cur_length = past_length + seq_length
+
+        # position ids
+        position_ids = torch.arange(past_length, cur_length, dtype=torch.long, device=device)
+        position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+
+        # Attention mask
+        if attention_mask is None:
+            attention_mask = torch.ones(batch_size, seq_length, device=device)
+        # get extended self-attention mask
+        if self.is_decoder:
+            if len(attention_mask.shape) == 2:
+                seq_ids = torch.arange(seq_length, device=device)
+                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
+                causal_mask = causal_mask.to(attention_mask.dtype)
+                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+            elif len(attention_mask.shape) == 3:
+                extended_attention_mask = attention_mask[:, None, :, :]
+            else:
+                raise ValueError
+        else:
+            extended_attention_mask = attention_mask[:, None, None, :]
+        extended_attention_mask = extended_attention_mask < 0.5
+        # get extended self-attention mask
+        # here we replace encoder_decoder_attention_mask with encoder_attention_mask
+        if self.is_decoder and encoder_hidden_states is not None:
+            if encoder_attention_mask is None:
+                encoder_seq_length = encoder_hidden_states.shape[1]
+                encoder_attention_mask = torch.ones(
+                    batch_size, encoder_seq_length, device=device, dtype=torch.long
+                )
+            extended_encoder_attention_mask = encoder_attention_mask[:, None, None, :]
+            extended_encoder_attention_mask = extended_encoder_attention_mask < 0.5
+        else:
+            extended_encoder_attention_mask = None
+        
+
+        # input embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # Prepare head mask if needed
+        head_mask = self.get_head_mask(head_mask, self.config.num_layers)
+        cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers)
+        present_key_value_states = () if use_cache else None
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        all_cross_attentions = () if (output_attentions and self.is_decoder) else None
+        hidden_states = inputs_embeds
+
+        for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)):
+            layer_head_mask = head_mask[i]
+            cross_attn_layer_head_mask = cross_attn_head_mask[i]
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            layer_outputs = layer_module(
+                hidden_states,
+                attention_mask=extended_attention_mask,
+                position_ids=position_ids,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=extended_encoder_attention_mask,
+                layer_head_mask=layer_head_mask,
+                cross_attn_layer_head_mask=cross_attn_layer_head_mask,
+                past_key_value=past_key_value,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+            )
+            # layer_outputs is a tuple with:
+            # hidden-states, key-value-states, (self-attention weights), (cross-attention weights)
+            if use_cache is False:
+                layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:]
+
+            hidden_states, present_key_value_state = layer_outputs[:2]
+            if use_cache:
+                present_key_value_states += (present_key_value_state,)
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[2],)
+                if self.is_decoder:
+                    all_cross_attentions = all_cross_attentions + (layer_outputs[3],)
+
+        hidden_states = self.final_layernorm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    present_key_value_states,
+                    all_hidden_states,
+                    all_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=present_key_value_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class OpenBAForConditionalGeneration(OpenBAPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [
+        r"encoder.embed_tokens.weight",
+        r"decoder.embed_tokens.weight",
+    ]
+    def __init__(self, config):
+        super().__init__(config)
+        self.shared_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.hidden_size = config.hidden_size
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.use_cache = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = OpenBAStack(encoder_config, self.shared_embedding)
+
+        decoder_config = copy.deepcopy(config)
+        decoder_config.is_decoder = True
+        decoder_config.is_encoder_decoder = False
+        decoder_config.num_layers = config.num_decoder_layers
+        decoder_config.max_seq_length = config.decoder_max_seq_length
+        self.decoder = OpenBAStack(decoder_config, self.shared_embedding)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=config.add_lm_head_bias)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    def get_input_embeddings(self):
+        return self.shared_embedding
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared_embedding = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+        self.decoder.set_input_embeddings(new_embeddings)
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.BoolTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        decoder_head_mask: Optional[torch.FloatTensor] = None,
+        cross_attn_head_mask: Optional[torch.Tensor] = None,
+        encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        decoder_inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
+
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Encode if needed (training, first prediction pass)
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                head_mask=head_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+        elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
+            encoder_outputs = BaseModelOutput(
+                last_hidden_state=encoder_outputs[0],
+                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,\
+            )
+
+        hidden_states = encoder_outputs[0]
+
+        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
+            # get decoder inputs from shifting lm labels to the right
+            decoder_input_ids = self._shift_right(labels)
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            head_mask=decoder_head_mask,
+            cross_attn_head_mask=cross_attn_head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = decoder_outputs[0]
+        # share embedding and softmax embedding
+        if self.config.tie_word_embeddings:
+            # Rescale output before projecting on vocab
+            sequence_output = sequence_output * (self.hidden_size ** -0.5)
+
+        lm_logits = self.lm_head(sequence_output).to(torch.float32)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            # move labels to correct device to enable PP
+            labels = labels.to(lm_logits.device)
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+            loss = loss.to(hidden_states.dtype)
+
+        if not return_dict:
+            output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs
+            return ((loss,) + output) if loss is not None else output
+
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        head_mask=None,
+        decoder_head_mask=None,
+        decoder_attention_mask=None,
+        cross_attn_head_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+        **kwargs,
+    ):
+        # cut decoder_input_ids if past is used
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "decoder_input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "encoder_outputs": encoder_outputs,
+            "attention_mask": attention_mask,
+            "head_mask": head_mask,
+            "decoder_head_mask": decoder_head_mask,
+            "decoder_attention_mask": decoder_attention_mask,
+            "cross_attn_head_mask": cross_attn_head_mask,
+            "use_cache": use_cache,
+        }
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return self._shift_right(labels)
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        # if decoder past is not included in output
+        # speedy decoding is disabled and no need to reorder
+        if past_key_values is None:
+            logger.warning("You might want to consider setting `use_cache=True` to speed up decoding")
+            return past_key_values
+
+        reordered_decoder_past = ()
+        for layer_past_states in past_key_values:
+            # get the correct batch idx from layer past batch dim
+            # batch dim of `past` is at 2nd position
+            reordered_layer_past_states = ()
+            for layer_past_state in layer_past_states:
+                # need to set correct `past` for each of the four key / value states
+                reordered_layer_past_states = reordered_layer_past_states + (
+                    layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)),
+                )
+
+            if reordered_layer_past_states[0].shape != layer_past_states[0].shape:
+                raise ValueError(
+                    f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched"
+                )
+            if len(reordered_layer_past_states) != len(layer_past_states):
+                raise ValueError(
+                    f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched"
+                )
+
+            reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,)
+        return reordered_decoder_past

pytorch_model-00001-of-00003.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e4abd82d21f572794eafe850d93113e021013ece482b5b5fa6da8333519056f
+size 9996790582

pytorch_model-00002-of-00003.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c331b78b1a2b2d7e16b13e8155bb50ffd50ffd045a51cda80921d82b7de7803
+size 9962486872

pytorch_model-00003-of-00003.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d3fd6bcd33b5ac6b2de93214cb389f2d3e64db551d421e78c1bcd36f570c845
+size 9180588837

pytorch_model.bin.index.json

@@ -0,0 +1,782 @@
+{
+  "metadata": {
+    "total_size": 29139601408
+  },
+  "weight_map": {
+    "decoder.block.0.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.inter_attn.kv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.inter_attn.kv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.inter_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.inter_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.inter_attn.q.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.inter_attn.q.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.0.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.post_inter_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.post_inter_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.0.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.inter_attn.kv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.inter_attn.kv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.inter_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.inter_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.inter_attn.q.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.inter_attn.q.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.1.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.post_inter_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.post_inter_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.1.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.10.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.10.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.10.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.11.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.11.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.12.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.12.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.13.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.13.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.14.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.14.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.15.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.15.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.16.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.16.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.17.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.17.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.18.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.18.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.19.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.19.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.2.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.inter_attn.kv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.inter_attn.kv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.inter_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.inter_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.inter_attn.q.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.inter_attn.q.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.2.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.post_inter_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.post_inter_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.2.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.20.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.20.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.20.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.21.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.21.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.22.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.22.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.23.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.23.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.24.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.24.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.25.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.25.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.26.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.26.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.27.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.27.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.28.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.28.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.29.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.29.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.29.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.29.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.29.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.29.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.29.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.3.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.inter_attn.kv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.inter_attn.kv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.inter_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.inter_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.inter_attn.q.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.inter_attn.q.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.3.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.post_inter_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.post_inter_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.3.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.30.input_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.inter_attn.kv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.inter_attn.kv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.inter_attn.q.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.inter_attn.q.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.post_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.post_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.self_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.self_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.self_attn.qkv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.30.self_attn.qkv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.input_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.inter_attn.kv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.inter_attn.kv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.inter_attn.q.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.inter_attn.q.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.post_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.post_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.self_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.self_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.self_attn.qkv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.31.self_attn.qkv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.input_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.inter_attn.kv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.inter_attn.kv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.inter_attn.q.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.inter_attn.q.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.post_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.post_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.self_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.self_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.self_attn.qkv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.32.self_attn.qkv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.input_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.inter_attn.kv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.inter_attn.kv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.inter_attn.q.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.inter_attn.q.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.post_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.post_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.self_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.self_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.self_attn.qkv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.33.self_attn.qkv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.input_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.inter_attn.kv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.inter_attn.kv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.inter_attn.q.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.inter_attn.q.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.post_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.post_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.self_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.self_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.self_attn.qkv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.34.self_attn.qkv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.input_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.input_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.inter_attn.kv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.inter_attn.kv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.inter_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.inter_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.inter_attn.q.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.inter_attn.q.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.mlp.fc_out.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.post_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.post_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.post_inter_attn_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.post_inter_attn_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.self_attn.o.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.self_attn.o.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.self_attn.qkv.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.35.self_attn.qkv.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.4.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.inter_attn.kv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.inter_attn.kv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.inter_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.inter_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.inter_attn.q.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.inter_attn.q.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.4.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.post_inter_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.post_inter_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.4.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.inter_attn.kv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.inter_attn.kv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.inter_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.inter_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.inter_attn.q.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.inter_attn.q.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.5.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.post_inter_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.post_inter_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.5.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.6.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.6.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.6.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.6.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.6.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.6.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "decoder.block.7.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.7.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.7.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.8.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.8.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.input_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.input_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.inter_attn.kv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.inter_attn.kv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.inter_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.inter_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.inter_attn.q.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.inter_attn.q.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.mlp.fc_in.weight": "pytorch_model-00003-of-00003.bin",
+    "decoder.block.9.mlp.fc_out.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.post_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.post_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.post_inter_attn_layernorm.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.post_inter_attn_layernorm.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.self_attn.o.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.self_attn.o.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.self_attn.qkv.bias": "pytorch_model-00002-of-00003.bin",
+    "decoder.block.9.self_attn.qkv.weight": "pytorch_model-00002-of-00003.bin",
+    "decoder.final_layernorm.bias": "pytorch_model-00003-of-00003.bin",
+    "decoder.final_layernorm.weight": "pytorch_model-00003-of-00003.bin",
+    "encoder.block.0.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.0.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.1.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.10.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.11.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.2.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.3.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.4.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.5.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.6.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.7.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.8.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.input_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.input_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.mlp.fc_in.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.mlp.fc_out.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.post_attn_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.post_attn_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.self_attn.o.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.self_attn.o.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.self_attn.qkv.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.block.9.self_attn.qkv.weight": "pytorch_model-00001-of-00003.bin",
+    "encoder.final_layernorm.bias": "pytorch_model-00001-of-00003.bin",
+    "encoder.final_layernorm.weight": "pytorch_model-00001-of-00003.bin",
+    "lm_head.bias": "pytorch_model-00003-of-00003.bin",
+    "lm_head.weight": "pytorch_model-00001-of-00003.bin",
+    "shared_embedding.weight": "pytorch_model-00001-of-00003.bin"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,294 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "\n",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "                                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef78f86560d809067d12bac6c09f19a462cb3af3f54d2b8acbba26e1433125d6
+size 4309802

tokenization_openba.py

@@ -0,0 +1,341 @@
+# coding=utf-8
+# Copyright 2018 T5 Authors and HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Tokenization class for model T5."""
+
+
+import os
+import re
+import warnings
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "spiece.model"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+    }
+}
+
+
+# TODO(PVP) - this should be removed in Transformers v5
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+}
+
+
+class OpenBATokenizer(PreTrainedTokenizer):
+    """
+    Construct a T5 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            [SentencePiece](https://github.com/google/sentencepiece) file (generally has a *.spm* extension) that
+            contains the vocabulary necessary to instantiate a tokenizer.
+        eos_token (`str`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+
+            <Tip>
+
+            When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+            The token used is the `sep_token`.
+
+            </Tip>
+
+        unk_token (`str`, *optional*, defaults to `"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        pad_token (`str`, *optional*, defaults to `"<pad>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        extra_ids (`int`, *optional*, defaults to 100):
+           Add a number of extra ids added to the vocabulary for use as sentinels. These tokens are
+            accessible as "<extra_id_{%d}>" where "{%d}" is a number between 0 and extra_ids-1. These tokens can be
+            retrieved by calling get_sentinel_tokens method and token ids can be by calling get_sentinel_token_ids
+            method
+         additional_special_tokens (`List[str]`, *optional*):
+            Additional special tokens used by the tokenizer.
+        sp_model_kwargs (`dict`, *optional*):
+            Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+            SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+            to set:
+
+            - `enable_sampling`: Enable subword regularization.
+            - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+
+              - `nbest_size = {0,1}`: No sampling is performed.
+              - `nbest_size > 1`: samples from the nbest_size results.
+              - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+                using forward-filtering-and-backward-sampling algorithm.
+
+            - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+              BPE-dropout.
+
+    Attributes:
+        sp_model (`SentencePieceProcessor`):
+            The *SentencePiece* processor that is used for every conversion (string, tokens and IDs).
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        eos_token="</s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+        extra_ids=100,
+        additional_special_tokens=None,
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ) -> None:
+        # Add extra_ids to the special token list
+        if extra_ids > 0 and additional_special_tokens is None:
+            additional_special_tokens = [f"<extra_id_{i}>" for i in range(extra_ids)]
+        elif extra_ids > 0 and additional_special_tokens is not None:
+            # Check that we have the right number of extra_id special tokens
+            extra_tokens = len(set(filter(lambda x: bool("extra_id" in str(x)), additional_special_tokens)))
+            if extra_tokens != extra_ids:
+                raise ValueError(
+                    f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"
+                    " provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids"
+                    " tokens"
+                )
+
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+
+        super().__init__(
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            extra_ids=extra_ids,
+            additional_special_tokens=additional_special_tokens,
+            sp_model_kwargs=self.sp_model_kwargs,
+            **kwargs,
+        )
+
+        self.vocab_file = vocab_file
+        self._extra_ids = extra_ids
+
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+
+    @staticmethod
+    def _eventually_correct_t5_max_length(pretrained_model_name_or_path, max_model_length, init_max_model_length):
+        if pretrained_model_name_or_path in OpenBATokenizer.max_model_input_sizes:
+            deprecated_max_model_length = OpenBATokenizer.max_model_input_sizes[pretrained_model_name_or_path]
+            if init_max_model_length is not None and init_max_model_length != max_model_length:
+                return init_max_model_length
+            elif init_max_model_length is None:
+                warnings.warn(
+                    "This tokenizer was incorrectly instantiated with a model max length of"
+                    f" {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this"
+                    " behavior is kept to avoid breaking backwards compatibility when padding/encoding with"
+                    " `truncation is True`.\n- Be aware that you SHOULD NOT rely on"
+                    f" {pretrained_model_name_or_path} automatically truncating your input to"
+                    f" {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences"
+                    f" longer than {deprecated_max_model_length} you can either instantiate this tokenizer with"
+                    " `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please"
+                    " instantiate this tokenizer with `model_max_length` set to your preferred value.",
+                    FutureWarning,
+                )
+
+        return max_model_length
+
+    @property
+    def vocab_size(self):
+        return self.sp_model.get_piece_size() + self._extra_ids
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        # normal case: some special tokens
+        if token_ids_1 is None:
+            return ([0] * len(token_ids_0)) + [1]
+        return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+
+    def get_sentinel_tokens(self):
+        return list(
+            set(filter(lambda x: bool(re.search(r"<extra_id_\d+>", x)) is not None, self.additional_special_tokens))
+        )
+
+    def get_sentinel_token_ids(self):
+        return [self._convert_token_to_id(token) for token in self.get_sentinel_tokens()]
+
+    def _add_eos_if_not_present(self, token_ids: List[int]) -> List[int]:
+        """Do not add eos again if user already added it."""
+        if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id:
+            warnings.warn(
+                f"This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated"
+                " eos tokens being added."
+            )
+            return token_ids
+        else:
+            return token_ids + [self.eos_token_id]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
+        use of token type ids, therefore a list of zeros is returned.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of zeros.
+        """
+        eos = [self.eos_token_id]
+
+        if token_ids_1 is None:
+            return len(token_ids_0 + eos) * [0]
+        return len(token_ids_0 + eos + token_ids_1 + eos) * [0]
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A sequence has the following format:
+
+        - single sequence: `X </s>`
+        - pair of sequences: `A </s> B </s>`
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        token_ids_0 = self._add_eos_if_not_present(token_ids_0)
+        if token_ids_1 is None:
+            return token_ids_0
+        else:
+            token_ids_1 = self._add_eos_if_not_present(token_ids_1)
+            return token_ids_0 + token_ids_1
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["sp_model"] = None
+        return state
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+
+        # for backward compatibility
+        if not hasattr(self, "sp_model_kwargs"):
+            self.sp_model_kwargs = {}
+
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(self.vocab_file)
+
+    def _tokenize(self, text: str) -> List[str]:
+        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        if token.startswith("<extra_id_"):
+            match = re.match(r"<extra_id_(\d+)>", token)
+            num = int(match.group(1))
+            return self.vocab_size - num - 1
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        if index < self.sp_model.get_piece_size():
+            token = self.sp_model.IdToPiece(index)
+        else:
+            token = f"<extra_id_{self.vocab_size - 1 - index}>"
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        current_sub_tokens = []
+        out_string = ""
+        prev_is_special = False
+        for token in tokens:
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                if not prev_is_special:
+                    out_string += " "
+                out_string += self.sp_model.decode(current_sub_tokens) + token
+                prev_is_special = True
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+                prev_is_special = False
+        out_string += self.sp_model.decode(current_sub_tokens)
+        return out_string.strip()
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)

tokenizer_config.json

@@ -0,0 +1,119 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_openba.OpenBATokenizer",
+      null
+      ]
+  },
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "OpenBATokenizer",
+  "tokenizer_file": null,
+  "unk_token": "<unk>"
+}