Upload 22 files

modeling/__init__.py

@@ -0,0 +1,37 @@
+#
+# Zhou Bo
+
+#
+
+""" Components for NN
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .tokenizers import *
+from .pooling import *
+from .mlm import MLMPredictionHead
+from .nnmodule import NNModule
+from .deberta import *
+from .disentangled_attention import *
+from .ops import *
+from .bert import *
+from .config import *
+from .cache_utils import *
+from .focal_loss import *
+# from .tokenization import BertTokenizer, BasicTokenizer, WordpieceTokenizer
+from .modeling import (BertConfig, BertModel, BertForPreTraining, BertForMaskedLM, 
+                        BertForNextSentencePrediction, PreTrainedBertModel,
+                        BertForSequenceClassification, BertForMultipleChoice, BertForTokenClassification, 
+                        BertForQuestionAnswering, BertForPreTrainingLossMask, BertPreTrainingPairRel, 
+                        BertPreTrainingPairTransform, BertPreTrainingHeads, MLMHead)
+# from .optimization import BertAdam, BertAdamFineTune
+try:
+    from .optimization_fp16 import FP16_Optimizer_State
+except:
+    pass
+from .file_utils import PYTORCH_PRETRAINED_BERT_CACHE
+from .flash import FlashQuadModel
+from .gat import GatModel

modeling/bert.py

@@ -0,0 +1,307 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This piece of code is modified based on https://github.com/huggingface/transformers
+
+import copy
+import torch
+from torch import nn
+from collections import Sequence
+from packaging import version
+import numpy as np
+import math
+import os
+import pdb
+
+import json
+from .ops import *
+from .disentangled_attention import *
+from .da_utils import *
+
+__all__ = ['BertEncoder', 'BertEmbeddings', 'ACT2FN', 'LayerNorm', 'BertLMPredictionHead']
+
+class BertSelfOutput(nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+    self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+    self.dropout = StableDropout(config.hidden_dropout_prob)
+    self.config = config
+
+  def forward(self, hidden_states, input_states, mask=None):
+    hidden_states = self.dense(hidden_states)
+    hidden_states = self.dropout(hidden_states)
+    hidden_states += input_states
+    hidden_states = MaskedLayerNorm(self.LayerNorm, hidden_states)
+    return hidden_states
+
+class BertAttention(nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.self = DisentangledSelfAttention(config)
+    self.output = BertSelfOutput(config)
+    self.config = config
+
+  def forward(self, hidden_states, attention_mask, return_att=False, query_states=None, relative_pos=None, rel_embeddings=None):
+    output = self.self(hidden_states, attention_mask, return_att, query_states=query_states, relative_pos=relative_pos, rel_embeddings=rel_embeddings)
+    self_output, att_matrix, att_logits_=output['hidden_states'], output['attention_probs'], output['attention_logits']
+    if query_states is None:
+      query_states = hidden_states
+    attention_output = self.output(self_output, query_states, attention_mask)
+
+    if return_att:
+      return (attention_output, att_matrix)
+    else:
+      return attention_output
+
+class BertIntermediate(nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+    self.intermediate_act_fn = ACT2FN[config.hidden_act] \
+      if isinstance(config.hidden_act, str) else config.hidden_act
+
+  def forward(self, hidden_states):
+    hidden_states = self.dense(hidden_states)
+    hidden_states = self.intermediate_act_fn(hidden_states)
+    return hidden_states
+
+class BertOutput(nn.Module):
+  def __init__(self, config):
+    super(BertOutput, self).__init__()
+    self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+    self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+    self.dropout = StableDropout(config.hidden_dropout_prob)
+    self.config = config
+
+  def forward(self, hidden_states, input_states, mask=None):
+    hidden_states = self.dense(hidden_states)
+    hidden_states = self.dropout(hidden_states)
+    hidden_states += input_states
+    hidden_states = MaskedLayerNorm(self.LayerNorm, hidden_states)
+    return hidden_states
+
+class BertLayer(nn.Module):
+  def __init__(self, config):
+    super(BertLayer, self).__init__()
+    self.attention = BertAttention(config)
+    self.intermediate = BertIntermediate(config)
+    self.output = BertOutput(config)
+
+  def forward(self, hidden_states, attention_mask, return_att=False, query_states=None, relative_pos=None, rel_embeddings=None):
+    attention_output = self.attention(hidden_states, attention_mask, return_att=return_att, \
+      query_states=query_states, relative_pos=relative_pos, rel_embeddings=rel_embeddings)
+    if return_att:
+      attention_output, att_matrix = attention_output
+    intermediate_output = self.intermediate(attention_output)
+    layer_output = self.output(intermediate_output, attention_output, attention_mask)
+    if return_att:
+      return (layer_output, att_matrix)
+    else:
+      return layer_output
+
+class ConvLayer(nn.Module):
+    def __init__(self, config):
+      super().__init__()
+      kernel_size = getattr(config, 'conv_kernel_size', 3)
+      groups = getattr(config, 'conv_groups', 1)
+      self.conv_act = getattr(config, 'conv_act', 'tanh')
+      self.conv = torch.nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size, padding = (kernel_size-1)//2, groups = groups)
+      self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+      self.dropout = StableDropout(config.hidden_dropout_prob)
+      self.config = config
+
+    def forward(self, hidden_states, residual_states, input_mask):
+        out = self.conv(hidden_states.permute(0,2,1).contiguous()).permute(0,2,1).contiguous()
+        if version.Version(torch.__version__) >= version.Version('1.2.0a'):
+            rmask = (1-input_mask).bool()
+        else:
+            rmask = (1-input_mask).byte()
+        out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0)
+        out = ACT2FN[self.conv_act](self.dropout(out))
+        output_states = MaskedLayerNorm(self.LayerNorm, residual_states + out, input_mask)
+
+        return output_states
+
+class BertEncoder(nn.Module):
+  """ Modified BertEncoder with relative position bias support
+  """
+  def __init__(self, config):
+    super().__init__()
+    #layer = BertLayer(config)
+    self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
+    self.relative_attention = getattr(config, 'relative_attention', False)
+    if self.relative_attention:
+      self.max_relative_positions = getattr(config, 'max_relative_positions', -1)
+      if self.max_relative_positions <1:
+        self.max_relative_positions = config.max_position_embeddings
+      self.position_buckets = getattr(config, 'position_buckets', -1)
+      pos_ebd_size = self.max_relative_positions*2
+      if self.position_buckets>0:
+        pos_ebd_size = self.position_buckets*2
+      self.rel_embeddings = nn.Embedding(pos_ebd_size, config.hidden_size)
+
+    self.norm_rel_ebd = [x.strip() for x in getattr(config, 'norm_rel_ebd', 'none').lower().split('|')]
+    if 'layer_norm' in self.norm_rel_ebd:
+      self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine = True)
+    kernel_size = getattr(config, 'conv_kernel_size', 0)
+    self.with_conv = False
+    if kernel_size > 0:
+      self.with_conv = True
+      self.conv = ConvLayer(config)
+
+  def get_rel_embedding(self):
+    rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
+    if rel_embeddings is not None and ('layer_norm' in self.norm_rel_ebd):
+      rel_embeddings = self.LayerNorm(rel_embeddings)
+    return rel_embeddings
+
+  def get_attention_mask(self, attention_mask):
+    if attention_mask.dim()<=2:
+      extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+      attention_mask = extended_attention_mask*extended_attention_mask.squeeze(-2).unsqueeze(-1)
+      attention_mask = attention_mask.byte()
+    elif attention_mask.dim()==3:
+      attention_mask = attention_mask.unsqueeze(1)
+
+    return attention_mask
+
+  def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
+    if self.relative_attention and relative_pos is None:
+      q = query_states.size(-2) if query_states is not None else hidden_states.size(-2)
+      relative_pos = build_relative_position(q, hidden_states.size(-2), bucket_size = self.position_buckets, max_position=self.max_relative_positions)
+    return relative_pos
+
+  def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True, return_att=False, query_states = None, relative_pos=None):
+    if attention_mask.dim()<=2:
+      input_mask = attention_mask
+    else:
+      input_mask = (attention_mask.sum(-2)>0).byte()
+    attention_mask = self.get_attention_mask(attention_mask)
+    relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
+
+    all_encoder_layers = []
+    att_matrices = []
+    if isinstance(hidden_states, Sequence):
+      next_kv = hidden_states[0]
+    else:
+      next_kv = hidden_states
+    rel_embeddings = self.get_rel_embedding()
+    for i, layer_module in enumerate(self.layer):
+      output_states = layer_module(next_kv, attention_mask, return_att, query_states = query_states, relative_pos=relative_pos, rel_embeddings=rel_embeddings)
+      if return_att:
+        output_states, att_m = output_states
+
+      if i == 0 and self.with_conv:
+        prenorm = output_states #output['prenorm_states']
+        output_states = self.conv(hidden_states, prenorm, input_mask)
+
+      if query_states is not None:
+        query_states = output_states
+        if isinstance(hidden_states, Sequence):
+          next_kv = hidden_states[i+1] if i+1 < len(self.layer) else None
+      else:
+        next_kv = output_states
+
+      if output_all_encoded_layers:
+        all_encoder_layers.append(output_states)
+        if return_att:
+          att_matrices.append(att_m)
+    if not output_all_encoded_layers:
+      all_encoder_layers.append(output_states)
+      if return_att:
+        att_matrices.append(att_m)
+    return {
+        'hidden_states': all_encoder_layers,
+        'attention_matrices': att_matrices
+        }
+
+class BertEmbeddings(nn.Module):
+  """Construct the embeddings from word, position and token_type embeddings.
+  """
+  def __init__(self, config):
+    super(BertEmbeddings, self).__init__()
+    padding_idx = getattr(config, 'padding_idx', 0)
+    self.embedding_size = getattr(config, 'embedding_size', config.hidden_size)
+    self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx = padding_idx)
+    self.position_biased_input = getattr(config, 'position_biased_input', True)
+    self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size)
+
+    if config.type_vocab_size>0:
+      self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size)
+    
+    if self.embedding_size != config.hidden_size:
+      self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False)
+    self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+    self.dropout = StableDropout(config.hidden_dropout_prob)
+    self.output_to_half = False
+    self.config = config
+
+  def forward(self, input_ids, token_type_ids=None, position_ids=None, mask = None):
+    seq_length = input_ids.size(1)
+    if position_ids is None:
+      position_ids = torch.arange(0, seq_length, dtype=torch.long, device=input_ids.device)
+      position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+    if token_type_ids is None:
+      token_type_ids = torch.zeros_like(input_ids)
+
+    words_embeddings = self.word_embeddings(input_ids)
+    position_embeddings = self.position_embeddings(position_ids.long())
+
+    embeddings = words_embeddings
+    if self.config.type_vocab_size>0:
+      token_type_embeddings = self.token_type_embeddings(token_type_ids)
+      embeddings += token_type_embeddings
+
+    if self.position_biased_input:
+      embeddings += position_embeddings
+
+    if self.embedding_size != self.config.hidden_size:
+      embeddings = self.embed_proj(embeddings)
+    embeddings = MaskedLayerNorm(self.LayerNorm, embeddings, mask)
+    embeddings = self.dropout(embeddings)
+    return {
+        'embeddings': embeddings,
+        'position_embeddings': position_embeddings}
+
+class BertLMPredictionHead(nn.Module):
+    def __init__(self, config, vocab_size):
+        super().__init__()
+        self.embedding_size = getattr(config, 'embedding_size', config.hidden_size)
+        self.dense = nn.Linear(config.hidden_size, self.embedding_size)
+        self.transform_act_fn = ACT2FN[config.hidden_act] \
+            if isinstance(config.hidden_act, str) else config.hidden_act
+
+        self.LayerNorm = LayerNorm(self.embedding_size, config.layer_norm_eps, elementwise_affine=True)
+
+        self.bias = nn.Parameter(torch.zeros(vocab_size))
+
+    def forward(self, hidden_states, embeding_weight):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        # b x s x d
+        hidden_states = MaskedLayerNorm(self.LayerNorm, hidden_states)
+
+        # b x s x v
+        logits = torch.matmul(hidden_states, embeding_weight.t().to(hidden_states)) + self.bias
+        return logits
+
+
+class AR_MASK(object):
+  def get_attention_mask(self, input_ids=None, token_type_ids=None ):
+    seq_len = input_ids.size(1)
+    # idxs = torch.arange(0, seq_len)
+    # mask = idxs[None, :] <= idxs[:, None]
+    mask = torch.tril(torch.ones((seq_len, seq_len), dtype=torch.uint8)).to(input_ids.device)
+    mask = mask.unsqueeze(0).expand(input_ids.size(0), seq_len, seq_len)
+    return mask
+    # torch.diagonal(torch.ones([input_ids.size(1), input_ids.size(1)])).byte().to(input_ids.device)
+
+class Prefix_MASK(object):
+  def get_attention_mask(self, input_ids=None, token_type_ids=None):
+    idxs = torch.cumsum(token_type_ids, axis=1)
+    mask = idxs[:, None, :] <= idxs[:, :, None]
+    return mask.byte().to(input_ids.device)

modeling/cache_utils.py

@@ -0,0 +1,132 @@
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 05/15/2020
+#
+
+import pdb
+import torch
+import os
+import requests
+from .config import ModelConfig
+import pathlib
+from ..utils import xtqdm as tqdm
+from zipfile import ZipFile
+import loguru
+# from ..utils import get_logger
+logger = loguru.logger
+
+__all__ = ['pretrained_models', 'load_model_state', 'load_vocab']
+
+class PretrainedModel:
+  def __init__(self, name, vocab, vocab_type, model='pytorch_model.bin', config='config.json', **kwargs):
+    self.__dict__.update(kwargs)
+    host = f'https://huggingface.co/microsoft/{name}/resolve/main/'
+    self.name = name
+    self.model_url = host + model
+    self.config_url = host + config
+    self.vocab_url = host + vocab
+    self.vocab_type = vocab_type
+  
+pretrained_models= {
+    'base': PretrainedModel('deberta-base', 'bpe_encoder.bin', 'gpt2'),
+    'large': PretrainedModel('deberta-large', 'bpe_encoder.bin', 'gpt2'),
+    'xlarge': PretrainedModel('deberta-xlarge', 'bpe_encoder.bin', 'gpt2'),
+    'base-mnli': PretrainedModel('deberta-base-mnli', 'bpe_encoder.bin', 'gpt2'),
+    'large-mnli': PretrainedModel('deberta-large-mnli', 'bpe_encoder.bin', 'gpt2'),
+    'xlarge-mnli': PretrainedModel('deberta-xlarge-mnli', 'bpe_encoder.bin', 'gpt2'),
+    'xlarge-v2': PretrainedModel('deberta-v2-xlarge', 'spm.model', 'spm'),
+    'xxlarge-v2': PretrainedModel('deberta-v2-xxlarge', 'spm.model', 'spm'),
+    'xlarge-v2-mnli': PretrainedModel('deberta-v2-xlarge-mnli', 'spm.model', 'spm'),
+    'xxlarge-v2-mnli': PretrainedModel('deberta-v2-xxlarge-mnli', 'spm.model', 'spm'),
+    'deberta-v3-small': PretrainedModel('deberta-v3-small', 'spm.model', 'spm'),
+    'deberta-v3-base': PretrainedModel('deberta-v3-base', 'spm.model', 'spm'),
+    'deberta-v3-large': PretrainedModel('deberta-v3-large', 'spm.model', 'spm'),
+    'mdeberta-v3-base': PretrainedModel('mdeberta-v3-base', 'spm.model', 'spm'),
+    'deberta-v3-xsmall': PretrainedModel('deberta-v3-xsmall', 'spm.model', 'spm'),
+  }
+
+def download_asset(url, name, tag=None, no_cache=False, cache_dir=None):
+  _tag = tag
+  if _tag is None:
+    _tag = 'latest'
+  if not cache_dir:
+    cache_dir = os.path.join(pathlib.Path.home(), f'.~DeBERTa/assets/{_tag}/')
+  os.makedirs(cache_dir, exist_ok=True)
+  output=os.path.join(cache_dir, name)
+  if os.path.exists(output) and (not no_cache):
+    return output
+
+  #repo=f'https://huggingface.co/microsoft/deberta-{name}/blob/main/bpe_encoder.bin'
+  headers = {}
+  headers['Accept'] = 'application/octet-stream'
+  resp = requests.get(url, stream=True, headers=headers)
+  if resp.status_code != 200:
+    raise Exception(f'Request for {url} return {resp.status_code}, {resp.text}')
+  
+  try:
+    with open(output, 'wb') as fs:
+      progress = tqdm(total=int(resp.headers['Content-Length']) if 'Content-Length' in resp.headers else -1, ncols=80, desc=f'Downloading {name}')
+      for c in resp.iter_content(chunk_size=1024*1024):
+        fs.write(c)
+        progress.update(len(c))
+      progress.close()
+  except:
+    os.remove(output)
+    raise
+
+  return output
+
+def load_model_state(path_or_pretrained_id, tag=None, no_cache=False, cache_dir=None):
+  model_path = path_or_pretrained_id
+  if model_path and (not os.path.exists(model_path)) and (path_or_pretrained_id.lower() in pretrained_models):
+    _tag = tag
+    pretrained = pretrained_models[path_or_pretrained_id.lower()]
+    if _tag is None:
+      _tag = 'latest'
+    if not cache_dir:
+      cache_dir = os.path.join(pathlib.Path.home(), f'.~DeBERTa/assets/{_tag}/{pretrained.name}')
+    os.makedirs(cache_dir, exist_ok=True)
+    model_path = os.path.join(cache_dir, 'pytorch_model.bin')
+    if (not os.path.exists(model_path)) or no_cache:
+      asset = download_asset(pretrained.model_url, 'pytorch_model.bin', tag=tag, no_cache=no_cache, cache_dir=cache_dir)
+      asset = download_asset(pretrained.config_url, 'model_config.json', tag=tag, no_cache=no_cache, cache_dir=cache_dir)
+  elif not model_path:
+    return None,None
+
+  config_path = os.path.join(os.path.dirname(model_path), 'model_config.json')
+  model_state = torch.load(model_path, map_location='cpu')
+  logger.info("Loaded pretrained model file {}".format(model_path))
+  if 'config' in model_state:
+    model_config = ModelConfig.from_dict(model_state['config'])
+  elif os.path.exists(config_path):
+    model_config = ModelConfig.from_json_file(config_path)
+  else:
+    model_config = None
+  return model_state, model_config
+
+def load_vocab(vocab_path=None, vocab_type=None, pretrained_id=None, tag=None, no_cache=False, cache_dir=None):
+  if pretrained_id and (pretrained_id.lower() in pretrained_models):
+    _tag = tag
+    if _tag is None:
+      _tag = 'latest'
+
+    pretrained = pretrained_models[pretrained_id.lower()]
+    if not cache_dir:
+      cache_dir = os.path.join(pathlib.Path.home(), f'.~DeBERTa/assets/{_tag}/{pretrained.name}')
+    os.makedirs(cache_dir, exist_ok=True)
+    vocab_type = pretrained.vocab_type
+    url = pretrained.vocab_url
+    outname = os.path.basename(url)
+    vocab_path =os.path.join(cache_dir, outname)
+    if (not os.path.exists(vocab_path)) or no_cache:
+      asset = download_asset(url, outname, tag=tag, no_cache=no_cache, cache_dir=cache_dir)
+  if vocab_type is None:
+    vocab_type = 'spm'
+  return vocab_path, vocab_type
+
+def test_download():
+  vocab = load_vocab()

modeling/config.py

@@ -0,0 +1,90 @@
+import json
+import copy
+
+__all__=['AbsModelConfig', 'ModelConfig']
+
+class AbsModelConfig(object):
+    def __init__(self):
+        pass
+
+    @classmethod
+    def from_dict(cls, json_object):
+        """Constructs a `ModelConfig` from a Python dictionary of parameters."""
+        config = cls()
+        for key, value in json_object.items():
+            if isinstance(value, dict):
+                value = AbsModelConfig.from_dict(value)
+            config.__dict__[key] = value
+        return config
+
+    @classmethod
+    def from_json_file(cls, json_file):
+        """Constructs a `ModelConfig` from a json file of parameters."""
+        with open(json_file, "r", encoding='utf-8') as reader:
+            text = reader.read()
+        return cls.from_dict(json.loads(text))
+
+    def __repr__(self):
+        return str(self.to_json_string())
+
+    def to_dict(self):
+        """Serializes this instance to a Python dictionary."""
+        output = copy.deepcopy(self.__dict__)
+        return output
+
+    def to_json_string(self):
+        """Serializes this instance to a JSON string."""
+        def _json_default(obj):
+            if isinstance(obj, AbsModelConfig):
+                return obj.__dict__
+        return json.dumps(self.__dict__, indent=2, sort_keys=True, default=_json_default) + "\n"
+
+class ModelConfig(AbsModelConfig):
+    """Configuration class to store the configuration of a :class:`~DeBERTa.deberta.DeBERTa` model.
+
+        Attributes:
+            hidden_size (int): Size of the encoder layers and the pooler layer, default: `768`.
+            num_hidden_layers (int): Number of hidden layers in the Transformer encoder, default: `12`.
+            num_attention_heads (int): Number of attention heads for each attention layer in
+                the Transformer encoder, default: `12`.
+            intermediate_size (int): The size of the "intermediate" (i.e., feed-forward)
+                layer in the Transformer encoder, default: `3072`.
+            hidden_act (str): The non-linear activation function (function or string) in the
+                encoder and pooler. If string, "gelu", "relu" and "swish" are supported, default: `gelu`.
+            hidden_dropout_prob (float): The dropout probabilitiy for all fully connected
+                layers in the embeddings, encoder, and pooler, default: `0.1`.
+            attention_probs_dropout_prob (float): The dropout ratio for the attention
+                probabilities, default: `0.1`.
+            max_position_embeddings (int): The maximum sequence length that this model might
+                ever be used with. Typically set this to something large just in case
+                (e.g., 512 or 1024 or 2048), default: `512`.
+            type_vocab_size (int): The vocabulary size of the `token_type_ids` passed into
+                `DeBERTa` model, default: `-1`.
+            initializer_range (int): The sttdev of the _normal_initializer for
+                initializing all weight matrices, default: `0.02`.
+            relative_attention (:obj:`bool`): Whether use relative position encoding, default: `False`.
+            max_relative_positions (int): The range of relative positions [`-max_position_embeddings`, `max_position_embeddings`], default: -1, use the same value as `max_position_embeddings`. 
+            padding_idx (int): The value used to pad input_ids, default: `0`.
+            position_biased_input (:obj:`bool`): Whether add absolute position embedding to content embedding, default: `True`.
+            pos_att_type (:obj:`str`): The type of relative position attention, it can be a combination of [`p2c`, `c2p`, `p2p`], e.g. "p2c", "p2c|c2p", "p2c|c2p|p2p"., default: "None".
+
+
+    """
+    def __init__(self):
+        """Constructs ModelConfig.
+
+        """
+        
+        self.hidden_size = 768
+        self.num_hidden_layers = 12
+        self.num_attention_heads = 12
+        self.hidden_act = "gelu"
+        self.intermediate_size = 3072
+        self.hidden_dropout_prob = 0.1
+        self.attention_probs_dropout_prob = 0.1
+        self.max_position_embeddings = 512
+        self.type_vocab_size = 0
+        self.initializer_range = 0.02
+        self.layer_norm_eps = 1e-7
+        self.padding_idx = 0
+        self.vocab_size = -1

modeling/da_utils.py

@@ -0,0 +1,36 @@
+import torch
+import pdb
+from functools import lru_cache
+import numpy as np
+
+__all__=['build_relative_position', 'make_log_bucket_position']
+
+def make_log_bucket_position(relative_pos, bucket_size, max_position):
+  sign = np.sign(relative_pos)
+  mid = bucket_size//2
+  abs_pos = np.where((relative_pos<mid) & (relative_pos > -mid), mid-1, np.abs(relative_pos))
+  log_pos = np.ceil(np.log(abs_pos/mid)/np.log((max_position-1)/mid) * (mid-1)) + mid
+  bucket_pos = np.where(abs_pos<=mid, relative_pos, log_pos*sign).astype(np.int)
+  return bucket_pos
+
+@lru_cache(maxsize=128)
+def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1):
+  q_ids = np.arange(0, query_size)
+  k_ids = np.arange(0, key_size)
+  rel_pos_ids = q_ids[:, None] - np.tile(k_ids, (q_ids.shape[0],1))
+  if bucket_size>0 and max_position > 0:
+    rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
+  rel_pos_ids = torch.tensor(rel_pos_ids, dtype=torch.long)
+  rel_pos_ids = rel_pos_ids[:query_size, :]
+  rel_pos_ids = rel_pos_ids.unsqueeze(0)
+  return rel_pos_ids
+
+def test_log_bucket():
+  x=np.arange(-511,511)
+  y=make_log_bucket_position(x, 128, 512)
+  # pdb.set_trace()
+
+
+if __name__ == '__main__':
+  test_log_bucket()
+  build_relative_position(query_size=16, key_size=16, bucket_size=4, max_position=16)

modeling/deberta.py

@@ -0,0 +1,149 @@
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 01/15/2020
+#
+
+import copy
+import torch
+import os
+
+import json
+from .ops import *
+from .bert import *
+from .config import ModelConfig
+from .cache_utils import load_model_state
+import pdb
+
+__all__ = ['DeBERTa']
+
+class DeBERTa(torch.nn.Module):
+  """ DeBERTa encoder
+  This module is composed of the input embedding layer with stacked transformer layers with disentangled attention.
+
+  Parameters:
+    config:
+      A model config class instance with the configuration to build a new model. The schema is similar to `BertConfig`, \
+          for more details, please refer :class:`~DeBERTa.deberta.ModelConfig`
+
+    pre_trained:
+      The pre-trained DeBERTa model, it can be a physical path of a pre-trained DeBERTa model or a released configurations, \
+          i.e. [**base, large, base_mnli, large_mnli**]
+
+  """
+
+  def __init__(self, config=None, pre_trained=None):
+    super().__init__()
+    state = None
+    if pre_trained is not None:
+      state, model_config = load_model_state(pre_trained)
+      if config is not None and model_config is not None:
+        for k in config.__dict__:
+          if k not in ['hidden_size',
+            'intermediate_size',
+            'num_attention_heads',
+            'num_hidden_layers',
+            'vocab_size',
+            'max_position_embeddings']:
+            model_config.__dict__[k] = config.__dict__[k]
+      config = copy.copy(model_config)
+    self.embeddings = BertEmbeddings(config)
+    self.encoder = BertEncoder(config)
+    self.config = config
+    self.pre_trained = pre_trained
+    self.apply_state(state)
+
+  def forward(self, input_ids, attention_mask=None, token_type_ids=None, output_all_encoded_layers=True, position_ids = None, return_att = False):
+    """
+    Args:
+      input_ids:
+        a torch.LongTensor of shape [batch_size, sequence_length] \
+      with the word token indices in the vocabulary
+
+      attention_mask:
+        an optional parameter for input mask or attention mask.
+
+        - If it's an input mask, then it will be torch.LongTensor of shape [batch_size, sequence_length] with indices \
+      selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max \
+      input sequence length in the current batch. It's the mask that we typically use for attention when \
+      a batch has varying length sentences.
+
+        - If it's an attention mask then it will be torch.LongTensor of shape [batch_size, sequence_length, sequence_length]. \
+      In this case, it's a mask indicate which tokens in the sequence should be attended by other tokens in the sequence.
+
+      token_type_ids:
+        an optional torch.LongTensor of shape [batch_size, sequence_length] with the token \
+      types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to \
+      a `sentence B` token (see BERT paper for more details).
+
+      output_all_encoded_layers:
+        whether to output results of all encoder layers, default, True
+
+    Returns:
+
+      - The output of the stacked transformer layers if `output_all_encoded_layers=True`, else \
+      the last layer of stacked transformer layers
+
+      - Attention matrix of self-attention layers if `return_att=True`
+
+
+    Example::
+
+      # Batch of wordPiece token ids.
+      # Each sample was padded with zero to the maxium length of the batch
+      input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
+      # Mask of valid input ids
+      attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
+
+      # DeBERTa model initialized with pretrained base model
+      bert = DeBERTa(pre_trained='base')
+
+      encoder_layers = bert(input_ids, attention_mask=attention_mask)
+
+    """
+
+    if attention_mask is None:
+      attention_mask = torch.ones_like(input_ids)
+    if token_type_ids is None:
+      token_type_ids = torch.zeros_like(input_ids)
+      token_mask = torch.ones_like(input_ids)
+    else:
+      idxs = torch.flip(torch.cumsum(torch.flip(token_type_ids, [-1]), axis=1), [-1])
+      token_mask = idxs > 0
+      token_mask = token_mask.byte()
+    ebd_output = self.embeddings(input_ids.to(torch.long), token_type_ids.to(torch.long), position_ids, token_mask)
+    embedding_output = ebd_output['embeddings']
+    encoder_output = self.encoder(embedding_output,
+                   attention_mask,
+                   output_all_encoded_layers=output_all_encoded_layers, return_att = return_att)
+    encoder_output.update(ebd_output)
+    return encoder_output
+
+  def apply_state(self, state = None):
+    """ Load state from previous loaded model state dictionary.
+
+      Args:
+        state (:obj:`dict`, optional): State dictionary as the state returned by torch.module.state_dict(), default: `None`. \
+            If it's `None`, then will use the pre-trained state loaded via the constructor to re-initialize \
+            the `DeBERTa` model
+    """
+    if self.pre_trained is None and state is None:
+      return
+    if state is None:
+      state, config = load_model_state(self.pre_trained)
+      self.config = config
+    
+    prefix = ''
+    for k in state:
+      if 'embeddings.' in k:
+        if not k.startswith('embeddings.'):
+          prefix = k[:k.index('embeddings.')]
+        break
+
+    missing_keys = []
+    unexpected_keys = []
+    error_msgs = []
+    self._load_from_state_dict(state, prefix = prefix, local_metadata=None, strict=True, missing_keys=missing_keys, unexpected_keys=unexpected_keys, error_msgs=error_msgs)

modeling/disentangled_attention.py

@@ -0,0 +1,212 @@
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 01/15/2020
+#
+
+"""
+  Disentangled SelfAttention module
+"""
+
+import numpy as np
+import math
+import torch
+from torch import nn
+import functools
+import pdb
+
+from .ops import *
+from .da_utils import build_relative_position
+
+import loguru
+logger=loguru.logger
+
+__all__=['DisentangledSelfAttention']
+class DisentangledSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        _attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.attention_head_size = getattr(config, 'attention_head_size', _attention_head_size)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+
+        self.share_att_key = getattr(config, 'share_att_key', False)
+        self.pos_att_type = [x.strip() for x in getattr(config, 'pos_att_type', 'c2p').lower().split('|')] # c2p|p2c
+        self.relative_attention = getattr(config, 'relative_attention', False)
+
+        if self.relative_attention:
+            self.position_buckets = getattr(config, 'position_buckets', -1)
+            self.max_relative_positions = getattr(config, 'max_relative_positions', -1)
+            if self.max_relative_positions <1:
+                self.max_relative_positions = config.max_position_embeddings
+            self.pos_ebd_size = self.max_relative_positions
+            if self.position_buckets>0:
+                self.pos_ebd_size = self.position_buckets
+                # For backward compitable
+
+            self.pos_dropout = StableDropout(config.hidden_dropout_prob)
+
+            if (not self.share_att_key):
+                if 'c2p' in self.pos_att_type or 'p2p' in self.pos_att_type:
+                    self.pos_key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+                if 'p2c' in self.pos_att_type or 'p2p' in self.pos_att_type:
+                    self.pos_query_proj = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = StableDropout(config.attention_probs_dropout_prob)
+        self._register_load_state_dict_pre_hook(self._pre_load_hook)
+
+    def transpose_for_scores(self, x, attention_heads):
+        new_x_shape = x.size()[:-1] + (attention_heads, -1)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
+
+    def forward(self, hidden_states, attention_mask, return_att=False, query_states=None, relative_pos=None, rel_embeddings=None):
+        if query_states is None:
+            query_states = hidden_states
+        query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads).float()
+        key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads).float()
+        value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
+        
+        rel_att = None
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        scale_factor = 1
+        if 'c2p' in self.pos_att_type:
+            scale_factor += 1
+        if 'p2c' in self.pos_att_type:
+            scale_factor += 1
+        if 'p2p' in self.pos_att_type:
+            scale_factor += 1
+        scale = 1/math.sqrt(query_layer.size(-1)*scale_factor)
+        attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)*scale)
+        if self.relative_attention:
+            rel_embeddings = self.pos_dropout(rel_embeddings)
+            rel_att = self.disentangled_attention_bias(query_layer, key_layer, relative_pos, rel_embeddings, scale_factor)
+
+        if rel_att is not None:
+            attention_scores = (attention_scores + rel_att)
+        attention_scores = (attention_scores - attention_scores.max(dim=-1, keepdim=True).values.detach()).to(hidden_states)
+        attention_scores = attention_scores.view(-1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1))
+
+        # bxhxlxd
+        _attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
+        attention_probs = self.dropout(_attention_probs)
+        context_layer = torch.bmm(attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer)
+        context_layer = context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1)).permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (-1,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        return {
+            'hidden_states': context_layer,
+            'attention_probs': _attention_probs,
+            'attention_logits': attention_scores
+            }
+
+    def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
+        if relative_pos is None:
+            q = query_layer.size(-2)
+            relative_pos = build_relative_position(q, key_layer.size(-2), bucket_size = self.position_buckets, max_position = self.max_relative_positions)
+        if relative_pos.dim()==2:
+            relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+        elif relative_pos.dim()==3:
+            relative_pos = relative_pos.unsqueeze(1)
+        # bxhxqxk
+        elif relative_pos.dim()!=4:
+            raise ValueError(f'Relative postion ids must be of dim 2 or 3 or 4. {relative_pos.dim()}')
+
+        att_span = self.pos_ebd_size
+        relative_pos = relative_pos.long().to(query_layer.device)
+
+        rel_embeddings = rel_embeddings[self.pos_ebd_size - att_span:self.pos_ebd_size + att_span, :].unsqueeze(0) #.repeat(query_layer.size(0)//self.num_attention_heads, 1, 1)
+        if self.share_att_key:
+            pos_query_layer = self.transpose_for_scores(self.query_proj(rel_embeddings), self.num_attention_heads)\
+                .repeat(query_layer.size(0)//self.num_attention_heads, 1, 1) #.split(self.all_head_size, dim=-1)
+            pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads)\
+                .repeat(query_layer.size(0)//self.num_attention_heads, 1, 1) #.split(self.all_head_size, dim=-1)
+        else:
+            if 'c2p' in self.pos_att_type or 'p2p' in self.pos_att_type:
+                pos_key_layer = self.transpose_for_scores(self.pos_key_proj(rel_embeddings), self.num_attention_heads)\
+                    .repeat(query_layer.size(0)//self.num_attention_heads, 1, 1) #.split(self.all_head_size, dim=-1)
+            if 'p2c' in self.pos_att_type or 'p2p' in self.pos_att_type:
+                pos_query_layer = self.transpose_for_scores(self.pos_query_proj(rel_embeddings), self.num_attention_heads)\
+                    .repeat(query_layer.size(0)//self.num_attention_heads, 1, 1) #.split(self.all_head_size, dim=-1)
+
+        score = 0
+        # content->position
+        if 'c2p' in self.pos_att_type:
+            scale = 1/math.sqrt(pos_key_layer.size(-1)*scale_factor)
+            c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2).to(query_layer)*scale)
+            c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span*2-1)
+            c2p_att = torch.gather(c2p_att, dim=-1, index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]))
+            try:
+                score += c2p_att
+            except:
+                print(c2p_att.size())
+
+        # position->content
+        if 'p2c' in self.pos_att_type or 'p2p' in self.pos_att_type:
+            scale = 1/math.sqrt(pos_query_layer.size(-1)*scale_factor)
+            if key_layer.size(-2) != query_layer.size(-2):
+                r_pos = build_relative_position(key_layer.size(-2), key_layer.size(-2), bucket_size = self.position_buckets, max_position = self.max_relative_positions).to(query_layer.device)
+                r_pos = r_pos.unsqueeze(0)
+            else:
+                r_pos = relative_pos
+
+            p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span*2-1)
+            if query_layer.size(-2) != key_layer.size(-2):
+                pos_index = relative_pos[:, :, :, 0].unsqueeze(-1)
+
+        if 'p2c' in self.pos_att_type:
+            p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2).to(key_layer)*scale)
+            p2c_att = torch.gather(p2c_att, dim=-1, index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)])).transpose(-1,-2)
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2c_att = torch.gather(p2c_att, dim=-2, index=pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2))))
+            score += p2c_att
+
+        # position->position
+        if 'p2p' in self.pos_att_type:
+            pos_query = pos_query_layer[:,:,att_span:,:]
+            p2p_att = torch.matmul(pos_query, pos_key_layer.transpose(-1, -2))
+            p2p_att = p2p_att.expand(query_layer.size()[:2] + p2p_att.size()[2:])
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2p_att = torch.gather(p2p_att, dim=-2, index=pos_index.expand(query_layer.size()[:2] + (pos_index.size(-2), p2p_att.size(-1))))
+            p2p_att = torch.gather(p2p_att, dim=-1, index=c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)]))
+            score += p2p_att
+
+        return score
+
+    def _pre_load_hook(self, state_dict, prefix, local_metadata, strict,
+        missing_keys, unexpected_keys, error_msgs):
+        self_state = self.state_dict()
+        if ((prefix + 'query_proj.weight') not in state_dict) and ((prefix + 'in_proj.weight') in state_dict):
+          v1_proj = state_dict[prefix+'in_proj.weight']
+          v1_proj = v1_proj.unsqueeze(0).reshape(self.num_attention_heads, -1, v1_proj.size(-1))
+          q,k,v=v1_proj.chunk(3, dim=1)
+          state_dict[prefix + 'query_proj.weight'] = q.reshape(-1, v1_proj.size(-1))
+          state_dict[prefix + 'key_proj.weight'] = k.reshape(-1, v1_proj.size(-1))
+          state_dict[prefix + 'key_proj.bias'] = self_state['key_proj.bias']
+          state_dict[prefix + 'value_proj.weight'] = v.reshape(-1, v1_proj.size(-1))
+          v1_query_bias = state_dict[prefix + 'q_bias']
+          state_dict[prefix + 'query_proj.bias'] = v1_query_bias
+          v1_value_bias = state_dict[prefix +'v_bias']
+          state_dict[prefix + 'value_proj.bias'] = v1_value_bias
+
+          v1_pos_key_proj = state_dict[prefix + 'pos_proj.weight']
+          state_dict[prefix + 'pos_key_proj.weight'] = v1_pos_key_proj
+          v1_pos_query_proj = state_dict[prefix + 'pos_q_proj.weight']
+          state_dict[prefix + 'pos_query_proj.weight'] = v1_pos_query_proj
+          v1_pos_query_proj_bias = state_dict[prefix + 'pos_q_proj.bias']
+          state_dict[prefix + 'pos_query_proj.bias'] = v1_pos_query_proj_bias
+          state_dict[prefix + 'pos_key_proj.bias'] = self_state['pos_key_proj.bias']
+
+          del state_dict[prefix + 'in_proj.weight']
+          del state_dict[prefix + 'q_bias']
+          del state_dict[prefix + 'v_bias']
+          del state_dict[prefix + 'pos_proj.weight']
+          del state_dict[prefix + 'pos_q_proj.weight']
+          del state_dict[prefix + 'pos_q_proj.bias']

modeling/file_utils.py

@@ -0,0 +1,239 @@
+"""
+Utilities for working with the local dataset cache.
+This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+Copyright by the AllenNLP authors.
+"""
+
+import os
+import logging
+import shutil
+import tempfile
+import json
+from urllib.parse import urlparse
+from pathlib import Path
+from typing import Optional, Tuple, Union, IO, Callable, Set
+from hashlib import sha256
+from functools import wraps
+
+from tqdm import tqdm
+
+import boto3
+from botocore.exceptions import ClientError
+import requests
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+PYTORCH_PRETRAINED_BERT_CACHE = Path(os.getenv('PYTORCH_PRETRAINED_BERT_CACHE',
+                                               Path.home() / '.pytorch_pretrained_bert'))
+
+
+def url_to_filename(url: str, etag: str = None) -> str:
+    """
+    Convert `url` into a hashed filename in a repeatable way.
+    If `etag` is specified, append its hash to the url's, delimited
+    by a period.
+    """
+    url_bytes = url.encode('utf-8')
+    url_hash = sha256(url_bytes)
+    filename = url_hash.hexdigest()
+
+    if etag:
+        etag_bytes = etag.encode('utf-8')
+        etag_hash = sha256(etag_bytes)
+        filename += '.' + etag_hash.hexdigest()
+
+    return filename
+
+
+def filename_to_url(filename: str, cache_dir: Union[str, Path] = None) -> Tuple[str, str]:
+    """
+    Return the url and etag (which may be ``None``) stored for `filename`.
+    Raise ``FileNotFoundError`` if `filename` or its stored metadata do not exist.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BERT_CACHE
+    if isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    cache_path = os.path.join(cache_dir, filename)
+    if not os.path.exists(cache_path):
+        raise FileNotFoundError("file {} not found".format(cache_path))
+
+    meta_path = cache_path + '.json'
+    if not os.path.exists(meta_path):
+        raise FileNotFoundError("file {} not found".format(meta_path))
+
+    with open(meta_path) as meta_file:
+        metadata = json.load(meta_file)
+    url = metadata['url']
+    etag = metadata['etag']
+
+    return url, etag
+
+
+def cached_path(url_or_filename: Union[str, Path], cache_dir: Union[str, Path] = None) -> str:
+    """
+    Given something that might be a URL (or might be a local path),
+    determine which. If it's a URL, download the file and cache it, and
+    return the path to the cached file. If it's already a local path,
+    make sure the file exists and then return the path.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BERT_CACHE
+    if isinstance(url_or_filename, Path):
+        url_or_filename = str(url_or_filename)
+    if isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    parsed = urlparse(url_or_filename)
+
+    if parsed.scheme in ('http', 'https', 's3'):
+        # URL, so get it from the cache (downloading if necessary)
+        return get_from_cache(url_or_filename, cache_dir)
+    elif os.path.exists(url_or_filename):
+        # File, and it exists.
+        return url_or_filename
+    elif parsed.scheme == '':
+        # File, but it doesn't exist.
+        raise FileNotFoundError("file {} not found".format(url_or_filename))
+    else:
+        # Something unknown
+        raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
+
+
+def split_s3_path(url: str) -> Tuple[str, str]:
+    """Split a full s3 path into the bucket name and path."""
+    parsed = urlparse(url)
+    if not parsed.netloc or not parsed.path:
+        raise ValueError("bad s3 path {}".format(url))
+    bucket_name = parsed.netloc
+    s3_path = parsed.path
+    # Remove '/' at beginning of path.
+    if s3_path.startswith("/"):
+        s3_path = s3_path[1:]
+    return bucket_name, s3_path
+
+
+def s3_request(func: Callable):
+    """
+    Wrapper function for s3 requests in order to create more helpful error
+    messages.
+    """
+
+    @wraps(func)
+    def wrapper(url: str, *args, **kwargs):
+        try:
+            return func(url, *args, **kwargs)
+        except ClientError as exc:
+            if int(exc.response["Error"]["Code"]) == 404:
+                raise FileNotFoundError("file {} not found".format(url))
+            else:
+                raise
+
+    return wrapper
+
+
+@s3_request
+def s3_etag(url: str) -> Optional[str]:
+    """Check ETag on S3 object."""
+    s3_resource = boto3.resource("s3")
+    bucket_name, s3_path = split_s3_path(url)
+    s3_object = s3_resource.Object(bucket_name, s3_path)
+    return s3_object.e_tag
+
+
+@s3_request
+def s3_get(url: str, temp_file: IO) -> None:
+    """Pull a file directly from S3."""
+    s3_resource = boto3.resource("s3")
+    bucket_name, s3_path = split_s3_path(url)
+    s3_resource.Bucket(bucket_name).download_fileobj(s3_path, temp_file)
+
+
+def http_get(url: str, temp_file: IO) -> None:
+    req = requests.get(url, stream=True)
+    content_length = req.headers.get('Content-Length')
+    total = int(content_length) if content_length is not None else None
+    progress = tqdm(unit="B", total=total)
+    for chunk in req.iter_content(chunk_size=1024):
+        if chunk: # filter out keep-alive new chunks
+            progress.update(len(chunk))
+            temp_file.write(chunk)
+    progress.close()
+
+
+def get_from_cache(url: str, cache_dir: Union[str, Path] = None) -> str:
+    """
+    Given a URL, look for the corresponding dataset in the local cache.
+    If it's not there, download it. Then return the path to the cached file.
+    """
+    if cache_dir is None:
+        cache_dir = PYTORCH_PRETRAINED_BERT_CACHE
+    if isinstance(cache_dir, Path):
+        cache_dir = str(cache_dir)
+
+    os.makedirs(cache_dir, exist_ok=True)
+
+    # Get eTag to add to filename, if it exists.
+    if url.startswith("s3://"):
+        etag = s3_etag(url)
+    else:
+        response = requests.head(url, allow_redirects=True)
+        if response.status_code != 200:
+            raise IOError("HEAD request failed for url {} with status code {}"
+                          .format(url, response.status_code))
+        etag = response.headers.get("ETag")
+
+    filename = url_to_filename(url, etag)
+
+    # get cache path to put the file
+    cache_path = os.path.join(cache_dir, filename)
+
+    if not os.path.exists(cache_path):
+        # Download to temporary file, then copy to cache dir once finished.
+        # Otherwise you get corrupt cache entries if the download gets interrupted.
+        with tempfile.NamedTemporaryFile() as temp_file:
+            logger.info("%s not found in cache, downloading to %s", url, temp_file.name)
+
+            # GET file object
+            if url.startswith("s3://"):
+                s3_get(url, temp_file)
+            else:
+                http_get(url, temp_file)
+
+            # we are copying the file before closing it, so flush to avoid truncation
+            temp_file.flush()
+            # shutil.copyfileobj() starts at the current position, so go to the start
+            temp_file.seek(0)
+
+            logger.info("copying %s to cache at %s", temp_file.name, cache_path)
+            with open(cache_path, 'wb') as cache_file:
+                shutil.copyfileobj(temp_file, cache_file)
+
+            logger.info("creating metadata file for %s", cache_path)
+            meta = {'url': url, 'etag': etag}
+            meta_path = cache_path + '.json'
+            with open(meta_path, 'w') as meta_file:
+                json.dump(meta, meta_file)
+
+            logger.info("removing temp file %s", temp_file.name)
+
+    return cache_path
+
+
+def read_set_from_file(filename: str) -> Set[str]:
+    '''
+    Extract a de-duped collection (set) of text from a file.
+    Expected file format is one item per line.
+    '''
+    collection = set()
+    with open(filename, 'r', encoding='utf-8') as file_:
+        for line in file_:
+            collection.add(line.rstrip())
+    return collection
+
+
+def get_file_extension(path: str, dot=True, lower: bool = True):
+    ext = os.path.splitext(path)[1]
+    ext = ext if dot else ext[1:]
+    return ext.lower() if lower else ext

modeling/flash.py

@@ -0,0 +1,794 @@
+#
+# Zhoubo
+# 
+"""
+    FLASH: https://arxiv.org/abs/2202.10447
+"""
+import copy
+import torch
+import os
+from collections import Sequence
+import json
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.activations import ACT2FN
+from .modeling import *
+from .ops import XSoftmax, sequence_masking
+
+from .bert import *
+from .config import ModelConfig
+from .cache_utils import load_model_state
+import einops
+
+
+class ScaleNorm(nn.Module):
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.scala = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        mean_square = (x ** 2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(mean_square + self.eps) * self.scala
+        return x
+
+
+
+class OffsetScale(nn.Module):
+    def __init__(self, dim, heads = 1):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+        self.beta = nn.Parameter(torch.zeros(heads, dim))
+        # nn.init.normal_(self.gamma, std = 0.02)
+        # nn.init.xavier_uniform_(self.gamma)
+
+    def forward(self, x):
+        out = (x * self.gamma) + self.beta
+        return out
+
+
+class ScaledSinuEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(1,))
+        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x):
+        n, device = x.shape[1], x.device
+        t = torch.arange(n, device = device).type_as(self.inv_freq)
+        sinu = torch.einsum('i , j -> i j', t, self.inv_freq)
+        emb = torch.cat((sinu.sin(), sinu.cos()), dim = -1)
+        return emb * self.scale
+
+
+def RoPE(x, dim):
+    """
+    :param x: input tensor
+    :param dim: oprate dimension
+    :return: tensor
+    """
+    shape = x.shape
+    if isinstance(dim, int):
+        dim = [dim]
+
+    spatial_shape = [shape[i] for i in dim]
+    total_len = 1
+    for i in spatial_shape:
+        total_len *= i
+    position = torch.reshape(torch.arange(total_len, dtype=torch.float, device=x.device), spatial_shape)
+
+    for i in range(dim[-1] + 1, len(shape) - 1, 1):
+        position = torch.unsqueeze(position, dim=-1)
+
+    half_size = shape[-1] // 2
+    freq_seq = -torch.arange(half_size, dtype=torch.float, device=x.device) / float(half_size)
+    inv_freq = 10000 ** -freq_seq
+    sinusoid = torch.einsum("...,d->...d", position, inv_freq)
+    sin = torch.sin(sinusoid).repeat_interleave(2, -1)
+    cos = torch.cos(sinusoid).repeat_interleave(2, -1)
+    tensor_cross = torch.stack([-x[..., 1:: 2], x[..., :: 2]], -1).reshape(x.shape)
+    # x1, x2 = torch.chunk(x, 2, dim=-1)
+    # return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+    return x * cos + tensor_cross * sin
+
+
+def rel_pos_bias(seq_len, s):
+    a = torch.rand([1, s], dtype=torch.float)
+    b = torch.rand([1, s], dtype=torch.float)
+    w = torch.rand([2 * seq_len - 1], dtype=torch.float)
+    if seq_len <= 512:
+        t = F.pad(w[: 2 * seq_len - 1], [0, seq_len]).repeat(seq_len)
+        t = t[..., :-seq_len].reshape(-1, seq_len, 3 * seq_len - 2)
+        r = (2 * seq_len - 1) // 2
+        t = t[..., r:-r]
+    else:
+        a = RoPE(a.repeat(seq_len, 1), dim=[0])
+        b = RoPE(b.repeat(seq_len, 1), dim=[0])
+        t = torch.einsum("mk,nk->mn", a, b)
+    return t
+
+def squared_relu(x, attention_mask, dim=-1):
+    rmask = ~(attention_mask.bool())
+    x = x.masked_fill(rmask, 0)
+    return torch.square(F.relu(x))
+
+
+def attention_normalize(a, axis=-1, mask=None, fn='softmax'):
+    if fn == 'softmax':
+        return XSoftmax.apply(a, mask, axis)
+    else:
+        mask_ = a > -float('inf') / 10
+        # mask_ = mask_.byte()
+        mask_ = torch.sum(mask_, axis=axis, keepdim=True)
+        l = torch.maximum(mask_, torch.ones_like(mask_))
+        if fn == 'squared_relu':
+            rmask = ~(mask.bool())
+            a = a.masked_fill(rmask, 0)
+            return torch.square(F.relu(a)) / l
+        elif fn == 'softmax_plus':
+            return XSoftmax.apply(a * torch.log(l) / np.log(512), mask, axis)
+    return a
+
+
+class GAULinear(nn.Linear):
+    def init_weight(self):
+        nn.init.xavier_uniform_(self.weight)
+
+
+class GatedAttentionUnit(nn.Module):
+    """
+    GAU Block: Gate Attention Unit
+    """
+    def __init__(
+        self,
+        max_seq_length,
+        hidden_size,
+        attention_key_size=128,
+        activation='swish',
+        use_bias=True,
+        attention_norm_type='squared_relu',
+        attention_scale=True,
+        dropout=0.1,
+        pre_norm=False,
+        norm_type="layer_norm", 
+        eps=1e-5,
+        shift_token=False, 
+        use_rel_bias=False,
+        add_residual=True,
+        **kwargs,):
+
+        super(GatedAttentionUnit, self).__init__(**kwargs)
+        self.max_seq_length = max_seq_length
+        self.units = hidden_size
+        self.intermediate_size = self.units * 2
+        self.key_size = attention_key_size
+        self.activation = activation
+        self.use_bias = use_bias
+        self.attention_norm_type = attention_norm_type
+        self.attention_scale = attention_scale
+        self.dropout = StableDropout(dropout)
+        self.i_dense = nn.Sequential(
+            nn.Linear(self.units, 2 * self.intermediate_size + self.key_size, bias=self.use_bias),
+            nn.SiLU()
+        )
+        self.o_dense = nn.Sequential(
+                        nn.Linear(self.intermediate_size, self.units, bias=self.use_bias),
+                        self.dropout)
+        self.q_scaleoffset = OffsetScale(self.key_size)
+        self.k_scaleoffset = OffsetScale(self.key_size)
+        self.pre_norm = pre_norm
+        self.norm = (nn.LayerNorm(hidden_size, eps=eps) if norm_type.lower() == "layer_norm" else ScaleNorm(eps=eps))
+        self.add_residual = add_residual
+
+    def forward(self, x, attention_mask=None, **kwargs):
+        shortcut = x
+        
+        if self.pre_norm:
+            x = self.norm(x)
+
+        x = self.i_dense(x)
+        u, v, qk = torch.split(x, [self.intermediate_size, self.intermediate_size, self.key_size], dim=-1)
+        q, k = self.q_scaleoffset(qk), self.k_scaleoffset(qk)
+        qk = RoPE(torch.stack([q, k], 2), dim=1)
+        q, k = qk[:, :, 0], qk[:, :, 1]
+        a = torch.einsum('bmd,bnd->bmn', q, k)
+        if self.attention_scale:
+            a = a / self.key_size**0.5
+        a = sequence_masking(a, attention_mask, '-inf', -1)
+        A = attention_normalize(a, -1, fn=self.attention_norm_type)
+        if self.dropout:
+            A = self.dropout(A)
+        out = self.o_dense(u * torch.einsum('bmn,bnd->bmd', A, v))
+
+        if self.add_residual:
+            out = out + shortcut
+        if not self.pre_norm:
+            out = self.norm(out)
+        return out
+        # # 加入RoPE
+        # if p_bias == 'rotary':
+        #     qk = K.stack([q, k], 2)
+        #     qk = apply_rotary_position_embeddings(inputs[n], qk)[0]
+        #     q, k = qk[:, :, 0], qk[:, :, 1]
+        # # Attention
+        # a = tf.einsum('bmd,bnd->bmn', q, k)
+        # if self.attention_scale:
+        #     a = a / self.key_size**0.5
+        # if a_bias is not None:
+        #     a = a + a_bias
+        # a = sequence_masking(a, mask, '-inf', -1)
+        # A = attention_normalize(a, -1, self.normalization)
+        # if self.attention_dropout:
+        #     A = Dropout(self.attention_dropout)(A)
+        # # 计算输出
+        # o = self.o_dense(u * tf.einsum('bmn,bnd->bmd', A, v))
+        
+        # return o
+
+class GAU(nn.Module):
+    def __init__(self, max_seq_length, hidden_size, expansion_factor=2, s=128, norm_type="layer_norm", eps=1e-5,
+                 hidden_act="silu", shift_token=False, use_rel_bias=False, attention_norm_type='softmax',
+                 pre_norm=False, dropout=0, add_residual = True):
+        super(GAU, self).__init__()
+        self.max_seq_length = max_seq_length
+        self.shift_token = shift_token
+        hidden_dim = int(expansion_factor * hidden_size)
+        self.norm = (nn.LayerNorm(hidden_size, eps=eps) if norm_type == "layer_norm" else ScaleNorm(eps=eps))
+        self.use_rel_bias = use_rel_bias
+        self.attention_norm_type = attention_norm_type
+        # if attention_norm_type == 'relu':
+        #     self.attention_norm_func = squared_relu
+        # else:
+        #     self.attention_norm_func = XSoftmax.apply
+        # self.norm = norm_klass(hidden_size)
+
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_hidden = nn.Sequential(
+            nn.Linear(hidden_size, hidden_dim * 2),
+            nn.SiLU()
+        )
+
+        self.to_qk = nn.Sequential(
+            nn.Linear(hidden_size, s),
+            nn.SiLU()
+        )
+
+        self.offsetscale = OffsetScale(s, heads = 2)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_size),
+            nn.Dropout(dropout)
+        )
+
+        self.add_residual = add_residual
+        self.act_fn = ACT2FN[hidden_act]
+        self.pre_norm = pre_norm
+
+
+    def forward(
+        self,
+        x,
+        relative_pos = None,
+        attention_mask = None
+    ):
+        seq_len, device = x.shape[-2], x.device
+        if self.pre_norm:
+            normed_x = self.norm(x)
+        else:
+            normed_x = x
+        v, gate = self.to_hidden(normed_x).chunk(2, dim = -1)
+
+        qk = self.to_qk(normed_x)
+        base = self.offsetscale(qk)
+        base = RoPE(base, 1)
+        q, k = base.unbind(dim = -2)
+        sim = torch.einsum('b i d, b j d -> b i j', q, k)
+
+        if relative_pos is not None:
+            sim = sim + relative_pos
+        if attention_mask is not None:
+            if attention_mask.dim() < 3:
+                attention_mask = einops.rearrange(attention_mask, 'b j -> b 1 j')
+            # attn = attn.masked_fill(~attention_mask.bool(), 0.)
+        attn = attention_normalize(sim, mask=attention_mask, fn=self.attention_norm_type)
+        # attn = F.relu(sim) ** 2 / seq_len# / q.size(-1)
+        # logger.info(attn.max())
+        attn = self.dropout(attn)
+        # if self.causal:
+        #     causal_mask = torch.ones((seq_len, seq_len), dtype = torch.bool, device = device).triu(1)
+        #     attn = attn.masked_fill(causal_mask, 0.)
+
+        out = torch.einsum('b i j, b j d -> b i d', attn, v)
+        out = out * gate
+
+        out = self.to_out(out)
+
+        if self.add_residual:
+            out = out + x
+        if not self.pre_norm:
+            out = self.norm(out)
+        return out
+
+
+class GAULayer(nn.Module):
+    def __init__(self, config, shift_token=False, use_ffn=False):
+        super(GAULayer, self).__init__()
+        self.attention = GatedAttentionUnit(config.max_position_embeddings, config.hidden_size, 
+                                    shift_token=shift_token, use_rel_bias=config.use_rel_bias, 
+                                    norm_type=config.norm_type, attention_norm_type=config.attention_norm_type,
+                                    pre_norm=config.pre_norm, dropout=config.hidden_dropout_prob)
+        if use_ffn:
+            self.intermediate = BertIntermediate(config)
+            self.output = BertOutput(config)
+        self.use_ffn = use_ffn
+
+    def forward(self, hidden_states, attention_mask, return_att=False, query_states=None, relative_pos=None, rel_embeddings=None):
+        attention_output = self.attention(hidden_states, attention_mask=attention_mask, relative_pos=relative_pos)
+        if self.use_ffn:
+            intermediate_output = self.intermediate(attention_output)
+            layer_output = self.output(intermediate_output, attention_output)
+            return layer_output
+        else:
+            return attention_output
+
+
+class FlashBlock(nn.Module):
+    """
+    FLASH Block: Fast Linear Attention with a Single Head
+    """
+
+    def __init__(self, model_size, sequence_length, chunk_size=256, expansion_factor=2, s=128, norm_type="layer_norm", eps=1e-5,
+                 hidden_act="silu"):
+        super(FlashBlock, self).__init__()
+        self.s = s
+        self.eps = eps
+        self.norm_type = norm_type
+        self.model_size = model_size
+        self.chunk_size = chunk_size
+        self.hidden_act = hidden_act
+        self.sequence_length = sequence_length
+        self.expansion_factor = expansion_factor
+        self.e = int(self.model_size * self.expansion_factor)
+
+        self.dense1 = nn.Linear(self.model_size, 2 * self.e + self.s, bias=True)
+        self.gamma = nn.Parameter(torch.rand((4, self.s)))
+        self.beta = nn.Parameter(torch.rand((4, self.s)))
+        self.dense2 = nn.Linear(self.e, self.model_size)
+        self.LayerNorm = (
+            nn.LayerNorm(model_size, eps=self.eps) if norm_type == "layer_norm" else ScaleNorm(eps=self.eps))
+
+        nn.init.xavier_normal_(self.dense1.weight)
+        self.act_fn = ACT2FN(self.hidden_act)
+
+    def global_linear_attention(self, query, key, value, causal):
+        if causal:
+            kv = torch.einsum("bgcs, bgce->bgse", key, value)
+            kv = torch.cumsum(kv, dim=1)
+            lin_v = torch.einsum("bgcs, bgse->bgce", query, kv)
+            return lin_v
+        else:
+            kv = torch.einsum("bgcs, bgce->bse", key, value)
+            lin_v = torch.einsum("bgcs, bse->bgce", query, kv)
+            return lin_v
+
+    def segment_ids_to_mask(self, segment_ids, causal=False):
+        """Generate the segment mask from the segment ids.
+        The segment mask is used to remove the attention between tokens in different documents.
+        """
+        min_ids, max_ids = torch.min(segment_ids, dim=-1).values, torch.max(segment_ids, dim=-1).values
+        # 1.0 indicates in the same group and 0.0 otherwise
+        mask = torch.logical_and(torch.less_equal(min_ids[:, :, None], max_ids[:, None, :]),
+                                 torch.greater_equal(max_ids[:, :, None], min_ids[:, None, :]))
+        mask = torch.tensor(mask, torch.float32)
+        if causal:
+            g = segment_ids.size()[1]
+            causal_mask = 1.0 - torch.triu(torch.ones([g, g], dtype=torch.float32))  # 保留主对角线以及主对角线以上的元素
+            mask *= causal_mask
+        mask = torch.div(mask, torch.sum(mask, dim=-1, keepdim=True))
+        return mask
+
+    def forward(self, x, causal=False, attention_mask=None, sequence_mask=None, **kwargs):
+        """
+        inputs: [batch_size, num_chunk, chunk_length, model_size]
+        """
+        _, g, n, d = x.size()
+        shortcut, x = x, self.LayerNorm(x)
+        # 通过线性变换得到Z，见论文公式(4)
+        uv = self.dense1(x)
+        # 将uv按最后一维切分，得到Ug:[C*e],Vg:[C*e], Zg:[C*s], 论文中的3.2部分
+        # u:[batch_size, num_chunk, chunk_length, self.e]
+        # v:[batch_size, num_chunk, chunk_length, self.e]
+        # z:[batch_size, num_chunk, chunk_length, self.s]
+        u, v, z = torch.split(self.act_fn(uv), [self.e, self.e, self.s], dim=-1)
+
+        # 生���quad_q, quad_k, lin_q, lin_k
+        # 首先进行简单的offset和scale,融入RoPE位置向量
+        z = torch.einsum("...r, hr->...hr", z, self.gamma) + self.beta
+        z = RoPE(z, dim=[1, 2])
+        quad_q, quad_k, lin_q, lin_k = torch.unbind(z, dim=-2)  # 按-2维进行分解得到quad_q, quad_k, lin_q和lin_k
+        # 计算global的lin_v
+        lin_v = self.global_linear_attention(lin_q, lin_k, v, causal)
+        if causal:
+            # 线性注意力部分
+            lin_kv = torch.einsum("bgnk, bgne->bgke", lin_k, lin_v) / torch.tensor(n, x.dtype)  # 见公式(7)
+            mask = self.segment_ids_to_mask(segment_ids=segment_ids, causal=causal)
+            cum_lin_kv = torch.einsum('bhke, bgh->bgke', lin_kv, mask)
+            linear = torch.einsum("bgnk, bgke->bgne", lin_kv, cum_lin_kv)
+            # 二次注意力
+            quad_qk = torch.einsum("bgnk, bgmk->bgnm", quad_q, quad_k)  # 论文Local attention per chunk部分
+            bias = rel_pos_bias(self.sequence_length, self.s)[:, :n, :n]
+            kernel = torch.square(F.relu(quad_qk / n + bias))  # 论文中的relu**2部分
+            causal_mask = torch.triu(torch.ones([n, n], dtype=x.dtype))
+            quadratic = torch.einsum("bgnm, bgme->bgne", kernel * causal_mask, v)
+        else:
+            lin_kv = torch.einsum("bgnk, bgne->bgke", lin_k, lin_v) / torch.tensor(n, x.dtype)  # 见公式(7)
+            mask = self.segment_ids_to_mask(segment_ids=segment_ids, causal=causal)
+            lin_kv = torch.einsum("bhke, bgh->bgke", lin_kv, mask)
+            linear = torch.einsum("bgnk, bgke->bgne", lin_q, lin_kv)
+            # 二次注意力
+            quad_qk = torch.einsum("bgnk, bgmk->bgnm", quad_q, quad_k)  # 论文Local attention per chunk部分
+            bias = rel_pos_bias(self.sequence_length, self.s)[:, :n, :n]
+            kernel = torch.square(F.relu(quad_qk / n + bias))  # 论文中的relu**2部分
+            quadratic = torch.einsum("bgnm, bgme->bgne", kernel, v)
+        x = u * (quadratic + linear)
+        x = self.dense2(x)
+        x = x + shortcut
+        return x
+
+class RelativePositionBias(nn.Module):
+    def __init__(
+        self,
+        scale,
+        causal = False,
+        num_buckets = 32,
+        max_distance = 128
+    ):
+        super().__init__()
+        self.scale = scale
+        self.causal = causal
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = nn.Embedding(num_buckets, 1)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position,
+        causal = True,
+        num_buckets = 32,
+        max_distance = 128
+    ):
+        ret = 0
+        n = -relative_position
+        if not causal:
+            num_buckets //= 2
+            ret += (n < 0).long() * num_buckets
+            n = torch.abs(n)
+        else:
+            n = torch.max(n, torch.zeros_like(n))
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = max_exact + (
+            torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+        ).long()
+        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, x):
+        i, j, device = *x.shape[-2:], x.device
+        q_pos = torch.arange(i, dtype = torch.long, device = device)
+        k_pos = torch.arange(j, dtype = torch.long, device = device)
+        rel_pos = rearrange(k_pos, 'j -> 1 j') - rearrange(q_pos, 'i -> i 1')
+        rp_bucket = self._relative_position_bucket(rel_pos, causal = self.causal, num_buckets = self.num_buckets, max_distance = self.max_distance)
+        values = self.relative_attention_bias(rp_bucket)
+        bias = rearrange(values, 'i j 1 -> i j')
+        return bias * self.scale
+
+
+class FlashEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, config, with_position=False):
+        super(FlashEmbeddings, self).__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size)
+        self.with_position = with_position
+        if with_position:
+            self.position_embeddings = ScaledSinuEmbedding(config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-5)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids, token_type_ids=None, position_ids=None, token_mask=None):
+        seq_length = input_ids.size(1)
+        if position_ids is None:
+            position_ids = torch.arange(
+                seq_length, dtype=torch.long, device=input_ids.device)
+            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        words_embeddings = self.word_embeddings(input_ids)
+        if self.with_position:
+            position_embeddings = self.position_embeddings(words_embeddings)
+        else:
+            position_embeddings = 0
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        # if self.num_pos_emb > 1:
+        #     num_batch = position_embeddings.size(0)
+        #     num_pos = position_embeddings.size(1)
+        #     position_embeddings = position_embeddings.view(
+        #         num_batch, num_pos, self.num_pos_emb, -1)[torch.arange(0, num_batch).long(), :, task_idx, :]
+
+        embeddings = words_embeddings + position_embeddings + token_type_embeddings
+        # if self.fp32_embedding:
+        #     embeddings = embeddings.half()
+        embeddings = MaskedLayerNorm(self.LayerNorm, embeddings, token_mask)
+        embeddings = self.dropout(embeddings)
+        return {
+                'embeddings': embeddings,
+                'position_embeddings': position_embeddings}
+
+
+class GAUEncoder(nn.Module):
+    def __init__(self, config, shift_token=False):
+        super().__init__()
+        layer = GAULayer(config, shift_token=shift_token)
+        self.layer = nn.ModuleList([copy.deepcopy(layer)
+                                    for _ in range(config.num_hidden_layers)])
+    
+    def get_attention_mask(self, attention_mask):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1)
+            attention_mask = extended_attention_mask*extended_attention_mask.squeeze(-2).unsqueeze(-1)
+            attention_mask = attention_mask #.byte()
+        return attention_mask
+
+    def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True, return_att=False, query_states = None, relative_pos=None):
+        all_encoder_layers = []
+        att_matrices = []
+        if isinstance(hidden_states, Sequence):
+            next_kv = hidden_states[0]
+        else:
+            next_kv = hidden_states
+        # rel_embeddings = self.get_rel_embedding()
+        for i, layer_module in enumerate(self.layer):
+            output_states = layer_module(next_kv, attention_mask, query_states = query_states, relative_pos=relative_pos)
+            if return_att:
+                    output_states, att_m = output_states
+
+            # if i == 0 and self.with_conv:
+            #     prenorm = output_states #output['prenorm_states']
+            #     output_states = self.conv(hidden_states, prenorm, input_mask)
+
+            if query_states is not None:
+                query_states = output_states
+                if isinstance(hidden_states, Sequence):
+                    next_kv = hidden_states[i+1] if i+1 < len(self.layer) else None
+            else:
+                next_kv = output_states
+
+            if output_all_encoded_layers:
+                all_encoder_layers.append(output_states)
+                if return_att:
+                    att_matrices.append(att_m)
+        if not output_all_encoded_layers:
+            all_encoder_layers.append(output_states)
+            if return_att:
+                    att_matrices.append(att_m)
+        return {
+            'hidden_states': all_encoder_layers,
+            'attention_matrices': att_matrices
+            }
+
+class FlashEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__(config)
+        layer = GateAttentionUnit(config.max_position_embeddings, config.hidden_size)
+        self.layer = nn.ModuleList([copy.deepcopy(layer)
+                                    for _ in range(config.num_hidden_layers)])
+    
+    def forward(self, hidden_states, attention_mask, token_mask=None,
+                    output_all_encoded_layers=True, 
+                    prev_embedding=None, prev_encoded_layers=None, mask_qkv=None, seg_ids=None):
+        # history embedding and encoded layer must be simultanously given
+        assert (prev_embedding is None) == (prev_encoded_layers is None)
+
+        all_encoder_layers = []
+        if (prev_embedding is not None) and (prev_encoded_layers is not None):
+            history_states = prev_embedding
+            for i, layer_module in enumerate(self.layer):
+                hidden_states = layer_module(
+                    hidden_states, attention_mask, history_states=history_states, mask_qkv=mask_qkv, seg_ids=seg_ids)
+                if output_all_encoded_layers:
+                    all_encoder_layers.append(hidden_states)
+                if prev_encoded_layers is not None:
+                    history_states = prev_encoded_layers[i]
+        else:
+            for layer_module in self.layer:
+                hidden_states = layer_module(
+                    hidden_states, attention_mask=attention_mask, mask_qkv=mask_qkv, seg_ids=seg_ids)
+                if output_all_encoded_layers:
+                    all_encoder_layers.append(hidden_states)
+        if not output_all_encoded_layers:
+            all_encoder_layers.append(hidden_states)
+        return all_encoder_layers
+
+# class FlashQuadModel(BertModel):
+#     def __init__(self, config, pooler=False, shift_token=False, causal=False) -> None:
+#         super().__init__(config)
+#         self.embeddings = FlashEmbeddings(config)
+#         self.encoder = GAUEncoder(config, causal=causal, shift_token=shift_token)
+#         if not pooler:
+#             self.pooler = None
+#         self.apply(self.init_bert_weights)
+
+
+class FlashQuadModel(torch.nn.Module):
+    """ 
+    Parameters:
+        config:
+        A model config class instance with the configuration to build a new model. The schema is similar to `BertConfig`, 
+
+        pre_trained:
+        The pre-trained DeBERTa model, it can be a physical path of a pre-trained DeBERTa model or a released configurations,
+            i.e. [**base, large, base_mnli, large_mnli**]
+
+    """
+
+    def __init__(self, config=None, pre_trained=None, pooler=False, shift_token=False, causal=False, **kwargs):
+        super().__init__()
+        state = None
+        if pre_trained is not None:
+            state, model_config = load_model_state(pre_trained)
+            if config is not None and model_config is not None:
+                for k in config.__dict__:
+                    if k not in ['hidden_size',
+                            'intermediate_size',
+                            'num_attention_heads',
+                            'num_hidden_layers',
+                            'vocab_size',
+                            'max_position_embeddings']:
+                            model_config.__dict__[k] = config.__dict__[k]
+            config = copy.copy(model_config)
+        self.embeddings = FlashEmbeddings(config, with_position=True)
+        self.encoder = GAUEncoder(config, shift_token=shift_token)
+        if not pooler:
+            self.pooler = None
+        self.config = config
+        self.pre_trained = pre_trained
+        self.apply_state(state)
+
+    def get_attention_mask(self, input_ids=None, token_type_ids=None, attention_mask=None, input_mask=None):
+        if attention_mask is None:
+            if input_mask is not None:
+                return input_mask.unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), input_mask.size(1))
+            else:
+                return torch.ones_like(input_ids, dtype=torch.uint8).unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), input_mask.size(1))
+        else:
+            if attention_mask.dim() == 2:
+                if input_mask is not None:
+                    attention_mask = attention_mask * input_mask
+                return attention_mask.unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), attention_mask.size(-1))
+            if attention_mask.dim() == 4:
+                attention_mask = attention_mask.squeeze(2)
+            if attention_mask.dim() == 3:
+                if input_mask is not None:
+                    return attention_mask * input_mask.unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), attention_mask.size(-1))
+                else:
+                    return attention_mask
+            
+
+    def forward(self, input_ids, input_mask, attention_mask=None, token_type_ids=None, 
+                output_all_encoded_layers=True, position_ids=None, return_att=False):
+        """
+        Args:
+            input_ids:
+                a torch.LongTensor of shape [batch_size, sequence_length] \
+            with the word token indices in the vocabulary
+
+            attention_mask:
+                an optional parameter for input mask or attention mask.
+
+                - If it's an input mask, then it will be torch.LongTensor of shape [batch_size, sequence_length] with indices \
+            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max \
+            input sequence length in the current batch. It's the mask that we typically use for attention when \
+            a batch has varying length sentences.
+
+                - If it's an attention mask then it will be torch.LongTensor of shape [batch_size, sequence_length, sequence_length]. \
+            In this case, it's a mask indicate which tokens in the sequence should be attended by other tokens in the sequence.
+
+            token_type_ids:
+                an optional torch.LongTensor of shape [batch_size, sequence_length] with the token \
+            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to \
+            a `sentence B` token (see BERT paper for more details).
+
+            output_all_encoded_layers:
+                whether to output results of all encoder layers, default, True
+
+            Returns:
+
+            - The output of the stacked transformer layers if `output_all_encoded_layers=True`, else \
+            the last layer of stacked transformer layers
+
+            - Attention matrix of self-attention layers if `return_att=True`
+
+
+            Example::
+
+            # Batch of wordPiece token ids.
+            # Each sample was padded with zero to the maxium length of the batch
+            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
+            # Mask of valid input ids
+            attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
+
+            # DeBERTa model initialized with pretrained base model
+            bert = DeBERTa(pre_trained='base')
+
+            encoder_layers = bert(input_ids, attention_mask=attention_mask)
+
+            """
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+            # input_mask = torch.ones_like(input_ids)
+
+        if input_mask is None:
+            idxs = torch.flip(torch.cumsum(torch.flip(token_type_ids, [-1]), axis=1), [-1])
+            input_mask = idxs > 0
+            if not torch.any(input_mask):
+                input_mask = torch.ones_like(input_ids)
+            input_mask = input_mask # .byte()
+        attention_mask = self.get_attention_mask(input_ids, token_type_ids, attention_mask, input_mask)
+        attention_mask = attention_mask #.byte()
+        embedding_output = self.embeddings(input_ids.to(torch.long), token_type_ids.to(torch.long), position_ids, input_mask)
+        encoder_output = self.encoder(embedding_output['embeddings'], attention_mask, output_all_encoded_layers=output_all_encoded_layers, return_att = return_att)
+        encoder_output.update(embedding_output)
+        return encoder_output
+
+    def apply_state(self, state = None):
+        """ Load state from previous loaded model state dictionary.
+
+        Args:
+            state (:obj:`dict`, optional): State dictionary as the state returned by torch.module.state_dict(), default: `None`. \
+                If it's `None`, then will use the pre-trained state loaded via the constructor to re-initialize \
+                the `DeBERTa` model
+        """
+        if self.pre_trained is None and state is None:
+            return
+        if state is None:
+            state, config = load_model_state(self.pre_trained)
+            self.config = config
+        
+        prefix = ''
+        for k in state:
+            if 'embeddings.' in k:
+                    if not k.startswith('embeddings.'):
+                        prefix = k[:k.index('embeddings.')]
+                    break
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+        self._load_from_state_dict(state, prefix = prefix, local_metadata=None, strict=True, missing_keys=missing_keys, unexpected_keys=unexpected_keys, error_msgs=error_msgs)
+
+
+class FlashModel(BertModel):
+    def __init__(self, config) -> None:
+        super().__init__(config)
+        self.encoder = FlashEncoder(config)
+        self.apply(self.init_bert_weights)
+
+if __name__ == '__main__':
+    model = FlashModel(768, 64)

modeling/focal_loss.py

@@ -0,0 +1,200 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.cuda.amp as amp
+
+
+##
+# version 1: use torch.autograd
+class FocalLossV1(nn.Module):
+
+    def __init__(self,
+                 alpha=0.25,
+                 gamma=2,
+                 reduction='mean',):
+        super(FocalLossV1, self).__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.reduction = reduction
+        self.crit = nn.BCEWithLogitsLoss(reduction='none')
+
+    def forward(self, logits, label):
+        '''
+        Usage is same as nn.BCEWithLogits:
+            >>> criteria = FocalLossV1()
+            >>> logits = torch.randn(8, 19, 384, 384)
+            >>> lbs = torch.randint(0, 2, (8, 19, 384, 384)).float()
+            >>> loss = criteria(logits, lbs)
+        '''
+        probs = torch.sigmoid(logits)
+        coeff = torch.abs(label - probs).pow(self.gamma).neg()
+        log_probs = torch.where(logits >= 0,
+                F.softplus(logits, -1, 50),
+                logits - F.softplus(logits, 1, 50))
+        log_1_probs = torch.where(logits >= 0,
+                -logits + F.softplus(logits, -1, 50),
+                -F.softplus(logits, 1, 50))
+        loss = label * self.alpha * log_probs + (1. - label) * (1. - self.alpha) * log_1_probs
+        loss = loss * coeff
+
+        if self.reduction == 'mean':
+            loss = loss.mean()
+        if self.reduction == 'sum':
+            loss = loss.sum()
+        return loss
+
+
+##
+# version 2: user derived grad computation
+class FocalSigmoidLossFuncV2(torch.autograd.Function):
+    '''
+    compute backward directly for better numeric stability
+    '''
+    @staticmethod
+    @amp.custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, logits, label, alpha, gamma):
+        #  logits = logits.float()
+
+        probs = torch.sigmoid(logits)
+        coeff = (label - probs).abs_().pow_(gamma).neg_()
+        log_probs = torch.where(logits >= 0,
+                F.softplus(logits, -1, 50),
+                logits - F.softplus(logits, 1, 50))
+        log_1_probs = torch.where(logits >= 0,
+                -logits + F.softplus(logits, -1, 50),
+                -F.softplus(logits, 1, 50))
+        ce_term1 = log_probs.mul_(label).mul_(alpha)
+        ce_term2 = log_1_probs.mul_(1. - label).mul_(1. - alpha)
+        ce = ce_term1.add_(ce_term2)
+        loss = ce * coeff
+
+        ctx.vars = (coeff, probs, ce, label, gamma, alpha)
+
+        return loss
+
+    @staticmethod
+    @amp.custom_bwd
+    def backward(ctx, grad_output):
+        '''
+        compute gradient of focal loss
+        '''
+        (coeff, probs, ce, label, gamma, alpha) = ctx.vars
+
+        d_coeff = (label - probs).abs_().pow_(gamma - 1.).mul_(gamma)
+        d_coeff.mul_(probs).mul_(1. - probs)
+        d_coeff = torch.where(label < probs, d_coeff.neg(), d_coeff)
+        term1 = d_coeff.mul_(ce)
+
+        d_ce = label * alpha
+        d_ce.sub_(probs.mul_((label * alpha).mul_(2).add_(1).sub_(label).sub_(alpha)))
+        term2 = d_ce.mul(coeff)
+
+        grads = term1.add_(term2)
+        grads.mul_(grad_output)
+
+        return grads, None, None, None
+
+
+class FocalLossV2(nn.Module):
+
+    def __init__(self,
+                 alpha=0.25,
+                 gamma=2,
+                 reduction='mean'):
+        super(FocalLossV2, self).__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, logits, label):
+        '''
+        Usage is same as nn.BCEWithLogits:
+            >>> criteria = FocalLossV2()
+            >>> logits = torch.randn(8, 19, 384, 384)
+            >>> lbs = torch.randint(0, 2, (8, 19, 384, 384)).float()
+            >>> loss = criteria(logits, lbs)
+        '''
+        loss = FocalSigmoidLossFuncV2.apply(logits, label, self.alpha, self.gamma)
+        if self.reduction == 'mean':
+            loss = loss.mean()
+        if self.reduction == 'sum':
+            loss = loss.sum()
+        return loss
+
+
+if __name__ == '__main__':
+    import torchvision
+    import torch
+    import numpy as np
+    import random
+    torch.manual_seed(15)
+    random.seed(15)
+    np.random.seed(15)
+    torch.backends.cudnn.deterministic = True
+
+    class Model(nn.Module):
+        def __init__(self):
+            super(Model, self).__init__()
+            net = torchvision.models.resnet18(pretrained=False)
+            self.conv1 = net.conv1
+            self.bn1 = net.bn1
+            self.maxpool = net.maxpool
+            self.relu = net.relu
+            self.layer1 = net.layer1
+            self.layer2 = net.layer2
+            self.layer3 = net.layer3
+            self.layer4 = net.layer4
+            self.out = nn.Conv2d(512, 3, 3, 1, 1)
+        def forward(self, x):
+            feat = self.conv1(x)
+            feat = self.bn1(feat)
+            feat = self.relu(feat)
+            feat = self.maxpool(feat)
+            feat = self.layer1(feat)
+            feat = self.layer2(feat)
+            feat = self.layer3(feat)
+            feat = self.layer4(feat)
+            feat = self.out(feat)
+            out = F.interpolate(feat, x.size()[2:], mode='bilinear', align_corners=True)
+            return out
+    net1 = Model()
+    net2 = Model()
+    net2.load_state_dict(net1.state_dict())
+
+    criteria1 = FocalLossV2()
+    # criteria2 = FocalLossV3()
+    net1.cuda()
+    net2.cuda()
+    net1.train()
+    net2.train()
+    net1.double()
+    net2.double()
+    criteria1.cuda()
+    # criteria2.cuda()
+
+    optim1 = torch.optim.SGD(net1.parameters(), lr=1e-2)
+    # optim2 = torch.optim.SGD(net2.parameters(), lr=1e-2)
+
+    bs = 16
+    for it in range(300000):
+        inten = torch.randn(bs, 3, 224, 244).cuda()
+        #  lbs = torch.randint(0, 2, (bs, 3, 224, 244)).float().cuda()
+        lbs = torch.randn(bs, 3, 224, 244).sigmoid().cuda()
+        inten = inten.double()
+        lbs = lbs.double()
+        logits = net1(inten)
+        loss1 = criteria1(logits, lbs)
+        optim1.zero_grad()
+        loss1.backward()
+        optim1.step()
+        # logits = net2(inten)
+        # loss2 = criteria2(logits, lbs)
+        # optim2.zero_grad()
+        # loss2.backward()
+        # optim2.step()
+        # with torch.no_grad():
+        #     if (it+1) % 50 == 0:
+        #         print('iter: {}, ================='.format(it+1))
+        #         print('out.weight: ', torch.mean(torch.abs(net1.out.weight - net2.out.weight)).item())
+        #         print('conv1.weight: ', torch.mean(torch.abs(net1.conv1.weight - net2.conv1.weight)).item())
+        #         print('loss: ', loss1.item() - loss2.item())

modeling/gat.py

@@ -0,0 +1,665 @@
+#
+# Zhoubo
+# 
+"""
+    FLASH: https://arxiv.org/abs/2202.10447
+"""
+import copy
+import torch
+import math
+import os
+from collections import Sequence
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.activations import ACT2FN
+from .ops import sequence_masking, XSoftmax, StableDropout, MaskedLayerNorm
+from .config import ModelConfig
+from .cache_utils import load_model_state
+import einops
+
+
+class ScaleNorm(nn.Module):
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.scala = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        mean_square = (x ** 2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(mean_square + self.eps) * self.scala
+        return x
+
+
+class BertLayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        """Construct a layernorm module in the TF style (epsilon inside the square root).
+        """
+        super(BertLayerNorm, self).__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.bias = nn.Parameter(torch.zeros(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, x):
+        u = x.mean(-1, keepdim=True)
+        s = (x - u).pow(2).mean(-1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+        return self.weight * x + self.bias
+
+
+class ScaledSinuEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(1,))
+        inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x):
+        n, device = x.shape[1], x.device
+        t = torch.arange(n, device = device).type_as(self.inv_freq)
+        sinu = torch.einsum('i , j -> i j', t, self.inv_freq)
+        emb = torch.cat((sinu.sin(), sinu.cos()), dim = -1)
+        return emb * self.scale
+
+
+def RoPE(x, dim):
+    """
+    :param x: input tensor
+    :param dim: oprate dimension
+    :return: tensor
+    """
+    shape = x.shape
+    if isinstance(dim, int):
+        dim = [dim]
+
+    spatial_shape = [shape[i] for i in dim]
+    total_len = 1
+    for i in spatial_shape:
+        total_len *= i
+    position = torch.reshape(torch.arange(total_len, dtype=torch.float, device=x.device), spatial_shape)
+
+    for i in range(dim[-1] + 1, len(shape) - 1, 1):
+        position = torch.unsqueeze(position, dim=-1)
+
+    half_size = shape[-1] // 2
+    freq_seq = -torch.arange(half_size, dtype=torch.float, device=x.device) / float(half_size)
+    inv_freq = 10000 ** -freq_seq
+    sinusoid = torch.einsum("...,d->...d", position, inv_freq)
+    sin = torch.sin(sinusoid)
+    cos = torch.cos(sinusoid)
+    x1, x2 = torch.chunk(x, 2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+def rel_pos_bias(seq_len, s):
+    a = torch.rand([1, s], dtype=torch.float)
+    b = torch.rand([1, s], dtype=torch.float)
+    w = torch.rand([2 * seq_len - 1], dtype=torch.float)
+    if seq_len <= 512:
+        t = F.pad(w[: 2 * seq_len - 1], [0, seq_len]).repeat(seq_len)
+        t = t[..., :-seq_len].reshape(-1, seq_len, 3 * seq_len - 2)
+        r = (2 * seq_len - 1) // 2
+        t = t[..., r:-r]
+    else:
+        a = RoPE(a.repeat(seq_len, 1), dim=[0])
+        b = RoPE(b.repeat(seq_len, 1), dim=[0])
+        t = torch.einsum("mk,nk->mn", a, b)
+    return t
+
+def squared_relu(x, attention_mask, dim=-1):
+    rmask = ~(attention_mask.bool())
+    x = x.masked_fill(rmask, 0)
+    return torch.square(F.relu(x))
+
+
+def attention_normalize(a, axis=-1, mask=None, fn='softmax'):
+    if fn == 'softmax':
+        return XSoftmax.apply(a, mask, axis)
+    else:
+        mask_ = a > -float('inf') / 10
+        # mask_ = mask_.byte()
+        mask_ = torch.sum(mask_, axis=axis, keepdim=True)
+        l = torch.maximum(mask_, torch.ones_like(mask_))
+        if fn == 'relu':
+            rmask = ~(mask.bool())
+            a = a.masked_fill(rmask, 0)
+            return torch.square(F.relu(a)) / l
+        elif fn == 'softmax_plus':
+            return XSoftmax.apply(a * torch.log(l) / np.log(512), mask, axis)
+    return a
+
+
+class GAULinear(nn.Linear):
+    def init_weight(self):
+        nn.init.xavier_uniform_(self.weight)
+
+
+class GatedAttentionUnit(nn.Module):
+    """
+    GAU Block: Gate Attention Unit
+    """
+    def __init__(
+        self,
+        max_seq_length,
+        hidden_size,
+        attention_key_size=128,
+        activation='swish',
+        use_bias=True,
+        attention_norm_type='squared_relu',
+        attention_scale=True,
+        dropout=0.1,
+        pre_norm=False,
+        norm_type="layer_norm", 
+        eps=1e-5,
+        shift_token=False, 
+        use_rel_bias=False,
+        add_residual=True,
+        **kwargs,):
+
+        super(GatedAttentionUnit, self).__init__(**kwargs)
+        self.max_seq_length = max_seq_length
+        self.units = hidden_size
+        self.intermediate_size = self.units * 2
+        self.key_size = attention_key_size
+        self.activation = activation
+        self.use_bias = use_bias
+        self.attention_norm_type = attention_norm_type
+        self.attention_scale = attention_scale
+        self.dropout = StableDropout(dropout)
+        self.i_dense = nn.Sequential(
+            nn.Linear(self.units, 2 * self.intermediate_size + self.key_size, bias=self.use_bias),
+            nn.SiLU()
+        )
+        self.o_dense = nn.Sequential(
+                        nn.Linear(self.intermediate_size, self.units, bias=self.use_bias),
+                        self.dropout)
+        self.q_scaleoffset = OffsetScale(self.key_size)
+        self.k_scaleoffset = OffsetScale(self.key_size)
+        self.pre_norm = pre_norm
+        self.norm = (nn.LayerNorm(hidden_size, eps=eps) if norm_type.lower() == "layer_norm" else ScaleNorm(eps=eps))
+        self.add_residual = add_residual
+
+    def forward(self, x, attention_mask=None, **kwargs):
+        shortcut = x
+        
+        if self.pre_norm:
+            x = self.norm(x)
+
+        x = self.i_dense(x)
+        u, v, qk = torch.split(x, [self.intermediate_size, self.intermediate_size, self.key_size], dim=-1)
+        q, k = self.q_scaleoffset(qk), self.k_scaleoffset(qk)
+        qk = RoPE(torch.stack([q, k], 2), dim=1)
+        q, k = qk[:, :, 0], qk[:, :, 1]
+        a = torch.einsum('bmd,bnd->bmn', q, k)
+        if self.attention_scale:
+            a = a / self.key_size**0.5
+        a = sequence_masking(a, attention_mask, '-inf', -1)
+        A = attention_normalize(a, -1, fn=self.attention_norm_type)
+        if self.dropout:
+            A = self.dropout(A)
+        out = self.o_dense(u * torch.einsum('bmn,bnd->bmd', A, v))
+
+        if self.add_residual:
+            out = out + shortcut
+        if not self.pre_norm:
+            out = self.norm(out)
+        return out
+
+
+class OffsetScale(nn.Module):
+    def __init__(self, dim, heads = 1):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+        self.beta = nn.Parameter(torch.zeros(heads, dim))
+        # nn.init.normal_(self.gamma, std = 0.02)
+        nn.init.xavier_uniform_(self.gamma)
+
+    def forward(self, x):
+        out = torch.einsum('... d, h d -> ... h d', x, self.gamma) + self.beta
+        return out
+
+
+class BertIntermediate(nn.Module):
+  def __init__(self, config):
+    super().__init__()
+    self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+    self.intermediate_act_fn = ACT2FN[config.hidden_act] \
+      if isinstance(config.hidden_act, str) else config.hidden_act
+
+  def forward(self, hidden_states):
+    hidden_states = self.dense(hidden_states)
+    hidden_states = self.intermediate_act_fn(hidden_states)
+    return hidden_states
+
+
+class BertOutput(nn.Module):
+  def __init__(self, config):
+    super(BertOutput, self).__init__()
+    self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+    self.LayerNorm = nn.LayerNorm(config.hidden_size, config.layer_norm_eps)
+    self.dropout = StableDropout(config.hidden_dropout_prob)
+    self.config = config
+
+  def forward(self, hidden_states, input_states, mask=None):
+    hidden_states = self.dense(hidden_states)
+    hidden_states = self.dropout(hidden_states)
+    hidden_states += input_states
+    hidden_states = MaskedLayerNorm(self.LayerNorm, hidden_states)
+    return hidden_states
+
+
+class GAU(nn.Module):
+    def __init__(self, max_seq_length, hidden_size, expansion_factor=2, s=128, norm_type="layer_norm", eps=1e-5,
+                 hidden_act="silu", shift_token=False, use_rel_bias=False, attention_norm_type='softmax',
+                 pre_norm=False, dropout=0, add_residual = True):
+        super(GAU, self).__init__()
+        self.max_seq_length = max_seq_length
+        self.shift_token = shift_token
+        hidden_dim = int(expansion_factor * hidden_size)
+        self.norm = (nn.LayerNorm(hidden_size, eps=eps) if norm_type == "layer_norm" else ScaleNorm(eps=eps))
+        self.use_rel_bias = use_rel_bias
+        self.attention_norm_type = attention_norm_type
+        # if attention_norm_type == 'relu':
+        #     self.attention_norm_func = squared_relu
+        # else:
+        #     self.attention_norm_func = XSoftmax.apply
+        # self.norm = norm_klass(hidden_size)
+
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_hidden = nn.Sequential(
+            nn.Linear(hidden_size, hidden_dim * 2),
+            nn.SiLU()
+        )
+
+        self.to_qk = nn.Sequential(
+            nn.Linear(hidden_size, s),
+            nn.SiLU()
+        )
+
+        self.offsetscale = OffsetScale(s, heads = 2)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_size),
+            nn.Dropout(dropout)
+        )
+
+        self.add_residual = add_residual
+        self.act_fn = ACT2FN[hidden_act]
+        self.pre_norm = pre_norm
+
+
+    def forward(
+        self,
+        x,
+        relative_pos = None,
+        attention_mask = None
+    ):
+        seq_len, device = x.shape[-2], x.device
+        if self.pre_norm:
+            normed_x = self.norm(x)
+        else:
+            normed_x = x
+        v, gate = self.to_hidden(normed_x).chunk(2, dim = -1)
+
+        qk = self.to_qk(normed_x)
+        base = self.offsetscale(qk)
+        base = RoPE(base, 1).half()
+        q, k = base.unbind(dim = -2)
+        sim = torch.einsum('b i d, b j d -> b i j', q, k)
+
+        if relative_pos is not None:
+            sim = sim + relative_pos
+        if attention_mask is not None:
+            if attention_mask.dim() < 3:
+                attention_mask = einops.rearrange(attention_mask, 'b j -> b 1 j')
+            # attn = attn.masked_fill(~attention_mask.bool(), 0.)
+        attn = attention_normalize(sim, mask=attention_mask, fn=self.attention_norm_type)
+        # attn = F.relu(sim) ** 2 / seq_len# / q.size(-1)
+        # logger.info(attn.max())
+        attn = self.dropout(attn)
+        # if self.causal:
+        #     causal_mask = torch.ones((seq_len, seq_len), dtype = torch.bool, device = device).triu(1)
+        #     attn = attn.masked_fill(causal_mask, 0.)
+
+        out = torch.einsum('b i j, b j d -> b i d', attn.half(), v)
+        out = out * gate
+
+        out = self.to_out(out)
+
+        if self.add_residual:
+            out = out + x
+        if not self.pre_norm:
+            out = self.norm(out)
+        return out
+
+
+class GatLayer(nn.Module):
+    def __init__(self, config, shift_token=False, use_ffn=False):
+        super(GatLayer, self).__init__()
+        self.attention = GatedAttentionUnit(config.max_position_embeddings, config.hidden_size, 
+                                    shift_token=shift_token, use_rel_bias=config.use_rel_bias, 
+                                    norm_type=config.norm_type, attention_norm_type=config.attention_norm_type,
+                                    pre_norm=config.pre_norm, dropout=config.hidden_dropout_prob)
+        if use_ffn:
+            self.intermediate = BertIntermediate(config)
+            self.output = BertOutput(config)
+        self.use_ffn = use_ffn
+
+    def forward(self, hidden_states, attention_mask, return_att=False, query_states=None, relative_pos=None, rel_embeddings=None):
+        attention_output = self.attention(hidden_states, attention_mask=attention_mask, relative_pos=relative_pos)
+        if self.use_ffn:
+            intermediate_output = self.intermediate(attention_output)
+            layer_output = self.output(intermediate_output, attention_output)
+            return layer_output
+        else:
+            return attention_output
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(
+        self,
+        scale,
+        causal = False,
+        num_buckets = 32,
+        max_distance = 128
+    ):
+        super().__init__()
+        self.scale = scale
+        self.causal = causal
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = nn.Embedding(num_buckets, 1)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position,
+        causal = True,
+        num_buckets = 32,
+        max_distance = 128
+    ):
+        ret = 0
+        n = -relative_position
+        if not causal:
+            num_buckets //= 2
+            ret += (n < 0).long() * num_buckets
+            n = torch.abs(n)
+        else:
+            n = torch.max(n, torch.zeros_like(n))
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = max_exact + (
+            torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)
+        ).long()
+        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, x):
+        i, j, device = *x.shape[-2:], x.device
+        q_pos = torch.arange(i, dtype = torch.long, device = device)
+        k_pos = torch.arange(j, dtype = torch.long, device = device)
+        rel_pos = einops.rearrange(k_pos, 'j -> 1 j') - einops.rearrange(q_pos, 'i -> i 1')
+        rp_bucket = self._relative_position_bucket(rel_pos, causal = self.causal, num_buckets = self.num_buckets, max_distance = self.max_distance)
+        values = self.relative_attention_bias(rp_bucket)
+        bias = einops.rearrange(values, 'i j 1 -> i j')
+        return bias * self.scale
+
+
+class GatEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, config, with_position=False):
+        super(GatEmbeddings, self).__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(
+            config.type_vocab_size, config.hidden_size)
+        self.with_position = with_position
+        if with_position:
+            self.position_embeddings = ScaledSinuEmbedding(config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=1e-5)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+
+    def forward(self, input_ids, token_type_ids=None, position_ids=None, token_mask=None):
+        seq_length = input_ids.size(1)
+        if position_ids is None:
+            position_ids = torch.arange(
+                seq_length, dtype=torch.long, device=input_ids.device)
+            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+
+        words_embeddings = self.word_embeddings(input_ids)
+        if self.with_position:
+            position_embeddings = self.position_embeddings(words_embeddings)
+        else:
+            position_embeddings = 0
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        # if self.num_pos_emb > 1:
+        #     num_batch = position_embeddings.size(0)
+        #     num_pos = position_embeddings.size(1)
+        #     position_embeddings = position_embeddings.view(
+        #         num_batch, num_pos, self.num_pos_emb, -1)[torch.arange(0, num_batch).long(), :, task_idx, :]
+
+        embeddings = words_embeddings + position_embeddings + token_type_embeddings
+        # if self.fp32_embedding:
+        #     embeddings = embeddings.half()
+        embeddings = MaskedLayerNorm(self.LayerNorm, embeddings, token_mask)
+        embeddings = self.dropout(embeddings)
+        return {
+                'embeddings': embeddings,
+                'position_embeddings': position_embeddings}
+
+
+class GatEncoder(nn.Module):
+    def __init__(self, config, shift_token=False):
+        super().__init__()
+        layer = GatLayer(config, shift_token=shift_token)
+        self.layer = nn.ModuleList([copy.deepcopy(layer)
+                                    for _ in range(config.num_hidden_layers)])
+    
+    def get_attention_mask(self, attention_mask):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1)
+            attention_mask = extended_attention_mask*extended_attention_mask.squeeze(-2).unsqueeze(-1)
+            attention_mask = attention_mask.byte()
+        return attention_mask
+
+    def forward(self, hidden_states, attention_mask, output_all_encoded_layers=True, return_att=False, query_states = None, relative_pos=None):
+        all_encoder_layers = []
+        att_matrices = []
+        if isinstance(hidden_states, Sequence):
+            next_kv = hidden_states[0]
+        else:
+            next_kv = hidden_states
+        # rel_embeddings = self.get_rel_embedding()
+        for i, layer_module in enumerate(self.layer):
+            output_states = layer_module(next_kv, attention_mask, query_states = query_states, relative_pos=relative_pos)
+            if return_att:
+                    output_states, att_m = output_states
+
+            # if i == 0 and self.with_conv:
+            #     prenorm = output_states #output['prenorm_states']
+            #     output_states = self.conv(hidden_states, prenorm, input_mask)
+
+            if query_states is not None:
+                query_states = output_states
+                if isinstance(hidden_states, Sequence):
+                    next_kv = hidden_states[i+1] if i+1 < len(self.layer) else None
+            else:
+                next_kv = output_states
+
+            if output_all_encoded_layers:
+                all_encoder_layers.append(output_states)
+                if return_att:
+                    att_matrices.append(att_m)
+        if not output_all_encoded_layers:
+            all_encoder_layers.append(output_states)
+            if return_att:
+                    att_matrices.append(att_m)
+        return {
+            'hidden_states': all_encoder_layers,
+            'attention_matrices': att_matrices
+            }
+
+
+class GatModel(torch.nn.Module):
+    """ 
+    Parameters:
+        config:
+        A model config class instance with the configuration to build a new model. The schema is similar to `BertConfig`, 
+
+        pre_trained:
+        The pre-trained DeBERTa model, it can be a physical path of a pre-trained DeBERTa model or a released configurations,
+            i.e. [**base, large, base_mnli, large_mnli**]
+
+    """
+
+    def __init__(self, config=None, pre_trained=None, pooler=False, shift_token=False, causal=False, **kwargs):
+        super().__init__()
+        state = None
+        if pre_trained is not None:
+            state, model_config = load_model_state(pre_trained)
+            if config is not None and model_config is not None:
+                for k in config.__dict__:
+                    if k not in ['hidden_size',
+                            'intermediate_size',
+                            'num_attention_heads',
+                            'num_hidden_layers',
+                            'vocab_size',
+                            'max_position_embeddings']:
+                            model_config.__dict__[k] = config.__dict__[k]
+            config = copy.copy(model_config)
+        self.embeddings = GatEmbeddings(config, with_position=True)
+        self.encoder = GatEncoder(config, shift_token=shift_token)
+        if not pooler:
+            self.pooler = None
+        self.config = config
+        self.pre_trained = pre_trained
+        self.apply_state(state)
+
+    def get_attention_mask(self, input_ids=None, token_type_ids=None, attention_mask=None, input_mask=None):
+        if attention_mask is None:
+            if input_mask is not None:
+                return input_mask.unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), input_mask.size(1))
+            else:
+                return torch.ones_like(input_ids, dtype=torch.uint8).unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), input_mask.size(1))
+        else:
+            if attention_mask.dim() == 2:
+                if input_mask is not None:
+                    attention_mask = attention_mask * input_mask
+                return attention_mask.unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), attention_mask.size(-1))
+            if attention_mask.dim() == 4:
+                attention_mask = attention_mask.squeeze(2)
+            if attention_mask.dim() == 3:
+                if input_mask is not None:
+                    return attention_mask * input_mask.unsqueeze(-1).expand(input_mask.size(0), input_mask.size(1), attention_mask.size(-1))
+                else:
+                    return attention_mask
+            
+
+    def forward(self, input_ids, input_mask, attention_mask=None, token_type_ids=None, 
+                output_all_encoded_layers=True, position_ids=None, return_att=False):
+        """
+        Args:
+            input_ids:
+                a torch.LongTensor of shape [batch_size, sequence_length] \
+            with the word token indices in the vocabulary
+
+            attention_mask:
+                an optional parameter for input mask or attention mask.
+
+                - If it's an input mask, then it will be torch.LongTensor of shape [batch_size, sequence_length] with indices \
+            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max \
+            input sequence length in the current batch. It's the mask that we typically use for attention when \
+            a batch has varying length sentences.
+
+                - If it's an attention mask then it will be torch.LongTensor of shape [batch_size, sequence_length, sequence_length]. \
+            In this case, it's a mask indicate which tokens in the sequence should be attended by other tokens in the sequence.
+
+            token_type_ids:
+                an optional torch.LongTensor of shape [batch_size, sequence_length] with the token \
+            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to \
+            a `sentence B` token (see BERT paper for more details).
+
+            output_all_encoded_layers:
+                whether to output results of all encoder layers, default, True
+
+            Returns:
+
+            - The output of the stacked transformer layers if `output_all_encoded_layers=True`, else \
+            the last layer of stacked transformer layers
+
+            - Attention matrix of self-attention layers if `return_att=True`
+
+
+            Example::
+
+            # Batch of wordPiece token ids.
+            # Each sample was padded with zero to the maxium length of the batch
+            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
+            # Mask of valid input ids
+            attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
+
+            # DeBERTa model initialized with pretrained base model
+            bert = DeBERTa(pre_trained='base')
+
+            encoder_layers = bert(input_ids, attention_mask=attention_mask)
+
+            """
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+            # input_mask = torch.ones_like(input_ids)
+
+        if input_mask is None:
+            idxs = torch.flip(torch.cumsum(torch.flip(token_type_ids, [-1]), axis=1), [-1])
+            input_mask = idxs > 0
+            if not torch.any(input_mask):
+                input_mask = torch.ones_like(input_ids)
+            input_mask = input_mask.byte()
+        attention_mask = self.get_attention_mask(input_ids, token_type_ids, attention_mask, input_mask)
+        attention_mask = attention_mask.byte()
+        embedding_output = self.embeddings(input_ids.to(torch.long), token_type_ids.to(torch.long), position_ids, input_mask)
+        encoder_output = self.encoder(embedding_output['embeddings'], attention_mask, output_all_encoded_layers=output_all_encoded_layers, return_att = return_att)
+        encoder_output.update(embedding_output)
+        return encoder_output
+
+    def apply_state(self, state = None):
+        """ Load state from previous loaded model state dictionary.
+
+        Args:
+            state (:obj:`dict`, optional): State dictionary as the state returned by torch.module.state_dict(), default: `None`. \
+                If it's `None`, then will use the pre-trained state loaded via the constructor to re-initialize \
+                the `DeBERTa` model
+        """
+        if self.pre_trained is None and state is None:
+            return
+        if state is None:
+            state, config = load_model_state(self.pre_trained)
+            self.config = config
+        
+        prefix = ''
+        for k in state:
+            if 'embeddings.' in k:
+                    if not k.startswith('embeddings.'):
+                        prefix = k[:k.index('embeddings.')]
+                    break
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+        self._load_from_state_dict(state, prefix = prefix, local_metadata=None, strict=True, missing_keys=missing_keys, unexpected_keys=unexpected_keys, error_msgs=error_msgs)
+
+
+if __name__ == '__main__':
+    model = GatModel(768, 64)

modeling/gpt2_bpe_utils.py

@@ -0,0 +1,163 @@
+"""
+Byte pair encoding utilities from GPT-2.
+
+Original source: https://github.com/openai/gpt-2/blob/master/src/encoder.py
+Original license: MIT
+"""
+
+from functools import lru_cache
+import json
+import random
+import unicodedata
+
+try:
+    import regex as re
+except ImportError:
+    raise ImportError('Please install regex with: pip install regex')
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+class Encoder:
+
+    def __init__(self, encoder, bpe_merges, errors='replace'):
+        self.encoder = encoder
+        self.decoder = {v:k for k,v in self.encoder.items()}
+        self.errors = errors # how to handle errors in decoding
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v:k for k, v in self.byte_encoder.items()}
+        self.bpe_ranks = dict(zip([tuple(k) for k in bpe_merges], range(len(bpe_merges))))
+        self.cache = {}
+        self.random = random.Random(0)
+
+        # Should haved added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
+        self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+    
+    def split_to_words(self, text):
+      return list(re.findall(self.pat, text))
+
+    def encode(self, text):
+        bpe_tokens = []
+        for token in self.split_to_words(text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors=self.errors)
+        return text
+
+def get_encoder(encoder, vocab):
+    return Encoder(
+        encoder=encoder,
+        bpe_merges=vocab,
+    )
+
+def _is_whitespace(char):
+    """Checks whether `chars` is a whitespace character."""
+    # \t, \n, and \r are technically contorl characters but we treat them
+    # as whitespace since they are generally considered as such.
+    if char == " " or char == "\t" or char == "\n" or char == "\r":
+        return True
+    cat = unicodedata.category(char)
+    if cat == "Zs":
+        return True
+    return False
+
+def _is_control(char):
+    """Checks whether `chars` is a control character."""
+    # These are technically control characters but we count them as whitespace
+    # characters.
+    if char == "\t" or char == "\n" or char == "\r":
+        return False
+    cat = unicodedata.category(char)
+    if cat.startswith("C"):
+        return True
+    return False
+
+def _is_punctuation(char):
+    """Checks whether `chars` is a punctuation character."""
+    cp = ord(char)
+    # We treat all non-letter/number ASCII as punctuation.
+    # Characters such as "^", "$", and "`" are not in the Unicode
+    # Punctuation class but we treat them as punctuation anyways, for
+    # consistency.
+    if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
+            (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
+        return True
+    cat = unicodedata.category(char)
+    if cat.startswith("P"):
+        return True
+    return False

modeling/gpt2_tokenizer.py

@@ -0,0 +1,216 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 01/15/2020
+#
+
+# This piece of code is derived from https://github.com/pytorch/fairseq/blob/master/fairseq/data/encoders/gpt2_bpe.py
+
+import torch
+import unicodedata
+import os
+from .gpt2_bpe_utils import get_encoder,_is_control,_is_whitespace,_is_punctuation
+from .cache_utils import load_vocab
+
+__all__ = ['GPT2Tokenizer']
+
+class GPT2Tokenizer(object):
+  """ A wrapper of GPT2 tokenizer with similar interface as BERT tokenizer
+
+  Args:
+    
+    vocab_file (:obj:`str`, optional):
+      The local path of vocabulary package or the release name of vocabulary in `DeBERTa GitHub releases <https://github.com/microsoft/DeBERTa/releases>`_, \
+          e.g. "bpe_encoder", default: `None`. 
+          
+          If it's `None`, then it will download the vocabulary in the latest release from GitHub. The vocabulary file is a \
+          state dictionary with three items, "dict_map", "vocab", "encoder" which correspond to three files used in `RoBERTa`, i.e. `dict.txt`, `vocab.txt` and `encoder.json`. \
+          
+          The difference between our wrapped GPT2 tokenizer and RoBERTa wrapped tokenizer are,
+
+          - Special tokens, unlike `RoBERTa` which use `<s>`, `</s>` as the `start` token and `end` token of a sentence. We use `[CLS]` and `[SEP]` as the `start` and `end`\
+              token of input sentence which is the same as `BERT`.
+
+          - We remapped the token ids in our dictionary with regarding to the new special tokens, `[PAD]` => 0, `[CLS]` => 1, `[SEP]` => 2, `[UNK]` => 3, `[MASK]` => 50264
+
+    do_lower_case (:obj:`bool`, optional):
+      Whether to convert inputs to lower case. **Not used in GPT2 tokenizer**.
+
+    special_tokens (:obj:`list`, optional):
+      List of special tokens to be added to the end of the vocabulary.
+
+
+  """
+  def __init__(self, vocab_file=None, do_lower_case=True, special_tokens=None):
+    self.pad_token='[PAD]'
+    self.sep_token='[SEP]'
+    self.unk_token='[UNK]'
+    self.cls_token='[CLS]'
+
+    self.symbols = []
+    self.count = []
+    self.indices = {}
+    self.pad_token_id = self.add_symbol(self.pad_token)
+    self.cls_token_id = self.add_symbol(self.cls_token)
+    self.sep_token_id = self.add_symbol(self.sep_token)
+    self.unk_token_id = self.add_symbol(self.unk_token)
+
+    self.gpt2_encoder = torch.load(vocab_file)
+    self.bpe = get_encoder(self.gpt2_encoder['encoder'], self.gpt2_encoder['vocab'])
+    for w,n in self.gpt2_encoder['dict_map']:
+      self.add_symbol(w, n)
+
+    self.mask_token='[MASK]'
+    self.mask_id = self.add_symbol(self.mask_token)
+    self.special_tokens = ['[MASK]', '[SEP]', '[PAD]', '[UNK]', '[CLS]']
+    if special_tokens is not None:
+      for t in special_tokens:
+        self.add_special_token(t)
+
+    self.vocab = self.indices
+    self.ids_to_tokens = self.symbols
+
+  def tokenize(self, text):
+    """ Convert an input text to tokens.
+      
+      Args:
+        
+        text (:obj:`str`): input text to be tokenized.
+
+      Returns:
+        A list of byte tokens where each token represent the byte id in GPT2 byte dictionary
+
+      Example::
+        
+        >>> tokenizer = GPT2Tokenizer()
+        >>> text = "Hello world!"
+        >>> tokens = tokenizer.tokenize(text)
+        >>> print(tokens)
+        ['15496', '995', '0']
+        
+    """
+    bpe = self._encode(text)
+
+    return [t for t in bpe.split(' ') if t]
+
+  def convert_tokens_to_ids(self, tokens):
+    """ Convert list of tokens to ids.
+      
+      Args:
+
+        tokens (:obj:`list<str>`): list of tokens
+
+      Returns:
+        
+        List of ids
+    """
+
+    return [self.vocab[t] for t in tokens]
+
+  def convert_ids_to_tokens(self, ids):
+    """ Convert list of ids to tokens.
+      
+      Args:
+
+        ids (:obj:`list<int>`): list of ids
+
+      Returns:
+        
+        List of tokens
+    """
+
+    tokens = []
+    for i in ids:
+      tokens.append(self.ids_to_tokens[i])
+    return tokens
+
+  def split_to_words(self, text):
+    return self.bpe.split_to_words(text)
+
+  def decode(self, tokens):
+    """ Decode list of tokens to text strings.
+    
+      Args:
+        
+        tokens (:obj:`list<str>`): list of tokens.
+
+      Returns:
+        
+        Text string corresponds to the input tokens.
+
+      Example::
+        
+        >>> tokenizer = GPT2Tokenizer()
+        >>> text = "Hello world!"
+        >>> tokens = tokenizer.tokenize(text)
+        >>> print(tokens)
+        ['15496', '995', '0']
+        
+        >>> tokenizer.decode(tokens)
+        'Hello world!'
+      
+    """
+    return self.bpe.decode([int(t) for t in tokens if t not in self.special_tokens])
+
+  def add_special_token(self, token):
+    """Adds a special token to the dictionary.
+    
+      Args:
+        token (:obj:`str`): Tthe new token/word to be added to the vocabulary.
+
+      Returns:
+        The id of new token in the vocabulary.
+
+    """
+    self.special_tokens.append(token)
+    return self.add_symbol(token)
+
+  def part_of_whole_word(self, token, is_bos=False):
+    if is_bos:
+      return True
+    s = self._decode(token)
+    if (len(s)==1 and (_is_whitespace(list(s)[0]) or _is_control(list(s)[0]) or _is_punctuation(list(s)[0]))):
+      return False
+
+    return not s.startswith(' ')
+
+  def sym(self, id):
+    return self.ids_to_tokens[id]
+
+  def id(self, sym):
+    return self.vocab[sym]
+
+  def _encode(self, x: str) -> str:
+    return ' '.join(map(str, self.bpe.encode(x)))
+
+  def _decode(self, x: str) -> str:
+    return self.bpe.decode(map(int, x.split()))
+
+  def add_symbol(self, word, n=1):
+    """Adds a word to the dictionary.
+    
+      Args:
+        word (:obj:`str`): Tthe new token/word to be added to the vocabulary.
+        n (int, optional): The frequency of the word.
+
+      Returns:
+        The id of the new word.
+
+    """
+    if word in self.indices:
+      idx = self.indices[word]
+      self.count[idx] = self.count[idx] + n
+      return idx
+    else:
+      idx = len(self.symbols)
+      self.indices[word] = idx
+      self.symbols.append(word)
+      self.count.append(n)
+      return idx
+
+  def save_pretrained(self, path: str):
+    torch.save(self.gpt2_encoder, path)

modeling/mlm.py

@@ -0,0 +1,38 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# This piece of code is modified based on https://github.com/huggingface/transformers
+
+import torch
+from torch import nn
+import pdb
+
+from .bert import LayerNorm,ACT2FN
+
+__all__ = ['MLMPredictionHead']
+
+class MLMPredictionHead(nn.Module):
+    def __init__(self, config, vocab_size):
+        super().__init__()
+        self.embedding_size = getattr(config, 'embedding_size', config.hidden_size)
+        self.dense = nn.Linear(config.hidden_size, self.embedding_size)
+        self.transform_act_fn = ACT2FN[config.hidden_act] \
+            if isinstance(config.hidden_act, str) else config.hidden_act
+
+        self.LayerNorm = LayerNorm(self.embedding_size, config.layer_norm_eps)
+        self.bias = nn.Parameter(torch.zeros(vocab_size))
+        self.pre_norm = PreLayerNorm(config)
+
+    def forward(self, hidden_states, embeding_weight):
+        hidden_states = self.pre_norm(hidden_states)
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        # b x s x d
+        hidden_states = MaskedLayerNorm(self.LayerNorm, hidden_states)
+
+        # b x s x v
+        logits = torch.matmul(hidden_states, embeding_weight.t().to(hidden_states)) + self.bias
+        return logits

modeling/nnmodule.py

@@ -0,0 +1,184 @@
+import pdb
+import os
+import torch
+import copy
+from torch import nn, tensor
+from .config import ModelConfig
+from ..utils import xtqdm as tqdm
+from .cache_utils import load_model_state
+from .flash import GAULinear
+
+from ..utils import get_logger
+logger = get_logger()
+
+__all__ = ['NNModule']
+
+def truncated_normal_(shape, mean=0, std=0.09):
+    with torch.no_grad():
+        tensor = torch.zeros(shape)
+        tmp = tensor.new_empty(shape + (4,)).normal_()
+        valid = (tmp < 2) & (tmp > -2)
+        ind = valid.max(-1, keepdim=True)[1]
+        tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
+        tensor.data.mul_(std).add_(mean)
+        return tensor
+
+class NNModule(nn.Module):
+  """ An abstract class to handle weights initialization and \
+    a simple interface for dowloading and loading pretrained models.
+
+  Args:
+    
+    config (:obj:`~DeBERTa.deberta.ModelConfig`): The model config to the module
+
+  """
+
+  def __init__(self, config, *inputs, **kwargs):
+    super().__init__()
+    self.config = config
+
+  def init_weights(self, module):
+    """ Apply Gaussian(mean=0, std=`config.initializer_range`) initialization to the module.
+
+    Args:
+      
+      module (:obj:`torch.nn.Module`): The module to apply the initialization.
+    
+    Example::
+      
+      class MyModule(NNModule):
+        def __init__(self, config):
+          # Add construction instructions
+          self.bert = DeBERTa(config)
+          
+          # Add other modules
+          ...
+
+          # Apply initialization
+          self.apply(self.init_weights)
+
+    """
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+      module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+    if isinstance(module, nn.Linear) and module.bias is not None:
+      module.bias.data.zero_()
+
+  def init_weights_gau(self, module):
+    """ Apply Gaussian(mean=0, std=`config.initializer_range`) initialization to the module.
+
+    Args:
+      
+      module (:obj:`torch.nn.Module`): The module to apply the initialization.
+    
+    Example::
+      
+      class MyModule(NNModule):
+        def __init__(self, config):
+          # Add construction instructions
+          self.bert = DeBERTa(config)
+          
+          # Add other modules
+          ...
+
+          # Apply initialization
+          self.apply(self.init_weights)
+
+    """
+    if isinstance(module, GAULinear):
+      module.init_weight()
+    else:
+      if isinstance(module, (nn.Linear, nn.Embedding)):
+        # module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        module.weight.data.copy_(self.initializer(module.weight.data.shape))
+    if isinstance(module, nn.Linear) and module.bias is not None:
+      module.bias.data.zero_()
+
+  def initializer(self, shape, dtype=None, order=3, gain=1.0):
+    if shape[1] > 10000 or shape[1] < 10:
+      hidden_size = shape[0]
+    else:
+      hidden_size = shape[1]
+    gain *= self.config.num_hidden_layers ** (-1.0 / order)
+    stddev = 1.13684723 / hidden_size**0.5 * gain
+    return torch.nn.init.trunc_normal_(torch.empty(shape, dtype=dtype), std=stddev)# truncated_normal_(shape, std=stddev)
+
+  @classmethod
+  def load_model(cls, model_path, model_config=None, tag=None, no_cache=False, cache_dir=None , *inputs, **kwargs):
+    """ Instantiate a sub-class of NNModule from a pre-trained model file.
+      
+    Args:
+
+      model_path (:obj:`str`): Path or name of the pre-trained model which can be either,
+        
+        - The path of pre-trained model
+
+        - The pre-trained DeBERTa model name in `DeBERTa GitHub releases <https://github.com/microsoft/DeBERTa/releases>`_, i.e. [**base, base_mnli, large, large_mnli**].
+
+        If `model_path` is `None` or `-`, then the method will create a new sub-class without initialing from pre-trained models.
+
+      model_config (:obj:`str`): The path of model config file. If it's `None`, then the method will try to find the the config in order:
+        
+        1. ['config'] in the model state dictionary.
+
+        2. `model_config.json` aside the `model_path`.
+        
+        If it failed to find a config the method will fail.
+
+      tag (:obj:`str`, optional): The release tag of DeBERTa, default: `None`.
+
+      no_cache (:obj:`bool`, optional): Disable local cache of downloaded models, default: `False`.
+
+      cache_dir (:obj:`str`, optional): The cache directory used to save the downloaded models, default: `None`. If it's `None`, then the models will be saved at `$HOME/.~DeBERTa`
+
+    Return:
+      
+      :obj:`NNModule` : The sub-class object.
+
+    """
+    # Load config
+    if model_config:
+      config = ModelConfig.from_json_file(model_config)
+    else:
+      config = None
+    model_config = None
+    model_state = None
+    if (model_path is not None) and (model_path.strip() == '-' or model_path.strip()==''):
+      model_path = None
+    try:
+      model_state, model_config = load_model_state(model_path, tag=tag, no_cache=no_cache, cache_dir=cache_dir)
+    except Exception as exp:
+      raise Exception(f'Failed to get model {model_path}. Exception: {exp}')
+    
+    if config is not None and model_config is not None:
+      for k in config.__dict__:
+        if k not in ['hidden_size',
+          'intermediate_size',
+          'num_attention_heads',
+          'num_hidden_layers',
+          'vocab_size',
+          'max_position_embeddings'] or (k not in  model_config.__dict__) or (model_config.__dict__[k] < 0):
+          model_config.__dict__[k] = config.__dict__[k]
+    if model_config is not None:
+      config = copy.copy(model_config)
+    vocab_size = config.vocab_size
+    # Instantiate model.
+    model = cls(config, *inputs, **kwargs)
+    if not model_state:
+      return model
+    # copy state_dict so _load_from_state_dict can modify it
+    state_dict = model_state.copy()
+
+    missing_keys = []
+    unexpected_keys = []
+    error_msgs = []
+    metadata = getattr(state_dict, '_metadata', None)
+    def load(module, prefix=''):
+      local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+      module._load_from_state_dict(
+        state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
+      for name, child in module._modules.items():
+        if child is not None:
+          load(child, prefix + name + '.')
+    load(model)
+    logger.warning(f'Missing keys: {missing_keys}, unexpected_keys: {unexpected_keys}, error_msgs: {error_msgs}')
+    return model

modeling/ops.py

@@ -0,0 +1,274 @@
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 01/15/2020
+#
+
+import pdb
+import math
+from packaging import version
+import torch
+from torch.nn import LayerNorm
+from wywLM.utils.jit_tracing import traceable
+
+if version.Version(torch.__version__) >= version.Version('1.0.0'):
+  from torch import _softmax_backward_data as _softmax_backward_data
+else:
+  from torch import softmax_backward_data as _softmax_backward_data
+
+__all__ = ['StableDropout', 'MaskedLayerNorm', 'XSoftmax', 'ACT2FN', 'LayerNorm']
+
+@traceable
+class XSoftmax(torch.autograd.Function):
+  """ Masked Softmax which is optimized for saving memory
+
+  Args:
+      
+    input (:obj:`torch.tensor`): The input tensor that will apply softmax.
+    mask (:obj:`torch.IntTensor`): The mask matrix where 0 indicate that element will be ignored in the softmax caculation.
+    dim (int): The dimenssion that will apply softmax.
+    
+  Example::
+
+    import torch
+    from DeBERTa.deberta import XSoftmax
+    # Make a tensor
+    x = torch.randn([4,20,100])
+    # Create a mask
+    mask = (x>0).int()
+    y = XSoftmax.apply(x, mask, dim=-1)
+      
+  """
+
+  @staticmethod
+  def forward(self, input, mask, dim):
+    """
+    """
+
+    self.dim = dim
+    if mask is None:
+      mask = torch.ones_like(input)
+    if version.Version(torch.__version__) >= version.Version('1.2.0a'):
+      rmask = ~(mask.bool())
+    else:
+      rmask = (1-mask).byte() # This line is not supported by Onnx tracing.
+
+    output = input.masked_fill(rmask, torch.finfo(input.dtype).min) # float('-inf')
+    output = torch.softmax(output, self.dim)
+    output.masked_fill_(rmask, 0)
+    self.save_for_backward(output)
+    return output
+
+  @staticmethod
+  def backward(self, grad_output):
+    """
+    """
+
+    output, = self.saved_tensors
+    if '1.11' in torch.__version__:
+      inputGrad = _softmax_backward_data(grad_output, output, self.dim, output.dtype)
+    else:
+      inputGrad = _softmax_backward_data(grad_output, output, self.dim, output.dtype)
+    return inputGrad, None, None
+
+  @staticmethod
+  def symbolic(g, self, mask, dim):
+      import torch.onnx.symbolic_helper as sym_help
+      from torch.onnx.symbolic_opset9 import masked_fill, softmax
+
+      mask_cast_value = g.op("Cast", mask, to_i=sym_help.cast_pytorch_to_onnx['Long'])
+      r_mask = g.op("Cast", g.op("Sub", g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)), mask_cast_value), to_i=sym_help.cast_pytorch_to_onnx['Byte'])
+      output = masked_fill(g, self, r_mask, g.op("Constant", value_t=torch.tensor(float('-inf'))))
+      output = softmax(g, output, dim)
+      return masked_fill(g, output, r_mask, g.op("Constant", value_t=torch.tensor(0, dtype=torch.uint8)))
+
+class DropoutContext(object):
+  def __init__(self):
+    self.dropout = 0
+    self.mask = None
+    self.scale = 1
+    self.reuse_mask = True
+
+def get_mask(input, local_context):
+  if not isinstance(local_context, DropoutContext):
+    dropout = local_context
+    mask = None
+  else:
+    dropout = local_context.dropout
+    dropout *= local_context.scale
+    mask = local_context.mask if local_context.reuse_mask else None
+
+  if dropout>0 and mask is None:
+    if version.Version(torch.__version__) >= version.Version('1.2.0a'):
+      mask=(1-torch.empty_like(input).bernoulli_(1-dropout)).bool()
+    else:
+      mask=(1-torch.empty_like(input).bernoulli_(1-dropout)).byte()
+  
+  if isinstance(local_context, DropoutContext):
+    if local_context.mask is None:
+      local_context.mask = mask
+
+  return mask, dropout
+
+@traceable
+class XDropout(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, input, local_ctx):
+    mask, dropout = get_mask(input, local_ctx)
+    ctx.scale=1.0/(1-dropout)
+    if dropout>0:
+      ctx.save_for_backward(mask)
+      return input.masked_fill(mask, 0)*ctx.scale
+    else:
+      return input
+
+  @staticmethod
+  def backward(ctx, grad_output):
+    if ctx.scale > 1:
+      mask, = ctx.saved_tensors
+      return grad_output.masked_fill(mask, 0)*ctx.scale, None
+    else:
+      return grad_output, None
+
+class StableDropout(torch.nn.Module):
+  """ Optimized dropout module for stabilizing the training
+
+  Args:
+
+    drop_prob (float): the dropout probabilities
+
+  """
+
+  def __init__(self, drop_prob):
+    super().__init__()
+    self.drop_prob = drop_prob
+    self.count = 0
+    self.context_stack = None
+
+  def forward(self, x):
+    """ Call the module
+
+    Args:
+      
+      x (:obj:`torch.tensor`): The input tensor to apply dropout
+
+
+    """
+    if self.training and self.drop_prob>0:
+      return XDropout.apply(x, self.get_context())
+    return x
+
+  def clear_context(self):
+    self.count = 0
+    self.context_stack = None
+
+  def init_context(self, reuse_mask=True, scale = 1):
+    if self.context_stack is None:
+      self.context_stack = []
+    self.count = 0
+    for c in self.context_stack:
+      c.reuse_mask = reuse_mask
+      c.scale = scale
+
+  def get_context(self):
+    if self.context_stack is not None:
+      if self.count >= len(self.context_stack):
+        self.context_stack.append(DropoutContext())
+      ctx = self.context_stack[self.count]
+      ctx.dropout = self.drop_prob
+      self.count += 1
+      return ctx
+    else:
+      return self.drop_prob
+
+def MaskedLayerNorm(layerNorm, input, mask = None):
+  """ Masked LayerNorm which will apply mask over the output of LayerNorm to avoid inaccurate updatings to the LayerNorm module.
+  
+  Args:
+    layernorm (:obj:`~DeBERTa.deberta.LayerNorm`): LayerNorm module or function
+    input (:obj:`torch.tensor`): The input tensor
+    mask (:obj:`torch.IntTensor`): The mask to applied on the output of LayerNorm where `0` indicate the output of that element will be ignored, i.e. set to `0`
+
+  Example::
+
+    # Create a tensor b x n x d
+    x = torch.randn([1,10,100])
+    m = torch.tensor([[1,1,1,0,0,0,0,0,0,0]], dtype=torch.int)
+    LayerNorm = DeBERTa.deberta.LayerNorm(100)
+    y = MaskedLayerNorm(LayerNorm, x, m)
+
+  """
+  output = layerNorm(input).to(input)
+  if mask is None:
+    return output
+  if mask.dim()!=input.dim():
+    if mask.dim()==4:
+      mask=mask.squeeze(1).squeeze(1)
+    mask = mask.unsqueeze(2)
+  mask = mask.to(output.dtype)
+  return output*mask
+
+def gelu(x):
+  """Implementation of the gelu activation function.
+    For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
+    0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+  """
+  return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+def swish(x):
+  return x * torch.sigmoid(x)
+
+def linear_act(x):
+  return x
+
+def sequence_masking(x, mask, value=0, axis=None):
+    """为序列条件mask的函数
+    mask: 形如(batch_size, seq_len)的0-1矩阵；
+    value: mask部分要被替换成的值，可以是'-inf'或'inf'；
+    axis: 序列所在轴，默认为1；
+    """
+    if mask is None:
+        return x
+    else:
+        x_dtype = x.dtype
+        if x_dtype == torch.bool:
+            x = x.to(torch.int32)
+        # if mask.dtype != x.dtype:
+        #     mask = mask.to(x.dtype)
+        if value == '-inf':
+            value = -float('inf')
+        elif value == 'inf':
+            value = float('inf')
+        if axis is None:
+            axis = 1
+        elif axis < 0:
+            axis = x.dim() + axis
+        assert axis > 0, 'axis must be greater than 0'
+        if mask.dim() != x.dim():
+            mask = align(mask, [0, axis], x.dim())
+        # value = value.to(x.dtype)
+        x = x.masked_fill_(~mask.bool(), value) # * mask + mask.fill_(value)
+        if x_dtype == torch.bool:
+            x = x.to(torch.bool)
+        return x
+
+def align(tensor, axes, ndim=None):
+    """重新对齐tensor（批量版expand_dims）
+    axes：原来的第i维对齐新tensor的第axes[i]维；
+    ndim：新tensor的维度。
+    """
+    assert len(axes) == tensor.dim()
+    assert ndim or min(axes) >= 0
+    ndim = ndim or max(axes) + 1
+    indices = [None] * ndim
+    for i in axes:
+        indices[i] = slice(None)
+    return tensor[indices]
+
+ACT2FN = {"gelu": torch.nn.functional.gelu, "relu": torch.nn.functional.relu, "swish": swish, "tanh": torch.tanh, "linear": linear_act, 'sigmoid': torch.sigmoid, 'silu': torch.nn.functional.silu}
+
+

modeling/pooling.py

@@ -0,0 +1,88 @@
+#
+# Zhou Bo
+# 
+#
+"""
+Pooling functions
+"""
+
+from torch import nn
+import copy
+import json
+import pdb
+from .bert import ACT2FN
+from .ops import StableDropout
+from .config import AbsModelConfig
+
+__all__ = ['PoolConfig', 'ContextPooler']
+
+class PoolConfig(AbsModelConfig):
+    """Configuration class to store the configuration of `pool layer`.
+
+        Parameters:
+        
+            config (:class:`~DeBERTa.deberta.ModelConfig`): The model config. The field of pool config will be initalized with the `pooling` field in model config.
+
+        Attributes:
+
+            hidden_size (int): Size of the encoder layers and the pooler layer, default: `768`.
+
+            dropout (float): The dropout rate applied on the output of `[CLS]` token,
+
+            hidden_act (:obj:`str`): The activation function of the projection layer, it can be one of ['gelu', 'tanh'].
+
+        Example::
+
+            # Here is the content of an exmple model config file in json format
+
+                {
+                  "hidden_size": 768,
+                  "num_hidden_layers" 12,
+                  "num_attention_heads": 12,
+                  "intermediate_size": 3072,
+                  ...
+                  "pooling": {
+                    "hidden_size":  768,
+                    "hidden_act": "gelu",
+                    "dropout": 0.1
+                  }
+                }
+
+    """
+    def __init__(self, config=None):
+        """Constructs PoolConfig.
+
+        Args:
+           `config`: the config of the model. The field of pool config will be initalized with the 'pooling' field in model config.
+        """
+        
+        self.hidden_size = 768
+        self.dropout = 0
+        self.hidden_act = 'gelu'
+        if config:
+            pool_config = getattr(config, 'pooling', config)
+            if isinstance(pool_config, dict):
+                pool_config = AbsModelConfig.from_dict(pool_config)
+            self.hidden_size = getattr(pool_config, 'hidden_size', config.hidden_size)
+            self.dropout = getattr(pool_config, 'dropout', 0)
+            self.hidden_act = getattr(pool_config, 'hidden_act', 'gelu')
+
+class ContextPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = StableDropout(config.dropout)
+        self.config = config
+
+    def forward(self, hidden_states, mask = None):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        
+        context_token = hidden_states[:, 0]
+        context_token = self.dropout(context_token)
+        pooled_output = self.dense(context_token)
+        pooled_output = ACT2FN[self.config.hidden_act](pooled_output)
+        return pooled_output
+
+    def output_dim(self):
+        return self.config.hidden_size

modeling/pretrained_models.py

@@ -0,0 +1,2 @@
+
+

modeling/spm_tokenizer.py

@@ -0,0 +1,322 @@
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 11/15/2020
+#
+
+
+import sentencepiece as sp
+import six
+import unicodedata
+import os
+import regex as re
+from .cache_utils import load_vocab
+import loguru
+logger=loguru.logger
+
+
+import pdb
+
+__all__ = ['SPMTokenizer']
+
+class SPMTokenizer:
+  def __init__(self, vocab_file, do_lower_case=False, special_tokens=None, bpe_dropout=0, split_by_punct=False):
+    self.split_by_punct = split_by_punct
+    spm = sp.SentencePieceProcessor()
+    assert os.path.exists(vocab_file)
+    spm.load(vocab_file)
+    bpe_vocab_size = spm.GetPieceSize()
+    # Token map
+    # <unk> 0+1
+    # <s> 1+1
+    # </s> 2+1
+    self.vocab = {spm.IdToPiece(i):i for i in range(bpe_vocab_size)}
+    self.id_to_tokens = [spm.IdToPiece(i) for i in range(bpe_vocab_size)]
+    #self.vocab['[PAD]'] = 0
+    #self.vocab['[CLS]'] = 1
+    #self.vocab['[SEP]'] = 2
+    #self.vocab['[UNK]'] = 3
+
+    _special_tokens = ['[MASK]', '[SEP]', '[PAD]', '[UNK]', '[CLS]']
+    self.special_tokens = []
+    if special_tokens is not None:
+      _special_tokens.extend(special_tokens)
+    for t in _special_tokens:
+      self.add_special_token(t)
+
+    self.spm = spm
+    self.pat = re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
+
+  def tokenize(self, text):
+    pieces = self._encode_as_pieces(text)
+    def _norm(x):
+      if x not in self.vocab or x=='<unk>':
+        return '[UNK]'
+      else:
+        return x
+    pieces = [_norm(p) for p in pieces]
+    return pieces
+
+  def convert_tokens_to_ids(self, tokens):
+    return [self.vocab[t] if t in self.vocab else 1 for t in tokens]
+
+  def convert_ids_to_tokens(self, ids):
+    tokens = []
+    for i in ids:
+      tokens.append(self.ids_to_tokens[i])
+    return tokens
+
+  def decode(self, tokens, start=-1, end=-1, raw_text=None):
+    if raw_text is None:
+      return self.spm.decode_pieces([t for t in tokens if t not in self.special_tokens])
+    else:
+      words = self.split_to_words(raw_text)
+      word_tokens = [self.tokenize(w) for w in words]
+      wt = [w for t in word_tokens for w in t]
+      #assert tokens == wt, f'{tokens} || {wt}'
+      if wt!=tokens:
+        for a,b in zip(wt, tokens):
+          if a!=b:
+            pdb.set_trace()
+      token2words = [0]*len(tokens)
+      tid = 0
+      for i,w in enumerate(word_tokens):
+        for k,t in enumerate(w):
+          token2words[tid] = i
+          tid += 1
+      word_start = token2words[start]
+      word_end = token2words[end] if end <len(tokens) else len(words)
+      text = ''.join(words[word_start:word_end])
+      return text
+
+  def add_special_token(self, token):
+    if token not in self.special_tokens:
+      self.special_tokens.append(token)
+      if token not in self.vocab:
+        self.vocab[token] = len(self.vocab)
+        self.id_to_tokens.append(token)
+    return self.id(token)
+
+  def part_of_whole_word(self, token, is_bos=False):
+    if is_bos:
+      return True
+    if (len(token)==1 and (_is_whitespace(list(token)[0]) or _is_control(list(token)[0]) or _is_punctuation(list(token)[0]))) or token in self.special_tokens:
+      return False
+
+    word_start = b'\xe2\x96\x81'.decode('utf-8')
+    return not token.startswith(word_start)
+
+  def pad(self):
+    return '[PAD]'
+
+  def bos(self):
+    return '[CLS]'
+
+  def eos(self):
+    return '[SEP]'
+
+  def unk(self):
+      return '[UNK]'
+
+  def mask(self):
+      return '[MASK]'
+
+  def sym(self, id):
+    return self.ids_to_tokens[id]
+
+  def id(self, sym):
+    return self.vocab[sym] if sym in self.vocab else 1
+
+  def _encode_as_pieces(self, text):
+    text = convert_to_unicode(text)
+    if self.split_by_punct:
+      words = self._run_split_on_punc(text)
+      pieces = [self.spm.encode_as_pieces(w) for w in words]
+      return [p for w in pieces for p in w]
+    else:
+      return self.spm.encode_as_pieces(text)
+
+  def split_to_words(self, text):
+    pieces = self._encode_as_pieces(text)
+    word_start = b'\xe2\x96\x81'.decode('utf-8')
+    words = []
+    offset = 0
+    prev_end = 0
+    for i,p in enumerate(pieces):
+      if p.startswith(word_start):
+        if offset>prev_end:
+          words.append(text[prev_end:offset])
+        prev_end = offset
+        w = p.replace(word_start, '')
+      else:
+        w = p
+      try:
+        s = text.index(w, offset)
+        pn = ""
+        k = i+1
+        while k < len(pieces):
+          pn = pieces[k].replace(word_start, '')
+          if len(pn)>0:
+            break
+          k += 1
+
+        if len(pn)>0 and pn in text[offset:s]:
+          offset = offset + 1
+        else:
+          offset = s + len(w)
+      except:
+        offset = offset + 1
+
+    if prev_end< offset:
+      words.append(text[prev_end:offset])
+
+    return words
+
+  def _run_strip_accents(self, text):
+    """Strips accents from a piece of text."""
+    text = unicodedata.normalize("NFD", text)
+    output = []
+    for char in text:
+      cat = unicodedata.category(char)
+      if cat == "Mn":
+        continue
+      output.append(char)
+    return "".join(output)
+
+  def _run_split_on_punc(self, text):
+    """Splits punctuation on a piece of text."""
+    #words = list(re.findall(self.pat, text))
+    chars = list(text)
+    i = 0
+    start_new_word = True
+    output = []
+    while i < len(chars):
+      char = chars[i]
+      if _is_punctuation(char):
+        output.append([char])
+        start_new_word = True
+      else:
+        if start_new_word:
+          output.append([])
+        start_new_word = False
+        output[-1].append(char)
+      i += 1
+
+    return ["".join(x) for x in output]
+  
+  def _tokenize_chinese_chars(self, text):
+    """Adds whitespace around any CJK character."""
+    output = []
+    for char in text:
+      cp = ord(char)
+      if self._is_chinese_char(cp):
+        output.append(" ")
+        output.append(char)
+        output.append(" ")
+      else:
+        output.append(char)
+    return "".join(output)
+
+  def _is_chinese_char(self, cp):
+    """Checks whether CP is the codepoint of a CJK character."""
+    # This defines a "chinese character" as anything in the CJK Unicode block:
+    #   https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+    #
+    # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+    # despite its name. The modern Korean Hangul alphabet is a different block,
+    # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+    # space-separated words, so they are not treated specially and handled
+    # like the all of the other languages.
+    if ((cp >= 0x4E00 and cp <= 0x9FFF) or  #
+      (cp >= 0x3400 and cp <= 0x4DBF) or  #
+      (cp >= 0x20000 and cp <= 0x2A6DF) or  #
+      (cp >= 0x2A700 and cp <= 0x2B73F) or  #
+      (cp >= 0x2B740 and cp <= 0x2B81F) or  #
+      (cp >= 0x2B820 and cp <= 0x2CEAF) or
+      (cp >= 0xF900 and cp <= 0xFAFF) or  #
+      (cp >= 0x2F800 and cp <= 0x2FA1F)):  #
+      return True
+  
+    return False
+  
+  def _clean_text(self, text):
+    """Performs invalid character removal and whitespace cleanup on text."""
+    output = []
+    for char in text:
+      cp = ord(char)
+      if cp == 0 or cp == 0xfffd or _is_control(char):
+        continue
+      if _is_whitespace(char):
+        output.append(" ")
+      else:
+        output.append(char)
+    return "".join(output)
+
+
+def _is_whitespace(char):
+    """Checks whether `chars` is a whitespace character."""
+    # \t, \n, and \r are technically contorl characters but we treat them
+    # as whitespace since they are generally considered as such.
+    if char == " " or char == "\t" or char == "\n" or char == "\r":
+        return True
+    cat = unicodedata.category(char)
+    if cat == "Zs":
+        return True
+    return False
+
+def _is_control(char):
+    """Checks whether `chars` is a control character."""
+    # These are technically control characters but we count them as whitespace
+    # characters.
+    if char == "\t" or char == "\n" or char == "\r":
+        return False
+    cat = unicodedata.category(char)
+    if cat.startswith("C"):
+        return True
+    return False
+
+def _is_punctuation(char):
+    """Checks whether `chars` is a punctuation character."""
+    cp = ord(char)
+    # We treat all non-letter/number ASCII as punctuation.
+    # Characters such as "^", "$", and "`" are not in the Unicode
+    # Punctuation class but we treat them as punctuation anyways, for
+    # consistency.
+    if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or
+            (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)):
+        return True
+    cat = unicodedata.category(char)
+    if cat.startswith("P"):
+        return True
+    return False
+
+def whitespace_tokenize(text):
+    """Runs basic whitespace cleaning and splitting on a peice of text."""
+    text = text.strip()
+    if not text:
+        return []
+    tokens = text.split()
+    return tokens
+
+def convert_to_unicode(text):
+  """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
+  if six.PY3:
+    if isinstance(text, str):
+      return text
+    elif isinstance(text, bytes):
+      return text.decode("utf-8", "ignore")
+    else:
+      raise ValueError("Unsupported string type: %s" % (type(text)))
+  elif six.PY2:
+    if isinstance(text, str):
+      return text.decode("utf-8", "ignore")
+    elif isinstance(text, unicode):
+      return text
+    else:
+      raise ValueError("Unsupported string type: %s" % (type(text)))
+  else:
+    raise ValueError("Not running on Python2 or Python 3?")
+

modeling/tokenizers.py

@@ -0,0 +1,18 @@
+#
+# Zhou Bo
+
+#
+
+""" tokenizers
+"""
+
+from .spm_tokenizer import *
+from .gpt2_tokenizer import GPT2Tokenizer
+from wywLM.models import BertTokenizer
+
+__all__ = ['tokenizers']
+tokenizers={
+    'gpt2': GPT2Tokenizer,
+    'spm': SPMTokenizer,
+    'bert': BertTokenizer
+    }

modeling/wywlm_modeling.py

@@ -0,0 +1,446 @@
+# Copyright (c) Microsoft, Inc. 2020
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Zhou Bo
+# Date: 01/15/2020
+#
+
+import copy
+import torch
+import os
+import random
+
+import json
+from .ops import *
+from .bert import *
+from .bert import BertLayer
+from .config import ModelConfig
+from .cache_utils import load_model_state
+from .nnmodule import NNModule
+
+# from ..utils.bad_grad_viz import register_hooks
+
+__all__ = ['WywLM']
+
+def flatten_states(q_states, mask_index):
+    q_states = q_states.reshape((-1, q_states.size(-1)))
+    q_states = q_states.index_select(0, mask_index)
+    return q_states
+
+
+class UGDecoder(torch.nn.Module):
+    def __init__(self, config, vocab_size):
+        super().__init__()
+        self.config = config
+        self.position_biased_input = getattr(config, 'position_biased_input', True)
+        # self.layer = torch.nn.ModuleList([BertLayer(config) for _ in range(2)])
+
+        # self.causal_mask = torch.tril(torch.ones((input_ids.dim(0), input_ids.dim(1), input_ids.dim(1))), diagonal=0)
+
+    def forward(self, ctx_layers, word_embedding, input_ids, z_states, attention_mask, \
+                encoder, target_ids=None, relative_pos=None, decode=False, s2s_idx=None):
+        causal_outputs, lm_outputs = self.emd_context_layer(ctx_layers, z_states, attention_mask, 
+                                                encoder, target_ids, input_ids, 
+                                                relative_pos=relative_pos, decode=decode,
+                                                word_embedding=word_embedding, s2s_idx=s2s_idx)
+        # loss_fct = torch.nn.CrossEntropyLoss(reduction='none')
+        
+        # ctx_layer = mlm_ctx_layers[-1]
+
+        # lm_logits = lm_logits.view(-1, lm_logits.size(-1))
+
+        return causal_outputs[-1], lm_outputs[-1]
+
+    def emd_context_layer(self, encoder_layers, z_states, attention_mask, encoder, target_ids, input_ids,\
+                          relative_pos=None, decode=False, word_embedding=None, s2s_idx=None):
+        # if decode:
+        #     attention_mask = torch.tril(torch.ones((input_ids.shape[0], 1, input_ids.shape[1], input_ids.shape[1])), diagonal=0).to(input_ids.device)
+        # else:
+        if attention_mask.dim()<=2:
+            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            att_mask = extended_attention_mask.byte()
+            attention_mask = att_mask*att_mask.squeeze(-2).unsqueeze(-1)
+        elif attention_mask.dim()==3:
+            attention_mask = attention_mask.unsqueeze(1)
+        
+        
+        if not self.position_biased_input: 
+
+            
+            lm_outputs = []
+            # else:
+            hidden_states = encoder_layers[-2]
+            layers = [encoder.layer[-1] for _ in range(2)]
+            z_states += hidden_states
+            query_states = z_states
+            query_mask = attention_mask
+            rel_embeddings = encoder.get_rel_embedding()
+            for layer in layers:
+                # TODO: pass relative pos ids
+                output = layer(hidden_states, query_mask, return_att=False, 
+                            query_states=query_states, relative_pos=relative_pos, 
+                            rel_embeddings=rel_embeddings)
+                query_states = output
+                lm_outputs.append(query_states)
+
+            # if decode:
+            attention_mask = torch.tril(torch.ones((input_ids.shape[0], 1, input_ids.shape[1], input_ids.shape[1])), 
+                                        diagonal=0).to(input_ids.device)
+            causal_outputs = []
+            # with torch.no_grad():
+            target_embd = word_embedding(target_ids)
+
+            target_embd += z_states.detach()
+            # self attention of target
+            output = layers[-2](target_embd, attention_mask, return_att=False, 
+                        query_states=target_embd, relative_pos=relative_pos, 
+                        rel_embeddings=encoder.get_rel_embedding())
+            causal_outputs.append(output)
+            # cross attention
+            output = layers[-1](output, attention_mask, return_att=False, 
+                        query_states=query_states, relative_pos=relative_pos, 
+                        rel_embeddings=encoder.get_rel_embedding())
+            causal_outputs.append(output)
+
+        else:
+            causal_outputs = [encoder_layers[-1]]
+            lm_outputs = [encoder_layers[-1]]
+        return causal_outputs, lm_outputs
+
+
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+    """
+    Shift input ids one token to the right.
+    """
+    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 WywLMLoss(torch.nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.loss_fn = torch.nn.CrossEntropyLoss(reduction='mean')
+        hidden_size = getattr(config, 'embedding_size', config.hidden_size)
+        self.compare = torch.nn.Linear(hidden_size * 3, 2)
+        # self.mlm_head = BertLMPredictionHead(config, config.vocab_size)
+        self.lm_head = BertLMPredictionHead(config, config.vocab_size)
+
+    def forward(self, logits, lm_logits, target_ids, dict_pos, input_ids, target_ids_s2s, decode=False, ebd_weight=None, task=0):
+        loss_compare = torch.tensor(0).to(logits).float()
+        mlm_loss = torch.tensor(0).to(logits).float()
+        lm_loss = torch.tensor(0).to(logits).float()
+
+        # else:
+        if task == 1:
+            compare_logits = []
+            compare_labels = []
+            for bi, sampel_pos in enumerate(dict_pos):
+                num_pos = int((sampel_pos > 0).sum().detach().cpu().numpy() / 4) - 1
+                if num_pos <= 1:
+                    continue
+                for pi in range(num_pos):
+                    pos = sampel_pos[pi]
+                    entry_logits = logits[bi][pos[0]: pos[1]]
+                    desc_logits = logits[bi][pos[2]: pos[3]]
+                    neg_num = random.randint(0, num_pos) # torch.randint(low=0, high=num_pos, size=(1,))
+                    ids_neg = input_ids[bi][sampel_pos[neg_num][0]: sampel_pos[neg_num][1]]
+                    ids_pos = input_ids[bi][pos[0]: pos[1]]
+                    if neg_num == pi or (ids_neg.shape == ids_pos.shape and torch.all(ids_neg == ids_pos)):
+                        neg_num = -1
+                        for ni in range(num_pos):
+                            neg_num = random.randint(0, num_pos)# torch.randint(low=0, high=num_pos, size=(1,))
+                            ids_neg = input_ids[bi][sampel_pos[neg_num][0]: sampel_pos[neg_num][1]]
+                            if neg_num != pi and (ids_neg.shape != ids_pos.shape or not torch.all(ids_neg == ids_pos)):
+                                break
+                            else:
+                                neg_num = -1
+                    if neg_num == -1:
+                        continue
+                    neg_desc_logits = logits[bi][sampel_pos[neg_num][2]: sampel_pos[neg_num][3]]
+                    if torch.any(torch.isnan(neg_desc_logits)):
+                        print('error')
+                    entry_logits = entry_logits.mean(dim=0, keepdim=True).float()
+                    desc_logits = desc_logits.mean(dim=0, keepdim=True).float()
+                    neg_desc_logits = neg_desc_logits.mean(dim=0, keepdim=True).float()
+                    compare_logits.append(torch.concat([entry_logits, desc_logits, entry_logits - desc_logits], dim=1))
+                    compare_logits.append(torch.concat([entry_logits, neg_desc_logits, entry_logits - neg_desc_logits], dim=1))
+                    compare_labels += [1, 0]
+            if len(compare_logits) > 0:
+                compare_logits = torch.concat(compare_logits, dim=0).to(logits.dtype)
+                compare_pred = self.compare(compare_logits)
+                loss_compare = self.loss_fn(compare_pred, torch.tensor(compare_labels, dtype=torch.long, device=compare_logits.device)).mean()
+
+        if torch.all(loss_compare == 0):
+            entry_logits = logits[0][0].unsqueeze(0)
+            compare_logits = torch.concat([entry_logits, entry_logits, entry_logits - entry_logits], dim=1)
+            compare_pred = self.compare(compare_logits)
+            compare_labels = [1]
+            loss_compare = self.loss_fn(compare_pred, torch.tensor(compare_labels, dtype=torch.long, device=compare_logits.device)).mean()
+
+        # if decode:
+        # lm_labels = target_ids_s2s.index_select(0, (target_ids_s2s.sum(-1) > 0).nonzero().view(-1)[0])
+        # lm_labels = lm_labels.repeat(logits.shape[0], 1).clone().view(-1)
+        # lm_labels = target_ids_s2s.clone()
+        # target_ids_s2s = shift_tokens_right(target_ids_s2s, 0, 1)
+        # target_ids_s2s.masked_fill_(target_ids_s2s==0, 3)
+        if task == 0:
+            _mask_index = (target_ids_s2s > 0).view(-1).nonzero().view(-1)
+            lm_logits_ = flatten_states(lm_logits, _mask_index)
+            lm_pred = self.lm_head(lm_logits_, ebd_weight).float()
+            lm_labels = target_ids_s2s.clone().reshape(-1)
+            lm_labels = lm_labels.index_select(0, _mask_index)
+            # lm_pred = torch.nn.functional.log_softmax(lm_pred)
+            # lm_loss = torch.nn.functional.nll_loss(lm_pred, lm_labels.long())
+            lm_loss = self.loss_fn(lm_pred, lm_labels.long())
+        # dot = register_hooks(lm_loss)
+        # lm_loss.backward()
+        # dot().save('tmp.dot')
+
+
+        _mask_index = (target_ids > 0).view(-1).nonzero().view(-1)
+        mlm_logits = flatten_states(logits, _mask_index)
+        mlm_pred = self.lm_head(mlm_logits, ebd_weight).float()
+        mlm_labels = target_ids.view(-1)
+        mlm_labels = mlm_labels.index_select(0, _mask_index)
+        mlm_loss = self.loss_fn(mlm_pred, mlm_labels.long())
+        return loss_compare, mlm_loss, lm_loss
+
+class WywLM(torch.nn.Module):
+    """ DeBERTa encoder
+    This module is composed of the input embedding layer with stacked transformer layers with disentangled attention.
+
+    Parameters:
+        config:
+            A model config class instance with the configuration to build a new model. The schema is similar to `BertConfig`, \
+                    for more details, please refer :class:`~DeBERTa.deberta.ModelConfig`
+
+        pre_trained:
+            The pre-trained DeBERTa model, it can be a physical path of a pre-trained DeBERTa model or a released configurations, \
+                    i.e. [**base, large, base_mnli, large_mnli**]
+
+    """
+
+    def __init__(self, config=None, pre_trained=None):
+        super().__init__()
+        state = None
+        if pre_trained is not None:
+            state, model_config = load_model_state(pre_trained)
+            if config is not None and model_config is not None:
+                for k in config.__dict__:
+                    if k not in ['hidden_size',
+                        'intermediate_size',
+                        'num_attention_heads',
+                        'num_hidden_layers',
+                        'vocab_size',
+                        'max_position_embeddings']:
+                        model_config.__dict__[k] = config.__dict__[k]
+            config = copy.copy(model_config)
+        self.embeddings = BertEmbeddings(config)
+        self.encoder = BertEncoder(config)
+        self.config = config
+        self.pre_trained = pre_trained
+        self.apply_state(state)
+
+    def forward(self, input_ids, attention_mask=None, token_type_ids=None, output_all_encoded_layers=True, position_ids = None, return_att = False):
+        """
+        Args:
+            input_ids:
+                a torch.LongTensor of shape [batch_size, sequence_length] \
+            with the word token indices in the vocabulary
+
+            attention_mask:
+                an optional parameter for input mask or attention mask.
+
+                - If it's an input mask, then it will be torch.LongTensor of shape [batch_size, sequence_length] with indices \
+            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max \
+            input sequence length in the current batch. It's the mask that we typically use for attention when \
+            a batch has varying length sentences.
+
+                - If it's an attention mask then it will be torch.LongTensor of shape [batch_size, sequence_length, sequence_length]. \
+            In this case, it's a mask indicate which tokens in the sequence should be attended by other tokens in the sequence.
+
+            token_type_ids:
+                an optional torch.LongTensor of shape [batch_size, sequence_length] with the token \
+            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to \
+            a `sentence B` token (see BERT paper for more details).
+
+            output_all_encoded_layers:
+                whether to output results of all encoder layers, default, True
+
+        Returns:
+
+            - The output of the stacked transformer layers if `output_all_encoded_layers=True`, else \
+            the last layer of stacked transformer layers
+
+            - Attention matrix of self-attention layers if `return_att=True`
+
+
+        Example::
+
+            # Batch of wordPiece token ids.
+            # Each sample was padded with zero to the maxium length of the batch
+            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
+            # Mask of valid input ids
+            attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
+
+            # DeBERTa model initialized with pretrained base model
+            bert = DeBERTa(pre_trained='base')
+
+            encoder_layers = bert(input_ids, attention_mask=attention_mask)
+
+        """
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros_like(input_ids)
+            token_mask = torch.ones_like(input_ids)
+        else:
+            idxs = torch.flip(torch.cumsum(torch.flip(token_type_ids, [-1]), axis=1), [-1])
+            token_mask = idxs > 0
+            token_mask = token_mask.byte()
+        ebd_output = self.embeddings(input_ids.to(torch.long), token_type_ids.to(torch.long), position_ids, token_mask)
+        embedding_output = ebd_output['embeddings']
+        encoder_output = self.encoder(embedding_output,
+                                     attention_mask,
+                                     output_all_encoded_layers=output_all_encoded_layers, return_att = return_att)
+        encoder_output.update(ebd_output)
+        return encoder_output
+
+    def apply_state(self, state = None):
+        """ Load state from previous loaded model state dictionary.
+
+            Args:
+                state (:obj:`dict`, optional): State dictionary as the state returned by torch.module.state_dict(), default: `None`. \
+                        If it's `None`, then will use the pre-trained state loaded via the constructor to re-initialize \
+                        the `DeBERTa` model
+        """
+        if self.pre_trained is None and state is None:
+            return
+        if state is None:
+            state, config = load_model_state(self.pre_trained)
+            self.config = config
+        
+        prefix = ''
+        for k in state:
+            if 'embeddings.' in k:
+                if not k.startswith('embeddings.'):
+                    prefix = k[:k.index('embeddings.')]
+                break
+
+        missing_keys = []
+        unexpected_keys = []
+        error_msgs = []
+        self._load_from_state_dict(state, prefix = prefix, local_metadata=None, strict=True, missing_keys=missing_keys, unexpected_keys=unexpected_keys, error_msgs=error_msgs)
+
+
+class MaskedLanguageModel(NNModule):
+    """ Masked language model
+    """
+    def __init__(self, config, *wargs, **kwargs):
+        super().__init__(config)
+        self.backbone = WywLM(config)
+
+        self.max_relative_positions = getattr(config, 'max_relative_positions', -1)
+        self.position_buckets = getattr(config, 'position_buckets', -1)
+        if self.max_relative_positions <1:
+            self.max_relative_positions = config.max_position_embeddings
+        # self.mlm_predictions = UGDecoder(self.backbone.config, self.backbone.embeddings.word_embeddings.weight.size(0))
+        self.lm_predictions = UGDecoder(self.backbone.config, self.backbone.embeddings.word_embeddings.weight.size(0))
+        self.device = None
+        self.loss = WywLMLoss(config)
+        # self.loss_lm = WywLMLoss(config)
+        self.apply(self.init_weights)
+
+    def forward(self, samples, position_ids=None):
+        task = samples['task']
+        if task == 0:
+            input_ids = samples['s2s_input_ids']
+            type_ids = samples['s2s_token_type_ids']
+            attention_mask = samples['s2s_attention_mask']
+            labels = samples['s2s_masked_lm_labels']
+            dict_pos = samples['dict_pos']
+            s2s_label = samples['s2s_label']
+        else:
+            input_ids = samples['input_ids']
+            type_ids = samples['token_type_ids']
+            attention_mask = samples['attention_mask']
+            labels = samples['masked_lm_labels']
+            dict_pos = samples['dict_pos']
+            s2s_label = samples['s2s_label']
+        
+        if self.device is None:
+            self.device = list(self.parameters())[0].device
+        
+        input_ids = input_ids.to(self.device)
+
+        type_ids = None
+        lm_labels = labels.to(self.device)
+        s2s_label = s2s_label.to(self.device)
+        attention_mask = attention_mask.to(self.device)
+
+        encoder_output = self.backbone(input_ids, attention_mask, type_ids, output_all_encoded_layers=True, position_ids = position_ids)
+        encoder_layers = encoder_output['hidden_states']
+        z_states = encoder_output['position_embeddings']
+        ctx_layer = encoder_layers[-1]
+        mlm_loss = torch.tensor(0).to(ctx_layer).float()
+        lm_loss = torch.tensor(0).to(ctx_layer).float()
+        lm_logits = None
+        label_inputs = None
+        loss = torch.tensor(0).to(ctx_layer).float()
+        loss_compare = torch.tensor(0).to(ctx_layer).float()
+
+        ebd_weight = self.backbone.embeddings.word_embeddings.weight
+        lm_logits, mlm_logits = self.lm_predictions(encoder_layers, self.backbone.embeddings.word_embeddings, 
+                                        input_ids, z_states, 
+                                        attention_mask, self.backbone.encoder,
+                                        target_ids=lm_labels)
+        # if lm_labels.detach().sum() != 0:
+        loss_compare, mlm_loss, lm_loss = self.loss(mlm_logits, 
+                                                    lm_logits,
+                                                    lm_labels, 
+                                                    dict_pos, 
+                                                    target_ids_s2s=s2s_label,
+                                                    decode=False, 
+                                                    ebd_weight=ebd_weight, 
+                                                    input_ids=input_ids,
+                                                    task=task)
+        loss = loss_compare * 10 + mlm_loss + lm_loss
+        # if s2s_label.detach().sum() != 0:
+        #     s2s_idx = (s2s_label.sum(-1)>0).nonzero().view(-1)
+        #     s2s_label = s2s_label.index_select(0, s2s_idx)
+        #     # ebd_weight = self.backbone.embeddings.word_embeddings.weight
+        #     # lm_logits = self.lm_predictions(encoder_layers[-3], self.backbone.embeddings.word_embeddings, 
+        #     #                                 input_ids.index_select(0, s2s_idx), z_states.index_select(0, s2s_idx), 
+        #     #                                 attention_mask.index_select(0, s2s_idx), self.backbone.encoder,
+        #     #                                 target_ids=s2s_label,
+        #     #                                 decode=True, s2s_idx=s2s_idx)
+        #     # lm_logits = encoder_layers[-1].detach().index_select(0, s2s_idx)
+        #     _, lm_loss = self.loss_lm(lm_logits, 
+        #                                 s2s_label, 
+        #                                 torch.zeros_like(dict_pos), 
+        #                                 decode=True, 
+        #                                 ebd_weight=ebd_weight, 
+        #                                 input_ids=input_ids.index_select(0, s2s_idx))
+            # lm_loss = lm_logits.max()
+            # loss = loss + lm_loss
+
+        return {
+                'logits' : lm_logits,
+                'labels' : lm_labels,
+                's2s_label': s2s_label,
+                'loss' : loss.float(),
+                'loss_compare': loss_compare.float(),
+                'lm_loss': lm_loss.float(),
+                'mlm_loss': mlm_loss.float()
+            }