upload

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+imgs/overview.png filter=lfs diff=lfs merge=lfs -text

INIT.sh

@@ -0,0 +1,23 @@
+sudo rm -f  /etc/conda/condarc && sudo touch /etc/conda/condarc
+conda create -n py38 python=3.8 -y
+conda install -n py38 ipykernel -y
+source activate py38
+# local env
+# conda create --prefix=conda_env/py38_env python=3.8 -y
+# conda install  --prefix=conda_env/py38_env ipykernel -y
+# source activate conda_env/py38_env
+
+
+python -m ipykernel install --user --name py38 --display-name "py38"
+# pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
+pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
+# pip install matplotlib datasets pandas transformers scikit-learn scipy tensorboard tqdm numpy seaborn fairseq tensorboardX
+
+pip install -r requirements.txt 
+pip install --user -U https://pai-dlc.oss-cn-zhangjiakou.aliyuncs.com/tunnel/common_io/common_io-0.4.1%2Btunnel-py2.py3-none-any.whl 
+pip install oss2
+pip install ujson cn2an whoosh openpyxl rapidfuzz numpy pandas tqdm jieba scikit-learn seaborn
+# pip install http://eas-data.oss-cn-shanghai.aliyuncs.com/sdk/allspark-0.15-py2.py3-none-any.whl
+pip install http://eas-data.oss-cn-shanghai.aliyuncs.com/sdk/allspark-0.15-py2.py3-none-any.whl
+# pip install "modelscope[nlp]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
+# pip install --user -U https://pai-dlc.oss-cn-zhangjiakou.aliyuncs.com/tunnel/common_io/common_io-0.4.1%2Btunnel-py2.py3-none-any.whl

README.md

@@ -0,0 +1,92 @@
+
+# TAAS
+
+## Introduction
+TAAS: A Text-based Delivery Address Analysis System in Logistics
+
+## System description
+TAAS is an integrated system for text-based address analysis in logistics field. TAAS supports several address perception tasks, as well as other logistics related tasks. Our system is based on a Geography-Graph Pre-trained model in logistics, termed G2PTL, which promotes the delivery address encoding by combining the semantic learning capabilities of text pre-training with the geographical-relationship encoding abilities of graph modeling. 
+
+![overview.png](./imgs/overview.png)
+
+## Supported Tasks
+
+1. **Address perception tasks**
+* Address Completion
+* Address Standardization
+* House Info Extraction
+* Address Entity Tokenization
+* Address embedding
+2. **Logistics related tasks**
+* Geo-locating From Text to Geospatial
+* Pick-up Estimation Time of Arrival
+* Pick-up and Delivery Route Prediction
+
+## How To Use
+
+Once installed, loading and using a fine-tuned model on any specific task can be done as follows:
+
+```python
+from transformers import AutoModel
+model = AutoModel.from_pretrained('Cainiao-AI/TAAS')
+model.eval()
+address = ['北京市马驹桥镇兴贸二街幸福家园1幢5单元1009室 注：放在门口即可']
+
+# Address completion
+output = model.addr_complet(address)
+print(output)
+```
+```python
+['北京市通州区马驹桥镇兴贸二街幸福家园1幢5单元1009室 注：放在门口即可']
+```
+```python
+# Address standardization
+output = model.addr_standardize(address)
+print(output)
+```
+```python
+['北京马驹桥镇兴贸二街幸福家园1幢5单元1009室']
+```
+```python
+# House info extraction
+output = model.house_info(address)
+print(output)
+```
+```python
+[{'楼栋': '1', '单元': '5', '门牌号': '1009'}]
+```
+```python
+# Address entity tokenization
+output = model.addr_entity(address)
+print(output)
+```
+```python
+[{'省': '北京', '市': '', '区': '马驹桥', '街道/镇': '镇兴贸二街', '道路': '', '道路号': '', 'poi': '幸福家园', '楼栋号': '1', '单元号': '5', '门牌号': '1009'}]
+```
+```python
+# Geo-locating from text to geospatial
+output = model.geolocate(address)
+```
+```python
+'s2网格化结果：453cf541fcb147b437433cf3cff43f470'
+```
+```python
+# Pick-up estimation time of arrival
+output = model.pickup_ETA(eta_data)
+# Users can get the address embeddings for their pick-up ETA model
+```
+```python
+# Pick-up and Delivery Route prediction
+output = model.route_predict(route_data)
+# Users can get the address embeddings for their route prediction model
+```
+
+## Requirements
+python>=3.8
+```shell
+tqdm==4.65.0
+torch==1.13.1
+transformers==4.27.4
+datasets==2.11.0
+fairseq==0.12.2
+```

TAAS_utils.py

@@ -0,0 +1,1544 @@
+#! python3
+# -*- encoding: utf-8 -*-
+
+from transformers.models.ernie.modeling_ernie import *
+import torch.utils.checkpoint
+from torch import nn
+from transformers.utils import logging
+import inspect
+from typing import Set, Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union
+import re
+import math
+from typing import Optional, Tuple
+
+import torch
+from fairseq import utils
+from fairseq.modules.fairseq_dropout import FairseqDropout
+from fairseq.modules.quant_noise import quant_noise
+from torch import Tensor, nn
+from torch.hub import load_state_dict_from_url
+import torch.distributed as dist
+
+logger = logging.get_logger(__name__)
+
+from torch.hub import load_state_dict_from_url
+import torch.distributed as dist
+
+PRETRAINED_MODEL_URLS = {
+  "pcqm4mv1_graphormer_base":"https://ml2md.blob.core.windows.net/graphormer-ckpts/checkpoint_best_pcqm4mv1.pt",
+  "pcqm4mv2_graphormer_base":"https://ml2md.blob.core.windows.net/graphormer-ckpts/checkpoint_best_pcqm4mv2.pt",
+  "oc20is2re_graphormer3d_base":"https://szheng.blob.core.windows.net/graphormer/modelzoo/oc20is2re/checkpoint_last_oc20_is2re.pt", # this pretrained model is temporarily unavailable
+  "pcqm4mv1_graphormer_base_for_molhiv":"https://ml2md.blob.core.windows.net/graphormer-ckpts/checkpoint_base_preln_pcqm4mv1_for_hiv.pt",
+}
+
+def load_pretrained_model(pretrained_model_name):
+    if pretrained_model_name not in PRETRAINED_MODEL_URLS:
+        raise ValueError("Unknown pretrained model name %s", pretrained_model_name)
+    if not dist.is_initialized():
+        return load_state_dict_from_url(PRETRAINED_MODEL_URLS[pretrained_model_name], progress=True)["model"]
+    else:
+        pretrained_model = load_state_dict_from_url(PRETRAINED_MODEL_URLS[pretrained_model_name], progress=True, file_name=f"{pretrained_model_name}_{dist.get_rank()}")["model"]
+        dist.barrier()
+        return pretrained_model
+
+
+class MultiheadAttention(nn.Module):
+    """Multi-headed attention.
+
+    See "Attention Is All You Need" for more details.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        self_attention=False,
+        q_noise=0.0,
+        qn_block_size=8,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout_module = FairseqDropout(
+            dropout, module_name=self.__class__.__name__
+        )
+
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim ** -0.5
+
+        self.self_attention = self_attention
+
+        assert self.self_attention, "Only support self attention"
+
+        assert not self.self_attention or self.qkv_same_dim, (
+            "Self-attention requires query, key and " "value to be of the same size"
+        )
+
+        self.k_proj = quant_noise(
+            nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size
+        )
+        self.v_proj = quant_noise(
+            nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size
+        )
+        self.q_proj = quant_noise(
+            nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
+        )
+
+        self.out_proj = quant_noise(
+            nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size
+        )
+
+        self.reset_parameters()
+
+        self.onnx_trace = False
+
+    def prepare_for_onnx_export_(self):
+        raise NotImplementedError
+
+    def reset_parameters(self):
+        if self.qkv_same_dim:
+            # Empirically observed the convergence to be much better with
+            # the scaled initialization
+            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
+        else:
+            nn.init.xavier_uniform_(self.k_proj.weight)
+            nn.init.xavier_uniform_(self.v_proj.weight)
+            nn.init.xavier_uniform_(self.q_proj.weight)
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+
+    def forward(
+        self,
+        query,
+        key: Optional[Tensor],
+        value: Optional[Tensor],
+        attn_bias: Optional[Tensor],
+        key_padding_mask: Optional[Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[Tensor] = None,
+        before_softmax: bool = False,
+        need_head_weights: bool = False,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """Input shape: Time x Batch x Channel
+
+        Args:
+            key_padding_mask (ByteTensor, optional): mask to exclude
+                keys that are pads, of shape `(batch, src_len)`, where
+                padding elements are indicated by 1s.
+            need_weights (bool, optional): return the attention weights,
+                averaged over heads (default: False).
+            attn_mask (ByteTensor, optional): typically used to
+                implement causal attention, where the mask prevents the
+                attention from looking forward in time (default: None).
+            before_softmax (bool, optional): return the raw attention
+                weights and values before the attention softmax.
+            need_head_weights (bool, optional): return the attention
+                weights for each head. Implies *need_weights*. Default:
+                return the average attention weights over all heads.
+        """
+        if need_head_weights:
+            need_weights = True
+
+        tgt_len, bsz, embed_dim = query.size()
+        src_len = tgt_len
+        assert embed_dim == self.embed_dim, f"query dim {embed_dim} != {self.embed_dim}"
+        assert list(query.size()) == [tgt_len, bsz, embed_dim]
+        if key is not None:
+            src_len, key_bsz, _ = key.size()
+            if not torch.jit.is_scripting():
+                assert key_bsz == bsz
+                assert value is not None
+                assert src_len, bsz == value.shape[:2]
+
+        q = self.q_proj(query)
+        k = self.k_proj(query)
+        v = self.v_proj(query)
+        q *= self.scaling
+
+        q = (
+            q.contiguous()
+            .view(tgt_len, bsz * self.num_heads, self.head_dim)
+            .transpose(0, 1)
+        )
+        if k is not None:
+            k = (
+                k.contiguous()
+                .view(-1, bsz * self.num_heads, self.head_dim)
+                .transpose(0, 1)
+            )
+        if v is not None:
+            v = (
+                v.contiguous()
+                .view(-1, bsz * self.num_heads, self.head_dim)
+                .transpose(0, 1)
+            )
+
+        assert k is not None
+        assert k.size(1) == src_len
+
+        # This is part of a workaround to get around fork/join parallelism
+        # not supporting Optional types.
+        if key_padding_mask is not None and key_padding_mask.dim() == 0:
+            key_padding_mask = None
+
+        if key_padding_mask is not None:
+            assert key_padding_mask.size(0) == bsz
+            assert key_padding_mask.size(1) == src_len
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+        attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
+
+        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_bias is not None:
+            attn_weights += attn_bias.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            # don't attend to padding symbols
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool),
+                float("-inf"),
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if before_softmax:
+            return attn_weights, v
+
+        attn_weights_float = utils.softmax(
+            attn_weights, dim=-1, onnx_trace=self.onnx_trace
+        )
+        attn_weights = attn_weights_float.type_as(attn_weights)
+        attn_probs = self.dropout_module(attn_weights)
+
+        assert v is not None
+        attn = torch.bmm(attn_probs, v)
+        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
+
+        attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn = self.out_proj(attn)
+
+        attn_weights: Optional[Tensor] = None
+        if need_weights:
+            attn_weights = attn_weights_float.view(
+                bsz, self.num_heads, tgt_len, src_len
+            ).transpose(1, 0)
+            if not need_head_weights:
+                # average attention weights over heads
+                attn_weights = attn_weights.mean(dim=0)
+
+        return attn, attn_weights
+
+    def apply_sparse_mask(self, attn_weights, tgt_len: int, src_len: int, bsz: int):
+        return attn_weights
+
+    def upgrade_state_dict_named(self, state_dict, name):
+        prefix = name + "." if name != "" else ""
+        items_to_add = {}
+        keys_to_remove = []
+        for k in state_dict.keys():
+            if k.endswith(prefix + "in_proj_weight"):
+                # in_proj_weight used to be q + k + v with same dimensions
+                dim = int(state_dict[k].shape[0] / 3)
+                items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim]
+                items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim]
+                items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :]
+
+                keys_to_remove.append(k)
+
+                k_bias = prefix + "in_proj_bias"
+                if k_bias in state_dict.keys():
+                    dim = int(state_dict[k].shape[0] / 3)
+                    items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]
+                    items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][
+                        dim : 2 * dim
+                    ]
+                    items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]
+
+                    keys_to_remove.append(prefix + "in_proj_bias")
+
+        for k in keys_to_remove:
+            del state_dict[k]
+
+        for key, value in items_to_add.items():
+            state_dict[key] = value
+
+
+def init_graphormer_params(module):
+    """
+    Initialize the weights specific to the Graphormer Model.
+    """
+
+    def normal_(data):
+        # with FSDP, module params will be on CUDA, so we cast them back to CPU
+        # so that the RNG is consistent with and without FSDP
+        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+    if isinstance(module, nn.Linear):
+        normal_(module.weight.data)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        normal_(module.weight.data)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        normal_(module.q_proj.weight.data)
+        normal_(module.k_proj.weight.data)
+        normal_(module.v_proj.weight.data)
+
+
+
+
+def add_start_docstrings(*docstr):
+    def docstring_decorator(fn):
+        fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
+        return fn
+
+    return docstring_decorator
+
+
+def add_start_docstrings_to_model_forward(*docstr):
+    def docstring_decorator(fn):
+        docstring = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
+        class_name = f"[`{fn.__qualname__.split('.')[0]}`]"
+        intro = f"   The {class_name} forward method, overrides the `__call__` special method."
+        note = r"""
+
+    <Tip>
+
+    Although the recipe for forward pass needs to be defined within this function, one should call the [`Module`]
+    instance afterwards instead of this since the former takes care of running the pre and post processing steps while
+    the latter silently ignores them.
+
+    </Tip>
+"""
+
+        fn.__doc__ = intro + note + docstring
+        return fn
+
+    return docstring_decorator
+
+
+def add_end_docstrings(*docstr):
+    def docstring_decorator(fn):
+        fn.__doc__ = (fn.__doc__ if fn.__doc__ is not None else "") + "".join(docstr)
+        return fn
+
+    return docstring_decorator
+
+
+PT_RETURN_INTRODUCTION = r"""
+    Returns:
+        [`{full_output_type}`] or `tuple(torch.FloatTensor)`: A [`{full_output_type}`] or a tuple of
+        `torch.FloatTensor` (if `return_dict=False` is passed or when `config.return_dict=False`) comprising various
+        elements depending on the configuration ([`{config_class}`]) and inputs.
+
+"""
+
+TF_RETURN_INTRODUCTION = r"""
+    Returns:
+        [`{full_output_type}`] or `tuple(tf.Tensor)`: A [`{full_output_type}`] or a tuple of `tf.Tensor` (if
+        `return_dict=False` is passed or when `config.return_dict=False`) comprising various elements depending on the
+        configuration ([`{config_class}`]) and inputs.
+
+"""
+
+
+def _get_indent(t):
+    """Returns the indentation in the first line of t"""
+    search = re.search(r"^(\s*)\S", t)
+    return "" if search is None else search.groups()[0]
+
+
+def _convert_output_args_doc(output_args_doc):
+    """Convert output_args_doc to display properly."""
+    # Split output_arg_doc in blocks argument/description
+    indent = _get_indent(output_args_doc)
+    blocks = []
+    current_block = ""
+    for line in output_args_doc.split("\n"):
+        # If the indent is the same as the beginning, the line is the name of new arg.
+        if _get_indent(line) == indent:
+            if len(current_block) > 0:
+                blocks.append(current_block[:-1])
+            current_block = f"{line}\n"
+        else:
+            # Otherwise it's part of the description of the current arg.
+            # We need to remove 2 spaces to the indentation.
+            current_block += f"{line[2:]}\n"
+    blocks.append(current_block[:-1])
+
+    # Format each block for proper rendering
+    for i in range(len(blocks)):
+        blocks[i] = re.sub(r"^(\s+)(\S+)(\s+)", r"\1- **\2**\3", blocks[i])
+        blocks[i] = re.sub(r":\s*\n\s*(\S)", r" -- \1", blocks[i])
+
+    return "\n".join(blocks)
+
+
+def _prepare_output_docstrings(output_type, config_class, min_indent=None):
+    """
+    Prepares the return part of the docstring using `output_type`.
+    """
+    output_docstring = output_type.__doc__
+
+    # Remove the head of the docstring to keep the list of args only
+    lines = output_docstring.split("\n")
+    i = 0
+    while i < len(lines) and re.search(r"^\s*(Args|Parameters):\s*$", lines[i]) is None:
+        i += 1
+    if i < len(lines):
+        params_docstring = "\n".join(lines[(i + 1):])
+        params_docstring = _convert_output_args_doc(params_docstring)
+
+    # Add the return introduction
+    full_output_type = f"{output_type.__module__}.{output_type.__name__}"
+    intro = TF_RETURN_INTRODUCTION if output_type.__name__.startswith("TF") else PT_RETURN_INTRODUCTION
+    intro = intro.format(full_output_type=full_output_type, config_class=config_class)
+    result = intro + params_docstring
+
+    # Apply minimum indent if necessary
+    if min_indent is not None:
+        lines = result.split("\n")
+        # Find the indent of the first nonempty line
+        i = 0
+        while len(lines[i]) == 0:
+            i += 1
+        indent = len(_get_indent(lines[i]))
+        # If too small, add indentation to all nonempty lines
+        if indent < min_indent:
+            to_add = " " * (min_indent - indent)
+            lines = [(f"{to_add}{line}" if len(line) > 0 else line) for line in lines]
+            result = "\n".join(lines)
+
+    return result
+
+
+PT_TOKEN_CLASSIFICATION_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer(
+    ...     "HuggingFace is a company based in Paris and New York", add_special_tokens=False, return_tensors="pt"
+    ... )
+
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> predicted_token_class_ids = logits.argmax(-1)
+
+    >>> # Note that tokens are classified rather then input words which means that
+    >>> # there might be more predicted token classes than words.
+    >>> # Multiple token classes might account for the same word
+    >>> predicted_tokens_classes = [model.config.id2label[t.item()] for t in predicted_token_class_ids[0]]
+    >>> predicted_tokens_classes
+    {expected_output}
+    ```
+
+    ```python
+    >>> labels = predicted_token_class_ids
+    >>> loss = model(**inputs, labels=labels).loss
+    >>> round(loss.item(), 2)
+    {expected_loss}
+    ```
+"""
+
+PT_QUESTION_ANSWERING_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+
+    >>> inputs = tokenizer(question, text, return_tensors="pt")
+    >>> with torch.no_grad():
+    ...     outputs = model(**inputs)
+
+    >>> answer_start_index = outputs.start_logits.argmax()
+    >>> answer_end_index = outputs.end_logits.argmax()
+
+    >>> predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1]
+    >>> tokenizer.decode(predict_answer_tokens)
+    {expected_output}
+    ```
+
+    ```python
+    >>> # target is "nice puppet"
+    >>> target_start_index = torch.tensor([{qa_target_start_index}])
+    >>> target_end_index = torch.tensor([{qa_target_end_index}])
+
+    >>> outputs = model(**inputs, start_positions=target_start_index, end_positions=target_end_index)
+    >>> loss = outputs.loss
+    >>> round(loss.item(), 2)
+    {expected_loss}
+    ```
+"""
+
+PT_SEQUENCE_CLASSIFICATION_SAMPLE = r"""
+    Example of single-label classification:
+
+    ```python
+    >>> import torch
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> predicted_class_id = logits.argmax().item()
+    >>> model.config.id2label[predicted_class_id]
+    {expected_output}
+    ```
+
+    ```python
+    >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
+    >>> num_labels = len(model.config.id2label)
+    >>> model = {model_class}.from_pretrained("{checkpoint}", num_labels=num_labels)
+
+    >>> labels = torch.tensor([1])
+    >>> loss = model(**inputs, labels=labels).loss
+    >>> round(loss.item(), 2)
+    {expected_loss}
+    ```
+
+    Example of multi-label classification:
+
+    ```python
+    >>> import torch
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}", problem_type="multi_label_classification")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> predicted_class_id = logits.argmax().item()
+    >>> model.config.id2label[predicted_class_id]
+    {expected_output}
+    ```
+
+    ```python
+    >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
+    >>> num_labels = len(model.config.id2label)
+    >>> model = {model_class}.from_pretrained(
+    ...     "{checkpoint}", num_labels=num_labels, problem_type="multi_label_classification"
+    ... )
+
+    >>> labels = torch.nn.functional.one_hot(torch.tensor([predicted_class_id]), num_classes=num_labels).to(
+    ...     torch.float
+    ... )
+    >>> loss = model(**inputs, labels=labels).loss
+    >>> loss.backward()  # doctest: +IGNORE_RESULT
+    ```
+"""
+
+PT_MASKED_LM_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("The capital of France is {mask}.", return_tensors="pt")
+
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> # retrieve index of {mask}
+    >>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
+
+    >>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
+    >>> tokenizer.decode(predicted_token_id)
+    {expected_output}
+    ```
+
+    ```python
+    >>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"]
+    >>> # mask labels of non-{mask} tokens
+    >>> labels = torch.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100)
+
+    >>> outputs = model(**inputs, labels=labels)
+    >>> round(outputs.loss.item(), 2)
+    {expected_loss}
+    ```
+"""
+
+PT_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+    >>> outputs = model(**inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    ```
+"""
+
+PT_MULTIPLE_CHOICE_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
+    >>> choice0 = "It is eaten with a fork and a knife."
+    >>> choice1 = "It is eaten while held in the hand."
+    >>> labels = torch.tensor(0).unsqueeze(0)  # choice0 is correct (according to Wikipedia ;)), batch size 1
+
+    >>> encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="pt", padding=True)
+    >>> outputs = model(**{{k: v.unsqueeze(0) for k, v in encoding.items()}}, labels=labels)  # batch size is 1
+
+    >>> # the linear classifier still needs to be trained
+    >>> loss = outputs.loss
+    >>> logits = outputs.logits
+    ```
+"""
+
+PT_CAUSAL_LM_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> import torch
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+    >>> outputs = model(**inputs, labels=inputs["input_ids"])
+    >>> loss = outputs.loss
+    >>> logits = outputs.logits
+    ```
+"""
+
+PT_SPEECH_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+    >>> from datasets import load_dataset
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> processor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
+    >>> with torch.no_grad():
+    ...     outputs = model(**inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    >>> list(last_hidden_states.shape)
+    {expected_output}
+    ```
+"""
+
+PT_SPEECH_CTC_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+    >>> import torch
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> processor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+    >>> predicted_ids = torch.argmax(logits, dim=-1)
+
+    >>> # transcribe speech
+    >>> transcription = processor.batch_decode(predicted_ids)
+    >>> transcription[0]
+    {expected_output}
+    ```
+
+    ```python
+    >>> inputs["labels"] = processor(text=dataset[0]["text"], return_tensors="pt").input_ids
+
+    >>> # compute loss
+    >>> loss = model(**inputs).loss
+    >>> round(loss.item(), 2)
+    {expected_loss}
+    ```
+"""
+
+PT_SPEECH_SEQ_CLASS_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+    >>> import torch
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = feature_extractor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
+
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> predicted_class_ids = torch.argmax(logits, dim=-1).item()
+    >>> predicted_label = model.config.id2label[predicted_class_ids]
+    >>> predicted_label
+    {expected_output}
+    ```
+
+    ```python
+    >>> # compute loss - target_label is e.g. "down"
+    >>> target_label = model.config.id2label[0]
+    >>> inputs["labels"] = torch.tensor([model.config.label2id[target_label]])
+    >>> loss = model(**inputs).loss
+    >>> round(loss.item(), 2)
+    {expected_loss}
+    ```
+"""
+
+PT_SPEECH_FRAME_CLASS_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+    >>> import torch
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = feature_extractor(dataset[0]["audio"]["array"], return_tensors="pt", sampling_rate=sampling_rate)
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> probabilities = torch.sigmoid(logits[0])
+    >>> # labels is a one-hot array of shape (num_frames, num_speakers)
+    >>> labels = (probabilities > 0.5).long()
+    >>> labels[0].tolist()
+    {expected_output}
+    ```
+"""
+
+PT_SPEECH_XVECTOR_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+    >>> import torch
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = feature_extractor(
+    ...     [d["array"] for d in dataset[:2]["audio"]], sampling_rate=sampling_rate, return_tensors="pt", padding=True
+    ... )
+    >>> with torch.no_grad():
+    ...     embeddings = model(**inputs).embeddings
+
+    >>> embeddings = torch.nn.functional.normalize(embeddings, dim=-1).cpu()
+
+    >>> # the resulting embeddings can be used for cosine similarity-based retrieval
+    >>> cosine_sim = torch.nn.CosineSimilarity(dim=-1)
+    >>> similarity = cosine_sim(embeddings[0], embeddings[1])
+    >>> threshold = 0.7  # the optimal threshold is dataset-dependent
+    >>> if similarity < threshold:
+    ...     print("Speakers are not the same!")
+    >>> round(similarity.item(), 2)
+    {expected_output}
+    ```
+"""
+
+PT_VISION_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+    >>> from datasets import load_dataset
+
+    >>> dataset = load_dataset("huggingface/cats-image")
+    >>> image = dataset["test"]["image"][0]
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = feature_extractor(image, return_tensors="pt")
+
+    >>> with torch.no_grad():
+    ...     outputs = model(**inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    >>> list(last_hidden_states.shape)
+    {expected_output}
+    ```
+"""
+
+PT_VISION_SEQ_CLASS_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import torch
+    >>> from datasets import load_dataset
+
+    >>> dataset = load_dataset("huggingface/cats-image")
+    >>> image = dataset["test"]["image"][0]
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = feature_extractor(image, return_tensors="pt")
+
+    >>> with torch.no_grad():
+    ...     logits = model(**inputs).logits
+
+    >>> # model predicts one of the 1000 ImageNet classes
+    >>> predicted_label = logits.argmax(-1).item()
+    >>> print(model.config.id2label[predicted_label])
+    {expected_output}
+    ```
+"""
+
+PT_SAMPLE_DOCSTRINGS = {
+    "SequenceClassification": PT_SEQUENCE_CLASSIFICATION_SAMPLE,
+    "QuestionAnswering": PT_QUESTION_ANSWERING_SAMPLE,
+    "TokenClassification": PT_TOKEN_CLASSIFICATION_SAMPLE,
+    "MultipleChoice": PT_MULTIPLE_CHOICE_SAMPLE,
+    "MaskedLM": PT_MASKED_LM_SAMPLE,
+    "LMHead": PT_CAUSAL_LM_SAMPLE,
+    "BaseModel": PT_BASE_MODEL_SAMPLE,
+    "SpeechBaseModel": PT_SPEECH_BASE_MODEL_SAMPLE,
+    "CTC": PT_SPEECH_CTC_SAMPLE,
+    "AudioClassification": PT_SPEECH_SEQ_CLASS_SAMPLE,
+    "AudioFrameClassification": PT_SPEECH_FRAME_CLASS_SAMPLE,
+    "AudioXVector": PT_SPEECH_XVECTOR_SAMPLE,
+    "VisionBaseModel": PT_VISION_BASE_MODEL_SAMPLE,
+    "ImageClassification": PT_VISION_SEQ_CLASS_SAMPLE,
+}
+
+TF_TOKEN_CLASSIFICATION_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer(
+    ...     "HuggingFace is a company based in Paris and New York", add_special_tokens=False, return_tensors="tf"
+    ... )
+
+    >>> logits = model(**inputs).logits
+    >>> predicted_token_class_ids = tf.math.argmax(logits, axis=-1)
+
+    >>> # Note that tokens are classified rather then input words which means that
+    >>> # there might be more predicted token classes than words.
+    >>> # Multiple token classes might account for the same word
+    >>> predicted_tokens_classes = [model.config.id2label[t] for t in predicted_token_class_ids[0].numpy().tolist()]
+    >>> predicted_tokens_classes
+    {expected_output}
+    ```
+
+    ```python
+    >>> labels = predicted_token_class_ids
+    >>> loss = tf.math.reduce_mean(model(**inputs, labels=labels).loss)
+    >>> round(float(loss), 2)
+    {expected_loss}
+    ```
+"""
+
+TF_QUESTION_ANSWERING_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+
+    >>> inputs = tokenizer(question, text, return_tensors="tf")
+    >>> outputs = model(**inputs)
+
+    >>> answer_start_index = int(tf.math.argmax(outputs.start_logits, axis=-1)[0])
+    >>> answer_end_index = int(tf.math.argmax(outputs.end_logits, axis=-1)[0])
+
+    >>> predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1]
+    >>> tokenizer.decode(predict_answer_tokens)
+    {expected_output}
+    ```
+
+    ```python
+    >>> # target is "nice puppet"
+    >>> target_start_index = tf.constant([{qa_target_start_index}])
+    >>> target_end_index = tf.constant([{qa_target_end_index}])
+
+    >>> outputs = model(**inputs, start_positions=target_start_index, end_positions=target_end_index)
+    >>> loss = tf.math.reduce_mean(outputs.loss)
+    >>> round(float(loss), 2)
+    {expected_loss}
+    ```
+"""
+
+TF_SEQUENCE_CLASSIFICATION_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf")
+
+    >>> logits = model(**inputs).logits
+
+    >>> predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0])
+    >>> model.config.id2label[predicted_class_id]
+    {expected_output}
+    ```
+
+    ```python
+    >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
+    >>> num_labels = len(model.config.id2label)
+    >>> model = {model_class}.from_pretrained("{checkpoint}", num_labels=num_labels)
+
+    >>> labels = tf.constant(1)
+    >>> loss = model(**inputs, labels=labels).loss
+    >>> round(float(loss), 2)
+    {expected_loss}
+    ```
+"""
+
+TF_MASKED_LM_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("The capital of France is {mask}.", return_tensors="tf")
+    >>> logits = model(**inputs).logits
+
+    >>> # retrieve index of {mask}
+    >>> mask_token_index = tf.where((inputs.input_ids == tokenizer.mask_token_id)[0])
+    >>> selected_logits = tf.gather_nd(logits[0], indices=mask_token_index)
+
+    >>> predicted_token_id = tf.math.argmax(selected_logits, axis=-1)
+    >>> tokenizer.decode(predicted_token_id)
+    {expected_output}
+    ```
+
+    ```python
+    >>> labels = tokenizer("The capital of France is Paris.", return_tensors="tf")["input_ids"]
+    >>> # mask labels of non-{mask} tokens
+    >>> labels = tf.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100)
+
+    >>> outputs = model(**inputs, labels=labels)
+    >>> round(float(outputs.loss), 2)
+    {expected_loss}
+    ```
+"""
+
+TF_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf")
+    >>> outputs = model(inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    ```
+"""
+
+TF_MULTIPLE_CHOICE_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
+    >>> choice0 = "It is eaten with a fork and a knife."
+    >>> choice1 = "It is eaten while held in the hand."
+
+    >>> encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="tf", padding=True)
+    >>> inputs = {{k: tf.expand_dims(v, 0) for k, v in encoding.items()}}
+    >>> outputs = model(inputs)  # batch size is 1
+
+    >>> # the linear classifier still needs to be trained
+    >>> logits = outputs.logits
+    ```
+"""
+
+TF_CAUSAL_LM_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf")
+    >>> outputs = model(inputs)
+    >>> logits = outputs.logits
+    ```
+"""
+
+TF_SPEECH_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> processor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="tf")
+    >>> outputs = model(**inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    >>> list(last_hidden_states.shape)
+    {expected_output}
+    ```
+"""
+
+TF_SPEECH_CTC_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+    >>> import tensorflow as tf
+
+    >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
+    >>> dataset = dataset.sort("id")
+    >>> sampling_rate = dataset.features["audio"].sampling_rate
+
+    >>> processor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> # audio file is decoded on the fly
+    >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="tf")
+    >>> logits = model(**inputs).logits
+    >>> predicted_ids = tf.math.argmax(logits, axis=-1)
+
+    >>> # transcribe speech
+    >>> transcription = processor.batch_decode(predicted_ids)
+    >>> transcription[0]
+    {expected_output}
+    ```
+
+    ```python
+    >>> inputs["labels"] = processor(text=dataset[0]["text"], return_tensors="tf").input_ids
+
+    >>> # compute loss
+    >>> loss = model(**inputs).loss
+    >>> round(float(loss), 2)
+    {expected_loss}
+    ```
+"""
+
+TF_VISION_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> from datasets import load_dataset
+
+    >>> dataset = load_dataset("huggingface/cats-image")
+    >>> image = dataset["test"]["image"][0]
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = feature_extractor(image, return_tensors="tf")
+    >>> outputs = model(**inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    >>> list(last_hidden_states.shape)
+    {expected_output}
+    ```
+"""
+
+TF_VISION_SEQ_CLASS_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+    >>> import tensorflow as tf
+    >>> from datasets import load_dataset
+
+    >>> dataset = load_dataset("huggingface/cats-image")
+    >>> image = dataset["test"]["image"][0]
+
+    >>> feature_extractor = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = feature_extractor(image, return_tensors="tf")
+    >>> logits = model(**inputs).logits
+
+    >>> # model predicts one of the 1000 ImageNet classes
+    >>> predicted_label = int(tf.math.argmax(logits, axis=-1))
+    >>> print(model.config.id2label[predicted_label])
+    {expected_output}
+    ```
+"""
+
+TF_SAMPLE_DOCSTRINGS = {
+    "SequenceClassification": TF_SEQUENCE_CLASSIFICATION_SAMPLE,
+    "QuestionAnswering": TF_QUESTION_ANSWERING_SAMPLE,
+    "TokenClassification": TF_TOKEN_CLASSIFICATION_SAMPLE,
+    "MultipleChoice": TF_MULTIPLE_CHOICE_SAMPLE,
+    "MaskedLM": TF_MASKED_LM_SAMPLE,
+    "LMHead": TF_CAUSAL_LM_SAMPLE,
+    "BaseModel": TF_BASE_MODEL_SAMPLE,
+    "SpeechBaseModel": TF_SPEECH_BASE_MODEL_SAMPLE,
+    "CTC": TF_SPEECH_CTC_SAMPLE,
+    "VisionBaseModel": TF_VISION_BASE_MODEL_SAMPLE,
+    "ImageClassification": TF_VISION_SEQ_CLASS_SAMPLE,
+}
+
+FLAX_TOKEN_CLASSIFICATION_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+
+    >>> outputs = model(**inputs)
+    >>> logits = outputs.logits
+    ```
+"""
+
+FLAX_QUESTION_ANSWERING_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+    >>> inputs = tokenizer(question, text, return_tensors="jax")
+
+    >>> outputs = model(**inputs)
+    >>> start_scores = outputs.start_logits
+    >>> end_scores = outputs.end_logits
+    ```
+"""
+
+FLAX_SEQUENCE_CLASSIFICATION_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+
+    >>> outputs = model(**inputs)
+    >>> logits = outputs.logits
+    ```
+"""
+
+FLAX_MASKED_LM_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("The capital of France is {mask}.", return_tensors="jax")
+
+    >>> outputs = model(**inputs)
+    >>> logits = outputs.logits
+    ```
+"""
+
+FLAX_BASE_MODEL_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax")
+    >>> outputs = model(**inputs)
+
+    >>> last_hidden_states = outputs.last_hidden_state
+    ```
+"""
+
+FLAX_MULTIPLE_CHOICE_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
+    >>> choice0 = "It is eaten with a fork and a knife."
+    >>> choice1 = "It is eaten while held in the hand."
+
+    >>> encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="jax", padding=True)
+    >>> outputs = model(**{{k: v[None, :] for k, v in encoding.items()}})
+
+    >>> logits = outputs.logits
+    ```
+"""
+
+FLAX_CAUSAL_LM_SAMPLE = r"""
+    Example:
+
+    ```python
+    >>> from transformers import {processor_class}, {model_class}
+
+    >>> tokenizer = {processor_class}.from_pretrained("{checkpoint}")
+    >>> model = {model_class}.from_pretrained("{checkpoint}")
+
+    >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np")
+    >>> outputs = model(**inputs)
+
+    >>> # retrieve logts for next token
+    >>> next_token_logits = outputs.logits[:, -1]
+    ```
+"""
+
+FLAX_SAMPLE_DOCSTRINGS = {
+    "SequenceClassification": FLAX_SEQUENCE_CLASSIFICATION_SAMPLE,
+    "QuestionAnswering": FLAX_QUESTION_ANSWERING_SAMPLE,
+    "TokenClassification": FLAX_TOKEN_CLASSIFICATION_SAMPLE,
+    "MultipleChoice": FLAX_MULTIPLE_CHOICE_SAMPLE,
+    "MaskedLM": FLAX_MASKED_LM_SAMPLE,
+    "BaseModel": FLAX_BASE_MODEL_SAMPLE,
+    "LMHead": FLAX_CAUSAL_LM_SAMPLE,
+}
+
+
+def add_code_sample_docstrings(
+        *docstr,
+        processor_class=None,
+        checkpoint=None,
+        output_type=None,
+        config_class=None,
+        mask="[MASK]",
+        qa_target_start_index=14,
+        qa_target_end_index=15,
+        model_cls=None,
+        modality=None,
+        expected_output="",
+        expected_loss="",
+):
+    def docstring_decorator(fn):
+        # model_class defaults to function's class if not specified otherwise
+        model_class = fn.__qualname__.split(".")[0] if model_cls is None else model_cls
+
+        if model_class[:2] == "TF":
+            sample_docstrings = TF_SAMPLE_DOCSTRINGS
+        elif model_class[:4] == "Flax":
+            sample_docstrings = FLAX_SAMPLE_DOCSTRINGS
+        else:
+            sample_docstrings = PT_SAMPLE_DOCSTRINGS
+
+        # putting all kwargs for docstrings in a dict to be used
+        # with the `.format(**doc_kwargs)`. Note that string might
+        # be formatted with non-existing keys, which is fine.
+        doc_kwargs = dict(
+            model_class=model_class,
+            processor_class=processor_class,
+            checkpoint=checkpoint,
+            mask=mask,
+            qa_target_start_index=qa_target_start_index,
+            qa_target_end_index=qa_target_end_index,
+            expected_output=expected_output,
+            expected_loss=expected_loss,
+        )
+
+        if "SequenceClassification" in model_class and modality == "audio":
+            code_sample = sample_docstrings["AudioClassification"]
+        elif "SequenceClassification" in model_class:
+            code_sample = sample_docstrings["SequenceClassification"]
+        elif "QuestionAnswering" in model_class:
+            code_sample = sample_docstrings["QuestionAnswering"]
+        elif "TokenClassification" in model_class:
+            code_sample = sample_docstrings["TokenClassification"]
+        elif "MultipleChoice" in model_class:
+            code_sample = sample_docstrings["MultipleChoice"]
+        elif "MaskedLM" in model_class or model_class in ["FlaubertWithLMHeadModel", "XLMWithLMHeadModel"]:
+            code_sample = sample_docstrings["MaskedLM"]
+        elif "LMHead" in model_class or "CausalLM" in model_class:
+            code_sample = sample_docstrings["LMHead"]
+        elif "CTC" in model_class:
+            code_sample = sample_docstrings["CTC"]
+        elif "AudioFrameClassification" in model_class:
+            code_sample = sample_docstrings["AudioFrameClassification"]
+        elif "XVector" in model_class and modality == "audio":
+            code_sample = sample_docstrings["AudioXVector"]
+        elif "Model" in model_class and modality == "audio":
+            code_sample = sample_docstrings["SpeechBaseModel"]
+        elif "Model" in model_class and modality == "vision":
+            code_sample = sample_docstrings["VisionBaseModel"]
+        elif "Model" in model_class or "Encoder" in model_class:
+            code_sample = sample_docstrings["BaseModel"]
+        elif "ImageClassification" in model_class:
+            code_sample = sample_docstrings["ImageClassification"]
+        else:
+            raise ValueError(f"Docstring can't be built for model {model_class}")
+
+        func_doc = (fn.__doc__ or "") + "".join(docstr)
+        output_doc = "" if output_type is None else _prepare_output_docstrings(output_type, config_class)
+        built_doc = code_sample.format(**doc_kwargs)
+        fn.__doc__ = func_doc + output_doc + built_doc
+        return fn
+
+    return docstring_decorator
+
+
+def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear:
+    """
+    Prune a linear layer to keep only entries in index.
+
+    Used to remove heads.
+
+    Args:
+        layer (`torch.nn.Linear`): The layer to prune.
+        index (`torch.LongTensor`): The indices to keep in the layer.
+        dim (`int`, *optional*, defaults to 0): The dimension on which to keep the indices.
+
+    Returns:
+        `torch.nn.Linear`: The pruned layer as a new layer with `requires_grad=True`.
+    """
+    index = index.to(layer.weight.device)
+    W = layer.weight.index_select(dim, index).clone().detach()
+    if layer.bias is not None:
+        if dim == 1:
+            b = layer.bias.clone().detach()
+        else:
+            b = layer.bias[index].clone().detach()
+    new_size = list(layer.weight.size())
+    new_size[dim] = len(index)
+    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
+    new_layer.weight.requires_grad = False
+    new_layer.weight.copy_(W.contiguous())
+    new_layer.weight.requires_grad = True
+    if layer.bias is not None:
+        new_layer.bias.requires_grad = False
+        new_layer.bias.copy_(b.contiguous())
+        new_layer.bias.requires_grad = True
+    return new_layer
+
+
+def apply_chunking_to_forward(
+        forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors
+) -> torch.Tensor:
+    """
+    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
+    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
+
+    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
+    applying `forward_fn` to `input_tensors`.
+
+    Args:
+        forward_fn (`Callable[..., torch.Tensor]`):
+            The forward function of the model.
+        chunk_size (`int`):
+            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
+        chunk_dim (`int`):
+            The dimension over which the `input_tensors` should be chunked.
+        input_tensors (`Tuple[torch.Tensor]`):
+            The input tensors of `forward_fn` which will be chunked
+
+    Returns:
+        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
+
+
+    Examples:
+
+    ```python
+    # rename the usual forward() fn to forward_chunk()
+    def forward_chunk(self, hidden_states):
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+    # implement a chunked forward function
+    def forward(self, hidden_states):
+        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
+    ```"""
+
+    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"
+
+    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
+    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
+    if num_args_in_forward_chunk_fn != len(input_tensors):
+        raise ValueError(
+            f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
+            "tensors are given"
+        )
+
+    if chunk_size > 0:
+        tensor_shape = input_tensors[0].shape[chunk_dim]
+        for input_tensor in input_tensors:
+            if input_tensor.shape[chunk_dim] != tensor_shape:
+                raise ValueError(
+                    f"All input tenors have to be of the same shape: {tensor_shape}, "
+                    f"found shape {input_tensor.shape[chunk_dim]}"
+                )
+
+        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
+            raise ValueError(
+                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
+                f"size {chunk_size}"
+            )
+
+        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size
+
+        # chunk input tensor into tuples
+        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
+        # apply forward fn to every tuple
+        output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks))
+        # concatenate output at same dimension
+        return torch.cat(output_chunks, dim=chunk_dim)
+
+    return forward_fn(*input_tensors)
+
+
+def find_pruneable_heads_and_indices(
+        heads: List[int], n_heads: int, head_size: int, already_pruned_heads: Set[int]
+) -> Tuple[Set[int], torch.LongTensor]:
+    """
+    Finds the heads and their indices taking `already_pruned_heads` into account.
+
+    Args:
+        heads (`List[int]`): List of the indices of heads to prune.
+        n_heads (`int`): The number of heads in the model.
+        head_size (`int`): The size of each head.
+        already_pruned_heads (`Set[int]`): A set of already pruned heads.
+
+    Returns:
+        `Tuple[Set[int], torch.LongTensor]`: A tuple with the remaining heads and their corresponding indices.
+    """
+    mask = torch.ones(n_heads, head_size)
+    heads = set(heads) - already_pruned_heads  # Convert to set and remove already pruned heads
+    for head in heads:
+        # Compute how many pruned heads are before the head and move the index accordingly
+        head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
+        mask[head] = 0
+    mask = mask.view(-1).contiguous().eq(1)
+    index: torch.LongTensor = torch.arange(len(mask))[mask].long()
+    return heads, index

chn_2_code.pkl

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

config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "TAAS"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_TAAS.TAASConfig",
+    "AutoModel": "modeling_TAAS.TAAS",
+    "AutoModelForMaskedLM": "modeling_TAAS.TAAS"
+  },
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 2048,
+  "model_type": "TAAS",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "task_type_vocab_size": 3,
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.1",
+  "type_vocab_size": 4,
+  "use_cache": true,
+  "use_task_id": true,
+  "vocab_size": 40000
+}

configuration_TAAS.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class TAASConfig(PretrainedConfig):
+    model_type = "Stellar"
+
+    def __init__(
+        self,
+        hidd_dropout=0.1,
+        intermediate_size=3072,
+        initialize_range=0.02,
+        max_pos_embeddings=2048,
+        hidd_act="gelu",
+        attention_dropout=0.1,
+        using_task_id=True,
+        vocabulary_size=40000,
+        hidd_size=768,
+        num_hidd_layers=12,
+        layer_norm_rate=1e-05,
+        num_atten_heads=12,
+        pad_token_id=0,
+        task_vocab_size=3,
+        classifier_drop=None,
+        pos_embedding="absolute",
+        use_cache=True,
+        vocab_size=4,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.vocab_size = vocabulary_size
+        self.max_position_embeddings = max_pos_embeddings
+        self.type_vocab_size = vocab_size
+        self.use_task_id = using_task_id
+        self.layer_norm_eps = layer_norm_rate
+        self.position_embedding_type = pos_embedding
+        self.num_attention_heads = num_atten_heads
+        self.hidden_size = hidd_size
+        self.attention_probs_dropout_prob = attention_dropout
+        self.initializer_range = initialize_range
+        self.hidden_act = hidd_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidd_dropout
+        self.use_cache = use_cache
+        self.classifier_dropout = classifier_drop
+        self.num_hidden_layers = num_hidd_layers
+        self.task_type_vocab_size = task_vocab_size
+
+        
+

graphormer.py

@@ -0,0 +1,393 @@
+from copy import deepcopy
+from torch.nn.init import xavier_uniform_
+import torch.nn.functional as F
+from torch.nn import Parameter
+from torch.nn.init import normal_
+import torch.utils.checkpoint
+from torch import Tensor, device
+from TAAS_utils import *
+from transformers.modeling_utils import ModuleUtilsMixin
+from fairseq import utils
+from fairseq.models import (
+    FairseqEncoder,
+    FairseqEncoderModel,
+    register_model,
+    register_model_architecture,
+)
+from fairseq.modules import (
+    LayerNorm,
+)
+from fairseq.utils import safe_hasattr
+
+def init_params(module, n_layers):
+    if isinstance(module, nn.Linear):
+        module.weight.data.normal_(mean=0.0, std=0.02 / math.sqrt(n_layers))
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+
+
+@torch.jit.script
+def softmax_dropout(input, dropout_prob: float, is_training: bool):
+    return F.dropout(F.softmax(input, -1), dropout_prob, is_training)
+
+
+class SelfMultiheadAttention(nn.Module):
+    def __init__(
+            self,
+            embed_dim,
+            num_heads,
+            dropout=0.0,
+            bias=True,
+            scaling_factor=1,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+
+        self.head_dim = embed_dim // num_heads
+        assert (self.head_dim * num_heads == self.embed_dim), "embed_dim must be divisible by num_heads"
+        self.scaling = (self.head_dim * scaling_factor) ** -0.5
+
+        self.linear_q = nn.Linear(self.embed_dim, self.num_heads * self.head_dim)
+        self.linear_k = nn.Linear(self.embed_dim, self.num_heads * self.head_dim)
+        self.linear_v = nn.Linear(self.embed_dim, self.num_heads * self.head_dim)
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=bias)
+
+    def forward(
+            self,
+            query: Tensor,
+            attn_bias: Tensor = None,
+    ) -> Tensor:
+        n_graph, n_node, embed_dim = query.size()
+        # q, k, v = self.in_proj(query).chunk(3, dim=-1)
+
+        _shape = (-1, n_graph * self.num_heads, self.head_dim)
+        q = self.linear_q(query).contiguous().view(n_graph, -1, self.num_heads, self.head_dim).transpose(1, 2) * self.scaling
+        k = self.linear_k(query).contiguous().view(n_graph, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.linear_v(query).contiguous().view(n_graph, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn_weights = torch.matmul(q, k.transpose(2, 3))
+        attn_weights = attn_weights + attn_bias
+        attn_probs = softmax_dropout(attn_weights, self.dropout, self.training)
+
+        attn = torch.matmul(attn_probs, v)
+        attn = attn.transpose(1, 2).contiguous().view(n_graph, -1, embed_dim)
+        attn = self.out_proj(attn)
+        return attn
+
+
+class Graphormer3DEncoderLayer(nn.Module):
+    """
+    Implements a Graphormer-3D Encoder Layer.
+    """
+
+    def __init__(
+            self,
+            embedding_dim: int = 768,
+            ffn_embedding_dim: int = 3072,
+            num_attention_heads: int = 8,
+            dropout: float = 0.1,
+            attention_dropout: float = 0.1,
+            activation_dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+
+        # Initialize parameters
+        self.embedding_dim = embedding_dim
+        self.num_attention_heads = num_attention_heads
+        self.attention_dropout = attention_dropout
+
+        self.dropout = dropout
+        self.activation_dropout = activation_dropout
+
+        self.self_attn = SelfMultiheadAttention(self.embedding_dim, num_attention_heads, dropout=attention_dropout)
+        # layer norm associated with the self attention layer
+        self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim)
+        self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
+        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embedding_dim)
+
+    def forward(self, x: Tensor, attn_bias: Tensor = None):
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        x = self.self_attn(query=x, attn_bias=attn_bias)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+
+        residual = x
+        x = self.final_layer_norm(x)
+        x = F.gelu(self.fc1(x))
+        x = F.dropout(x, p=self.activation_dropout, training=self.training)
+        x = self.fc2(x)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = residual + x
+        return x
+
+
+from fairseq.models import (
+    BaseFairseqModel,
+    register_model,
+    register_model_architecture,
+)
+
+
+class Graphormer3D(BaseFairseqModel):
+    def __init__(self):
+        super().__init__()
+        self.atom_types = 64
+        self.edge_types = 64 * 64
+        self.embed_dim = 768
+        self.layer_nums = 12
+        self.ffn_embed_dim = 768
+        self.blocks = 4
+        self.attention_heads = 48
+        self.input_dropout = 0.0
+        self.dropout = 0.1
+        self.attention_dropout = 0.1
+        self.activation_dropout = 0.0
+        self.node_loss_weight = 15
+        self.min_node_loss_weight = 1
+        self.eng_loss_weight = 1
+        self.num_kernel = 128
+        self.atom_encoder = nn.Embedding(self.atom_types, self.embed_dim, padding_idx=0)
+        self.edge_embedding = nn.Embedding(32, self.attention_heads, padding_idx=0)
+        self.input_dropout = nn.Dropout(0.1)
+        self.layers = nn.ModuleList(
+            [
+                Graphormer3DEncoderLayer(
+                    self.embed_dim,
+                    self.ffn_embed_dim,
+                    num_attention_heads=self.attention_heads,
+                    dropout=self.dropout,
+                    attention_dropout=self.attention_dropout,
+                    activation_dropout=self.activation_dropout,
+                )
+                for _ in range(self.layer_nums)
+            ]
+        )
+        self.atom_encoder = nn.Embedding(512 * 9 + 1, self.embed_dim, padding_idx=0)
+        self.edge_encoder = nn.Embedding(512 * 3 + 1, self.attention_heads, padding_idx=0)
+        self.edge_type = 'multi_hop'
+        if self.edge_type == 'multi_hop':
+            self.edge_dis_encoder = nn.Embedding(16 * self.attention_heads * self.attention_heads, 1)
+        self.spatial_pos_encoder = nn.Embedding(512, self.attention_heads, padding_idx=0)
+        self.in_degree_encoder = nn.Embedding(512, self.embed_dim, padding_idx=0)
+        self.out_degree_encoder = nn.Embedding(512, self.embed_dim, padding_idx=0)
+        self.node_position_ids_encoder = nn.Embedding(10, self.embed_dim, padding_idx=0)
+
+        self.final_ln: Callable[[Tensor], Tensor] = nn.LayerNorm(self.embed_dim)
+
+        self.engergy_proj: Callable[[Tensor], Tensor] = NonLinear(self.embed_dim, 1)
+        self.energe_agg_factor: Callable[[Tensor], Tensor] = nn.Embedding(3, 1)
+        nn.init.normal_(self.energe_agg_factor.weight, 0, 0.01)
+
+        self.graph_token = nn.Embedding(1, 768)
+        self.graph_token_virtual_distance = nn.Embedding(1, self.attention_heads)
+
+        K = self.num_kernel
+
+        self.gbf: Callable[[Tensor, Tensor], Tensor] = GaussianLayer(K, self.edge_types)
+        self.bias_proj: Callable[[Tensor], Tensor] = NonLinear(K, self.attention_heads)
+        self.edge_proj: Callable[[Tensor], Tensor] = nn.Linear(K, self.embed_dim)
+        self.node_proc: Callable[[Tensor, Tensor, Tensor], Tensor] = NodeTaskHead(self.embed_dim, self.attention_heads)
+
+    def forward(self, node_feature, spatial_pos, in_degree, out_degree, edge_type_matrix, edge_input, node_position_ids):
+        """
+        attn_bias：图中节点对之间的最短路径距离超过最短路径限制最大距离(spatial_pos_max)的位置为-inf，其余位置为0，形状为(n_graph, n_node+1, n_node+1)
+        spatial_pos：图中节点对之间的最短路径长度，形状为(n_graph, n_node, n_node)
+        x：图中节点的特征，形状为(n_graph, n_node, n_node_features)
+        in_degree：图中节点的入度，形状为(n_graph, n_node)
+        out_degree：图中节点的出度，形状为(n_graph, n_node)
+        edge_input：图中节点对之间的最短路径(限制最短路径最大跳数为multi_hop_max_dist)上的边的特征，形状为(n_graph, n_node, n_node, multi_hop_max_dist, n_edge_features)
+        attn_edge_type：图的边特征，形状为(n_graph, n_node, n_node, n_edge_features)
+        :param batch_data:
+        :return:
+        """
+        # attn_bias, spatial_pos, x = batch_data.attn_bias, batch_data.spatial_pos, batch_data.x
+        # in_degree, out_degree = batch_data.in_degree, batch_data.out_degree
+        # edge_input, attn_edge_type = batch_data.edge_input, batch_data.attn_edge_type
+        # graph_attn_bias
+        attn_edge_type = self.edge_embedding(edge_type_matrix)
+        edge_input = self.edge_embedding(edge_input)#.mean(-2)
+        # 添加虚拟节点表示全图特征表示，之后按照图中正常节点处理
+        n_graph, n_node = node_feature.size()[:2]
+        # graph_attn_bias = attn_bias.clone()
+        # graph_attn_bias = graph_attn_bias.unsqueeze(1).repeat(1, self.attention_heads, 1, 1)  # [n_graph, n_head, n_node+1, n_node+1]
+
+        # spatial pos
+        # 空间编码,节点之间最短路径长度对应的可学习标量
+        # [n_graph, n_node, n_node, n_head] -> [n_graph, n_head, n_node, n_node]
+        spatial_pos_bias = self.spatial_pos_encoder(spatial_pos).permute(0, 3, 1, 2)
+        # graph_attn_bias[:, :, 1:, 1:] = graph_attn_bias[:, :, 1:, 1:] + spatial_pos_bias
+        # graph_attn_bias = spatial_pos_bias
+        # reset spatial pos here
+        # 所有节点都和虚拟节点直接有边相连，则所有节点和虚拟节点之间的最短路径长度为1
+        # t = self.graph_token_virtual_distance.weight.view(1, self.attention_heads, 1)
+        # graph_attn_bias[:, :, 1:, 0] = graph_attn_bias[:, :, 1:, 0] + t
+        # graph_attn_bias[:, :, 0, :] = graph_attn_bias[:, :, 0, :] + t
+        # edge feature
+        # 每个节点对沿最短路径计算边特征和可学习嵌入点积的平均值，并作为偏置项添加到注意模块中
+        if self.edge_type == 'multi_hop':
+            spatial_pos_ = spatial_pos.clone()
+            spatial_pos_[spatial_pos_ == 0] = 1  # set pad to 1
+            # set 1 to 1, x > 1 to x - 1
+            spatial_pos_ = torch.where(spatial_pos_ > 1, spatial_pos_ - 1, spatial_pos_)
+            # if self.multi_hop_max_dist > 0:
+            #     spatial_pos_ = spatial_pos_.clamp(0, self.multi_hop_max_dist)
+            #     edge_input = edge_input[:, :, :, :self.multi_hop_max_dist, :]
+            # [n_graph, n_node, n_node, max_dist, n_head]
+            # edge_input = self.edge_encoder(edge_input).mean(-2)
+            max_dist = edge_input.size(-2)
+            edge_input_flat = edge_input.permute(3, 0, 1, 2, 4).reshape(max_dist, -1, self.attention_heads)
+            edge_input_flat = torch.bmm(edge_input_flat, self.edge_dis_encoder.weight.reshape(-1, self.attention_heads, self.attention_heads)[:max_dist, :, :])
+            edge_input = edge_input_flat.reshape(max_dist, n_graph, n_node, n_node, self.attention_heads).permute(1, 2, 3, 0, 4)
+            edge_input = (edge_input.sum(-2) / (spatial_pos_.float().unsqueeze(-1))).permute(0, 3, 1, 2)
+        else:
+            # [n_graph, n_node, n_node, n_head] -> [n_graph, n_head, n_node, n_node]
+            edge_input = self.edge_encoder(attn_edge_type).mean(-2).permute(0, 3, 1, 2)
+
+        # graph_attn_bias[:, :, 1:, 1:] = graph_attn_bias[:, :, 1:, 1:] + edge_input
+        graph_attn_bias = spatial_pos_bias + edge_input
+        # graph_attn_bias = graph_attn_bias + attn_bias.unsqueeze(1)  # reset
+        # graph_attn_bias = graph_attn_bias.contiguous().view(-1, 6, 6)
+        # node feauture + graph token
+        # node_feature = x # self.atom_encoder(x).sum(dim=-2)  # [n_graph, n_node, n_hidden]
+        # if self.flag and perturb is not None:
+        #     node_feature += perturb
+
+        node_position_embedding = self.node_position_ids_encoder(node_position_ids)
+        node_position_embedding = node_position_embedding.contiguous().view(n_graph, n_node, self.embed_dim)
+        # print(node_position_embedding.shape)
+        # 根据节点的入度、出度为每个节点分配两个实值嵌入向量，添加到节点特征中作为输入
+        node_feature = node_feature + self.in_degree_encoder(in_degree) + \
+                       self.out_degree_encoder(out_degree) + node_position_embedding
+        # print(node_feature.shape)
+        # graph_token_feature = self.graph_token.weight.unsqueeze(0).repeat(n_graph, 1, 1)
+        # graph_node_feature = torch.cat([graph_token_feature, node_feature], dim=1)
+
+        # transfomrer encoder
+        output = self.input_dropout(node_feature)#.permute(1, 0, 2)
+        for enc_layer in self.layers:
+            output = enc_layer(output, graph_attn_bias)
+        output = self.final_ln(output)
+
+        # output part
+        # 整个图的表示是最后一层虚拟节点的特征
+        # if self.dataset_name == 'PCQM4M-LSC':
+        #     # get whole graph rep
+        #     output = self.out_proj(output[:, 0, :])
+        # else:
+        #     output = self.downstream_out_proj(output[:, 0, :])
+        # print(output.shape)
+        return output
+
+
+@torch.jit.script
+def gaussian(x, mean, std):
+    pi = 3.14159
+    a = (2 * pi) ** 0.5
+    return torch.exp(-0.5 * (((x - mean) / std) ** 2)) / (a * std)
+
+
+class GaussianLayer(nn.Module):
+    def __init__(self, K=128, edge_types=1024):
+        super().__init__()
+        self.K = K
+        self.means = nn.Embedding(1, K)
+        self.stds = nn.Embedding(1, K)
+        self.mul = nn.Embedding(edge_types, 1)
+        self.bias = nn.Embedding(edge_types, 1)
+        nn.init.uniform_(self.means.weight, 0, 3)
+        nn.init.uniform_(self.stds.weight, 0, 3)
+        nn.init.constant_(self.bias.weight, 0)
+        nn.init.constant_(self.mul.weight, 1)
+
+    def forward(self, x, edge_types):
+        mul = self.mul(edge_types)
+        bias = self.bias(edge_types)
+        x = mul * x.unsqueeze(-1) + bias
+        x = x.expand(-1, -1, -1, self.K)
+        mean = self.means.weight.float().view(-1)
+        std = self.stds.weight.float().view(-1).abs() + 1e-5
+        return gaussian(x.float(), mean, std).type_as(self.means.weight)
+
+
+class RBF(nn.Module):
+    def __init__(self, K, edge_types):
+        super().__init__()
+        self.K = K
+        self.means = nn.parameter.Parameter(torch.empty(K))
+        self.temps = nn.parameter.Parameter(torch.empty(K))
+        self.mul: Callable[..., Tensor] = nn.Embedding(edge_types, 1)
+        self.bias: Callable[..., Tensor] = nn.Embedding(edge_types, 1)
+        nn.init.uniform_(self.means, 0, 3)
+        nn.init.uniform_(self.temps, 0.1, 10)
+        nn.init.constant_(self.bias.weight, 0)
+        nn.init.constant_(self.mul.weight, 1)
+
+    def forward(self, x: Tensor, edge_types):
+        mul = self.mul(edge_types)
+        bias = self.bias(edge_types)
+        x = mul * x.unsqueeze(-1) + bias
+        mean = self.means.float()
+        temp = self.temps.float().abs()
+        return ((x - mean).square() * (-temp)).exp().type_as(self.means)
+
+
+class NonLinear(nn.Module):
+    def __init__(self, input, output_size, hidden=None):
+        super(NonLinear, self).__init__()
+        if hidden is None:
+            hidden = input
+        self.layer1 = nn.Linear(input, hidden)
+        self.layer2 = nn.Linear(hidden, output_size)
+
+    def forward(self, x):
+        x = F.gelu(self.layer1(x))
+        x = self.layer2(x)
+        return x
+
+
+class NodeTaskHead(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.q_proj: Callable[[Tensor], Tensor] = nn.Linear(embed_dim, embed_dim)
+        self.k_proj: Callable[[Tensor], Tensor] = nn.Linear(embed_dim, embed_dim)
+        self.v_proj: Callable[[Tensor], Tensor] = nn.Linear(embed_dim, embed_dim)
+        self.num_heads = num_heads
+        self.scaling = (embed_dim // num_heads) ** -0.5
+        self.force_proj1: Callable[[Tensor], Tensor] = nn.Linear(embed_dim, 1)
+        self.force_proj2: Callable[[Tensor], Tensor] = nn.Linear(embed_dim, 1)
+        self.force_proj3: Callable[[Tensor], Tensor] = nn.Linear(embed_dim, 1)
+
+    def forward(
+            self,
+            query: Tensor,
+            attn_bias: Tensor,
+            delta_pos: Tensor,
+    ) -> Tensor:
+        bsz, n_node, _ = query.size()
+        q = (self.q_proj(query).view(bsz, n_node, self.num_heads, -1).transpose(1, 2) * self.scaling)
+        k = self.k_proj(query).view(bsz, n_node, self.num_heads, -1).transpose(1, 2)
+        v = self.v_proj(query).view(bsz, n_node, self.num_heads, -1).transpose(1, 2)
+        attn = q @ k.transpose(-1, -2)  # [bsz, head, n, n]
+        attn_probs = softmax_dropout(attn.view(-1, n_node, n_node) + attn_bias, 0.1, self.training).view(bsz, self.num_heads, n_node, n_node)
+        rot_attn_probs = attn_probs.unsqueeze(-1) * delta_pos.unsqueeze(1).type_as(attn_probs)  # [bsz, head, n, n, 3]
+        rot_attn_probs = rot_attn_probs.permute(0, 1, 4, 2, 3)
+        x = rot_attn_probs @ v.unsqueeze(2)  # [bsz, head , 3, n, d]
+        x = x.permute(0, 3, 2, 1, 4).contiguous().view(bsz, n_node, 3, -1)
+        f1 = self.force_proj1(x[:, :, 0, :]).view(bsz, n_node, 1)
+        f2 = self.force_proj2(x[:, :, 1, :]).view(bsz, n_node, 1)
+        f3 = self.force_proj3(x[:, :, 2, :]).view(bsz, n_node, 1)
+        cur_force = torch.cat([f1, f2, f3], dim=-1).float()
+        return cur_force
+

htc_loss.py

@@ -0,0 +1,145 @@
+#! python3
+# -*- encoding: utf-8 -*-
+
+import torch
+import torch.nn.functional as F
+import pandas as pd
+import sys
+import os
+
+from transformers.utils.hub import cached_file
+
+resolved_module_file = cached_file(
+    'Cainiao-AI/TAAS', 
+    'htc_mask_dict_old.pkl'
+)
+
+htc_weights = [0.067, 0.133, 0.2, 0.267, 0.333]
+htc_mask_dict = pd.read_pickle(resolved_module_file)
+import numpy as np
+import operator
+def calculate_multi_htc_acc_batch(predicted_htc, y, sequence_len = 6):
+    acc_cnt = np.array([0, 0, 0, 0, 0])
+    y = y.view(-1, sequence_len, 5).tolist()
+    predicted = np.array(predicted_htc).reshape(-1, sequence_len, 5).tolist()
+    batch_size = len(y)
+    total_cnt = np.array([0, 0, 0, 0, 0])
+    for batch_i in range(batch_size):
+        for index, s2 in enumerate(y[batch_i]):
+            for c, i in enumerate(range(5)):
+                y_l10 = y[batch_i][index][:i+1]
+                p_l10 = predicted[batch_i][index][:i+1]
+                if -100 in y_l10:
+                    break
+                
+                if operator.eq(y_l10, p_l10):
+                    acc_cnt[c] += 1
+                total_cnt[c] += 1
+        
+    return acc_cnt, total_cnt
+    
+
+class HTCLoss(torch.nn.Module):
+    def __init__(self, device, reduction='mean', using_htc = True):
+        super(HTCLoss, self).__init__()
+        self.reduction = reduction
+        self.htc_weights = htc_weights
+        self.device = device
+        self.using_htc = using_htc
+        self.htc_mask_dict = htc_mask_dict
+        for key, value in self.htc_mask_dict.items():
+            # self.htc_mask_dict[key] = torch.tensor(value).to(self.device)
+            self.htc_mask_dict[key] = torch.tensor(value).clone().detach().to(self.device)
+
+    def forward(self, logits, target):  # [bs,num_class]  CE=q*-log(p), q*log(1-p),p=softmax(logits)
+        # target相关变量都在cuda上
+        target = target.reshape(-1, 1)
+        target_mask = target != -100
+        target_mask = target_mask.squeeze()
+        target_mask_idx = torch.where(target == -100)
+        target_new = target.clone()
+        target_new[target_mask_idx] = 0
+        predict_res = []
+        if not self.using_htc:
+            log_pro = -1.0 * F.log_softmax(logits, dim=1)
+            # one_hot = torch.zeros(logits.shape[0], logits.shape[1]).to(self.device)  # .cuda()
+            # one_hot = one_hot.scatter_(1, target_new, 1)
+            # loss = torch.mul(log_pro, one_hot).sum(dim=1)
+            # loss = loss*target_mask
+        else:
+            # _, predicted = torch.max(logits[:, :32], 1)
+            logits_reshaped = logits.clone() 
+            logits_reshaped = logits_reshaped.reshape(-1, 5, 100) 
+            _, aa_predicted = torch.max(logits_reshaped[:,0,1:32], 1) 
+            aa_predicted += 1
+            logits_new = -5 * torch.ones_like(logits_reshaped).to(self.device)
+            logits_new[:,0,1:32] = logits_reshaped[:,0,1:32]
+            for sample_idx, aa in enumerate(aa_predicted):
+                bb_idx = htc_mask_dict['{:02d}'.format(aa)]
+                _, bb_idy = torch.max(logits_reshaped[sample_idx,1,bb_idx], 0)
+                bb = bb_idx[bb_idy]
+                logits_new[sample_idx,1,bb_idx] = logits_reshaped[sample_idx,1,bb_idx]
+                cc_idx = htc_mask_dict['{:02d}{:02d}'.format(aa, bb)]
+                _, cc_idy = torch.max(logits_reshaped[sample_idx,2,cc_idx], 0)
+                logits_new[sample_idx,2,cc_idx] = logits_reshaped[sample_idx,2,cc_idx]
+                cc = cc_idx[cc_idy]
+                d_idx = htc_mask_dict['{:02d}{:02d}{:02d}'.format(aa, bb, cc)]
+                _, d_idy = torch.max(logits_reshaped[sample_idx,3,d_idx], 0)
+                logits_new[sample_idx,3,d_idx] = logits_reshaped[sample_idx,3,d_idx]
+                d = d_idx[d_idy]
+                ee_idx = htc_mask_dict['{:02d}{:02d}{:02d}{:01d}'.format(aa, bb, cc, d)]
+                _, ee_idy = torch.max(logits_reshaped[sample_idx,4,ee_idx], 0)
+                logits_new[sample_idx,4,ee_idx] = logits_reshaped[sample_idx,4,ee_idx]
+                ee = ee_idx[ee_idy]
+                predict_res.extend([aa.item(), bb.item(), cc.item(), d.item(), ee.item()])
+            # predicted = predicted.reshape(-1, 5)
+            # aa = predicted[:, 0]
+            # aa = ['{:02d}'.format(i) for i in aa]
+            # bb_activate = [htc_mask_dict[i] for i in aa]
+            logits_new = logits_new.reshape(-1, 100)
+            log_pro = -1.0 * F.log_softmax(logits_new, dim=1)
+        logits = logits.contiguous().view(-1, 100)
+        one_hot = torch.zeros(logits.shape[0], logits.shape[1]).to(self.device)  # .cuda()
+        one_hot = one_hot.scatter_(1, target_new, 1)
+        loss = torch.mul(log_pro, one_hot).sum(dim=1)
+        loss = loss*target_mask
+        bs = int(loss.shape[0] / 5)
+        w_loss = []
+        for i in range(bs):
+            w_loss.extend(self.htc_weights)
+        w_loss = torch.FloatTensor(w_loss).to(self.device)
+        loss = loss.mul(w_loss) * 5
+        if self.reduction == 'mean':
+            loss = loss[torch.where(loss>0)].mean()
+        elif self.reduction == 'sum':
+            loss = loss[torch.where(loss>0)].sum()
+        return loss, predict_res
+
+    def get_htc_code(self, logits):  # [bs,num_class]  CE=q*-log(p), q*log(1-p),p=softmax(logits)
+        logits_reshaped = logits.clone() 
+        logits_reshaped = logits_reshaped.reshape(-1, 5, 100) 
+        _, aa_predicted = torch.max(logits_reshaped[:,0,1:32], 1) 
+        aa_predicted += 1
+        logits_new = -5 * torch.ones_like(logits_reshaped).to(self.device)
+        logits_new[:,0,1:32] = logits_reshaped[:,0,1:32]
+        predict_res = []
+        for sample_idx, aa in enumerate(aa_predicted):
+            bb_idx = htc_mask_dict['{:02d}'.format(aa)]
+            _, bb_idy = torch.max(logits_reshaped[sample_idx,1,bb_idx], 0)
+            bb = bb_idx[bb_idy]
+            logits_new[sample_idx,1,bb_idx] = logits_reshaped[sample_idx,1,bb_idx]
+            cc_idx = htc_mask_dict['{:02d}{:02d}'.format(aa, bb)]
+            _, cc_idy = torch.max(logits_reshaped[sample_idx,2,cc_idx], 0)
+            logits_new[sample_idx,2,cc_idx] = logits_reshaped[sample_idx,2,cc_idx]
+            cc = cc_idx[cc_idy]
+            d_idx = htc_mask_dict['{:02d}{:02d}{:02d}'.format(aa, bb, cc)]
+            _, d_idy = torch.max(logits_reshaped[sample_idx,3,d_idx], 0)
+            logits_new[sample_idx,3,d_idx] = logits_reshaped[sample_idx,3,d_idx]
+            d = d_idx[d_idy]
+            ee_idx = htc_mask_dict['{:02d}{:02d}{:02d}{:01d}'.format(aa, bb, cc, d)]
+            _, ee_idy = torch.max(logits_reshaped[sample_idx,4,ee_idx], 0)
+            logits_new[sample_idx,4,ee_idx] = logits_reshaped[sample_idx,4,ee_idx]
+            ee = ee_idx[ee_idy]
+            predict_res.extend([aa.item(), bb.item(), cc.item(), d.item(), ee.item()])
+        return predict_res
+    

htc_mask_dict_old.pkl

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

modeling_TAAS.py

@@ -0,0 +1,1034 @@
+#! python3
+# -*- encoding: utf-8 -*-
+
+from copy import deepcopy
+from torch.nn.init import xavier_uniform_
+import torch.nn.functional as F
+from torch.nn import Parameter
+from torch.nn.init import normal_
+import torch.utils.checkpoint
+from torch import Tensor, device
+from TAAS_utils import * 
+from transformers.modeling_utils import ModuleUtilsMixin
+from transformers import AutoTokenizer, AutoModel, BertTokenizer
+from graphormer import Graphormer3D 
+import pickle
+import torch
+import sys
+from ner_model import NER_model
+import numpy as np
+
+
+from htc_loss import HTCLoss
+from transformers.utils.hub import cached_file
+remap_code_2_chn_file_path = cached_file( 
+    'Cainiao-AI/TAAS',
+    'remap_code_2_chn.pkl'
+)
+s2_label_dict_remap = {
+ 0: '0',
+ 1: '1',
+ 2: '2',
+ 3: '3',
+ 4: '4',
+ 5: '5',
+ 6: '6',
+ 7: '7',
+ 8: '8',
+ 9: '9',
+ 10: 'a',
+ 11: 'b',
+ 12: 'c',
+ 13: 'd',
+ 14: 'e',
+ 15: 'f'}
+
+class StellarEmbedding(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+        self.ner_type_embeddings = nn.Embedding(10, config.hidden_size)
+        self.use_task_id = config.use_task_id
+        if config.use_task_id:
+            self.task_type_embeddings = nn.Embedding(config.task_type_vocab_size, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+        self.register_buffer("token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long),
+                             persistent=False)
+        self._reset_parameters()
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            token_type_ids: Optional[torch.LongTensor] = None,
+            ner_type_ids: Optional[torch.LongTensor] = None,
+            task_type_ids: Optional[torch.LongTensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            past_key_values_length: int = 0,
+    ) -> torch.Tensor:
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, past_key_values_length: seq_length + past_key_values_length]
+
+        # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+        # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+        # issue #5664
+        if token_type_ids is None:
+            if hasattr(self, "token_type_ids"):
+                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
+                token_type_ids = buffered_token_type_ids_expanded
+            else:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+        if ner_type_ids is not None:
+            ner_type_embeddings = self.ner_type_embeddings(ner_type_ids)
+
+            embeddings = inputs_embeds + token_type_embeddings + ner_type_embeddings
+        else:
+            embeddings = inputs_embeds + token_type_embeddings
+        if self.position_embedding_type == "absolute":
+            position_embeddings = self.position_embeddings(position_ids)
+            embeddings += position_embeddings
+
+        # add `task_type_id` for ERNIE model
+        if self.use_task_id:
+            if task_type_ids is None:
+                task_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+            task_type_embeddings = self.task_type_embeddings(task_type_ids)
+            embeddings += task_type_embeddings
+
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                normal_(p, mean=0.0, std=0.02)
+
+    def set_pretrained_weights(self, path):
+        pre_train_weights = torch.load(path, map_location=torch.device('cpu'))
+        new_weights = dict()
+        for layer in self.state_dict().keys():
+            if layer == 'position_ids':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_ids']
+            elif layer == 'word_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.word_embeddings.weight']
+            elif layer == 'position_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_embeddings.weight']
+            elif layer == 'token_type_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.token_type_embeddings.weight']
+            elif layer == 'task_type_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.task_type_embeddings.weight']
+            elif layer == 'LayerNorm.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.weight']
+            elif layer == 'LayerNorm.bias':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.bias']
+            else:
+                new_weights[layer] = self.state_dict()[layer]
+        self.load_state_dict(new_weights)
+
+    def save_weights(self, path):
+        torch.save(self.state_dict(), path)
+
+    def load_weights(self, path):
+        self.load_state_dict(torch.load(path))
+
+
+# Copied from transformers.models.bert.modeling_bert.BertLayer
+class StellarLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = ErnieAttention(config)
+        self.is_decoder = config.is_decoder
+        self.add_cross_attention = config.add_cross_attention
+        if self.add_cross_attention:
+            if not self.is_decoder:
+                raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
+            self.crossattention = ErnieAttention(config, position_embedding_type="absolute")
+        self.intermediate = ErnieIntermediate(config)
+        self.output = ErnieOutput(config)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.FloatTensor] = None,
+            head_mask: Optional[torch.FloatTensor] = None,
+            encoder_hidden_states: Optional[torch.FloatTensor] = None,
+            encoder_attention_mask: Optional[torch.FloatTensor] = None,
+            past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+            output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+
+        # if decoder, the last output is tuple of self-attn cache
+        if self.is_decoder:
+            outputs = self_attention_outputs[1:-1]
+            present_key_value = self_attention_outputs[-1]
+        else:
+            outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+
+        cross_attn_present_key_value = None
+        if self.is_decoder and encoder_hidden_states is not None:
+            if not hasattr(self, "crossattention"):
+                raise ValueError(
+                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
+                    " by setting `config.add_cross_attention=True`"
+                )
+
+            # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
+            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                cross_attn_past_key_value,
+                output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:-1]  # add cross attentions if we output attention weights
+
+            # add cross-attn cache to positions 3,4 of present_key_value tuple
+            cross_attn_present_key_value = cross_attention_outputs[-1]
+            present_key_value = present_key_value + cross_attn_present_key_value
+
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+
+        # if decoder, return the attn key/values as the last output
+        if self.is_decoder:
+            outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+class StellarEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([StellarLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.FloatTensor] = None,
+            head_mask: Optional[torch.FloatTensor] = None,
+            encoder_hidden_states: Optional[torch.FloatTensor] = None,
+            encoder_attention_mask: Optional[torch.FloatTensor] = None,
+            past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = False,
+            output_hidden_states: Optional[bool] = False,
+            return_dict: Optional[bool] = True,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+
+        next_decoder_cache = () if use_cache else None
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                if use_cache:
+                    logger.warning(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+                if self.config.add_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler
+class StellarPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class StellarModel(nn.Module):
+    """
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__()
+        self.config = config
+        self.encoder = StellarEncoder(config)
+        self.pooler = StellarPooler(config) if add_pooling_layer else None
+        # Initialize weights and apply final processing
+        self._reset_parameters()
+
+    # Copied from transformers.models.bert.modeling_bert.BertModel._prune_heads
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def forward(
+            self,
+            h_input,
+            input_ids: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            token_type_ids: Optional[torch.Tensor] = None,
+            task_type_ids: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.Tensor] = None,
+            head_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            encoder_hidden_states: Optional[torch.Tensor] = None,
+            encoder_attention_mask: Optional[torch.Tensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
+        r"""
+        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.config.is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        batch_size, seq_length = input_shape
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        # past_key_values_length
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+        if attention_mask is None:
+            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+
+        if token_type_ids is None:
+            if hasattr(self.embeddings, "token_type_ids"):
+                buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
+                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
+                token_type_ids = buffered_token_type_ids_expanded
+            else:
+                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        encoder_outputs = self.encoder(
+            h_input,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+    def get_extended_attention_mask(
+            self, attention_mask: Tensor, input_shape: Tuple[int], device: device = None, dtype: torch.float = None
+    ) -> Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (`Tuple[int]`):
+                The shape of the input to the model.
+
+        Returns:
+            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
+        """
+        if dtype is None:
+            dtype = torch.float32
+
+        if not (attention_mask.dim() == 2 and self.config.is_decoder):
+            # show warning only if it won't be shown in `create_extended_attention_mask_for_decoder`
+            if device is not None:
+                warnings.warn(
+                    "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning
+                )
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            if self.config.is_decoder:
+                extended_attention_mask = ModuleUtilsMixin.create_extended_attention_mask_for_decoder(
+                    input_shape, attention_mask, device
+                )
+            else:
+                extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and the dtype's smallest value for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
+        return extended_attention_mask
+
+    def get_head_mask(
+            self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False
+    ) -> Tensor:
+        """
+        Prepare the head mask if needed.
+
+        Args:
+            head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*):
+                The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard).
+            num_hidden_layers (`int`):
+                The number of hidden layers in the model.
+            is_attention_chunked: (`bool`, *optional*, defaults to `False`):
+                Whether or not the attentions scores are computed by chunks or not.
+
+        Returns:
+            `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with
+            `[None]` for each layer.
+        """
+        if head_mask is not None:
+            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
+            if is_attention_chunked is True:
+                head_mask = head_mask.unsqueeze(-1)
+        else:
+            head_mask = [None] * num_hidden_layers
+
+        return head_mask
+
+    def _reset_parameters(self):
+        r"""Initiate parameters in the transformer model."""
+        for p in self.parameters():
+            if p.dim() > 1:
+                normal_(p, mean=0.0, std=self.config.initializer_range)
+
+    def save_weights(self, path):
+        torch.save(self.state_dict(), path)
+
+    def load_weights(self, path):
+        self.load_state_dict(torch.load(path))
+
+
+class TAAS(PreTrainedModel):
+    def __init__(self, config, return_last_hidden_state=False):
+        super(TAAS, self).__init__(config)
+
+        """
+        :param d_model:  d_k = d_v = d_model/nhead = 64, 模型中向量的维度，论文默认值为 512
+        :param nhead:               多头注意力机制中多头的数量，论文默认为值 8
+        :param num_encoder_layers:  encoder堆叠的数量，也就是论文中的N，论文默认值为6
+        :param num_decoder_layers:  decoder堆叠的数量，也就是论文中的N，论文默认值为6
+        :param dim_feedforward:     全连接中向量的维度，论文默认值为 2048
+        :param dropout:             丢弃率，论文中的默认值为 0.1
+        """
+
+        self.config = deepcopy(config)
+        self.return_last_hidden_state = return_last_hidden_state
+        self.dropout = nn.Dropout(self.config.hidden_dropout_prob)
+        #  ================ StellarEmbedding =====================
+        self.embedding = StellarEmbedding(self.config)
+        self.embedding_weights = Parameter(torch.ones(1, 1, self.config.hidden_size))
+        #  ================ StellarModel =====================
+        self.stellar_config = deepcopy(config)
+        self.stellar_model = StellarModel(self.stellar_config)
+        #  ================ TranSAGE =====================
+        # self.transage_layer = TranSAGE()
+        self.graphormer = Graphormer3D()
+        # ================ 解码部分 =====================
+        self.encoder_config = deepcopy(config)
+        self.encoder_config.num_hidden_layers = 1
+        self.encoder = StellarModel(self.encoder_config)
+        self.encoder_out_dim = self.encoder_config.hidden_size
+        # ================ GC任务部分 =====================
+        self.gc_trans = nn.Linear(self.encoder_out_dim, 16 * 33, bias=True)
+        # ================ MLM任务部分 =====================
+        self.cls = ErnieForMaskedLM(self.stellar_config).cls
+        # ================ alias任务部分 =====================
+        self.down_hidden_dim = 512
+        self.down_kernel_num = 128
+        self.alias_trans = nn.Linear(self.encoder_out_dim, self.down_hidden_dim, bias=True)
+        self.alias_trans2 = torch.nn.Conv2d(1, self.down_kernel_num, (2, self.down_hidden_dim), stride=1, bias=True)
+        self.alias_layer = nn.Linear(self.down_kernel_num * 5, 2 * 5, bias=True)
+        # ================ AOI任务部分 =====================
+        self.aoi_trans = nn.Linear(self.encoder_out_dim, self.down_hidden_dim, bias=True)
+        self.aoi_trans2 = torch.nn.Conv2d(1, self.down_kernel_num, (2, self.down_hidden_dim), stride=1, bias=True)
+        self.aoi_layer = nn.Linear(self.down_kernel_num * 5, 2 * 5, bias=True)
+        
+        # ================ HTC任务部分 =====================
+        self.htc_trans = nn.Linear(self.encoder_out_dim, 5 * 100, bias=True)
+
+        # ================ NER任务部分 =====================
+        # self.ner_model = torch.load('ner.pth')
+        self.ner_model = NER_model(vocab_size=11)
+        # self.ner_model.load_state_dict(torch.load('ner.pth'))
+
+
+    def forward(self,
+                input_ids,
+                attention_mask,
+                token_type_ids,
+                node_position_ids,
+                spatial_pos, in_degree, out_degree, edge_type_matrix, edge_input,
+                prov_city_mask: Optional[torch.Tensor] = None,
+                sequence_len=6,
+                labels: Optional[torch.Tensor] = None
+                ):
+        """
+        :param input_ids:   [sequence_len * batch_size, src_len]
+        :param attention_mask:  [sequence_len * batch_size, src_len]
+        :param token_type_ids:  [sequence_len * batch_size, src_len]
+        :param sequence_len:  int
+        :param labels:
+        :param is_eval: bool
+        :return:
+        """
+        batch_size_input = int(input_ids.shape[0] / sequence_len)
+
+        embedding_output = self.embedding(input_ids=input_ids, token_type_ids=token_type_ids)
+        
+        stellar_predictions = self.stellar_model(embedding_output,
+                                                 input_ids=input_ids,
+                                                 token_type_ids=token_type_ids,
+                                                 attention_mask=attention_mask)
+        last_hidden_state = stellar_predictions[0].contiguous().view(batch_size_input, sequence_len, -1,
+                                                                     self.encoder_out_dim)
+        pooler_output = stellar_predictions[1].contiguous().view(batch_size_input, sequence_len, self.encoder_out_dim)
+        h_ = self.graphormer(pooler_output, spatial_pos, in_degree, out_degree, edge_type_matrix, edge_input, node_position_ids)
+        h_ = h_.unsqueeze(2)
+        new_hidden_state = torch.cat((h_, last_hidden_state[:, :, 1:, :]), dim=2)
+        new_hidden_state = new_hidden_state.contiguous().view(batch_size_input * sequence_len, -1, self.encoder_out_dim)
+        encoder_outputs = self.encoder(new_hidden_state,
+                                       input_ids=input_ids,
+                                       token_type_ids=token_type_ids,
+                                       attention_mask=attention_mask)
+        final_hidden_state = encoder_outputs[0]
+        final_pooler_output = encoder_outputs[1].contiguous().view(batch_size_input, sequence_len, self.encoder_out_dim)
+        prediction_scores = self.cls(final_hidden_state)  # 用于 MLM 任务
+
+        gc_layer_out = self.gc_trans(final_pooler_output)
+        gc_layer_out = gc_layer_out.contiguous().view(-1, 16)
+        
+        htc_layer_out = self.htc_trans(final_pooler_output)
+        htc_layer_out = htc_layer_out.contiguous().view(-1, 100)
+
+
+        # MLM loss
+        if labels is not None:
+            # masked_lm_loss = None
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+            return [gc_layer_out, masked_lm_loss, prediction_scores, htc_layer_out]
+
+        if self.return_last_hidden_state:
+            return final_pooler_output, pooler_output
+
+        return gc_layer_out, final_pooler_output, final_hidden_state, prediction_scores, last_hidden_state, htc_layer_out
+
+    def get_htc_code(self, htc_layer_out):
+        htc_loss_fct = HTCLoss(device=self.device, reduction='mean')
+        htc_pred = htc_loss_fct.get_htc_code(htc_layer_out)
+        return htc_pred
+
+    def decode_htc_code_2_chn(self, htc_pred):
+        arr = htc_pred
+        with open(remap_code_2_chn_file_path, 'rb') as fr:
+            remap_code_2_chn = pickle.loads(fr.read())
+        return remap_code_2_chn['{:02d}{:02d}{:02d}{:01d}{:02d}'.format(arr[0], arr[1], arr[2], arr[3], arr[4])]
+
+    # Address Standarization
+    def addr_standardize(self, address):
+        tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh')
+        encoded_input = tokenizer(address, return_tensors='pt', padding='max_length',
+                                truncation=True,  # 超过最大长度截断
+                                max_length=60,
+                                add_special_tokens=True).to(self.device)
+        word_ids = encoded_input['input_ids']
+        attention_mask = encoded_input['attention_mask']
+
+        length = len(word_ids)
+        node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device)
+        in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device)
+        edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device)
+
+        logits = self.ner_model(**encoded_input,
+                                    node_position_ids = node_position_ids,
+                                    spatial_pos = spatial_pos, 
+                                    in_degree = in_degree, 
+                                    out_degree = out_degree,
+                                    edge_type_matrix = edge_type_matrix,
+                                    edge_input = edge_input,)[0]
+        output = []
+        ner_labels = torch.argmax(logits, dim=-1)
+        if len(address) == 1:
+            ner_labels = ner_labels.unsqueeze(0)
+        for i in range(len(address)):
+            ner_label = ner_labels[i]
+            word_id = word_ids[i]
+            # cut padding
+            idx = torch.where(attention_mask[i]>0)
+            ner_label = ner_label[idx][1:-1]
+            word_id = word_id[idx][1:-1]
+            # cut other info 
+            idx1 = torch.where(ner_label != 0)
+            ner_label = ner_label[idx1].tolist()
+            word_id = word_id[idx1].tolist()
+            # add house info
+            if 8 in ner_label:
+                idx2 = ''.join([str(i) for i in ner_label]).rfind('8')
+                word_id.insert(idx2+1, 2770)
+                ner_label.insert(idx2+1, 8)
+            if 9 in ner_label:
+                idx2 = ''.join([str(i) for i in ner_label]).rfind('9')
+                word_id.insert(idx2+1, 269)
+                word_id.insert(idx2+2, 183)
+                ner_label.insert(idx2+1, 9)
+                ner_label.insert(idx2+2, 9)
+            if 10 in ner_label:
+                idx2 = ''.join([str(i) for i in ner_label]).rfind('10')
+                word_id.insert(idx2+1, 485)
+                ner_label.insert(idx2+1, 10)
+
+            output.append(tokenizer.decode(word_id).replace(' ', ''))
+            
+        return output
+
+    # Address Entity Tokenization
+    def addr_entity(self, address):
+        tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh')
+        encoded_input = tokenizer(address, return_tensors='pt', padding='max_length',
+                                truncation=True,  # 超过最大长度截断
+                                max_length=60,
+                                add_special_tokens=True).to(self.device)
+        word_ids = encoded_input['input_ids']
+        attention_mask = encoded_input['attention_mask']
+
+        length = len(word_ids)
+        node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device)
+        in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device)
+        edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device)
+
+        logits = self.ner_model(**encoded_input,
+                                    node_position_ids = node_position_ids,
+                                    spatial_pos = spatial_pos, 
+                                    in_degree = in_degree, 
+                                    out_degree = out_degree,
+                                    edge_type_matrix = edge_type_matrix,
+                                    edge_input = edge_input,)[0]
+
+        ner_labels = torch.argmax(logits, dim=-1)
+        if len(address) == 1:
+            ner_labels = ner_labels.unsqueeze(0)
+
+        output = []
+        tmp = {1:'省', 2:'市', 3:'区', 4:'街道/镇', 5:'道路', 6:'道路号', 7:'poi', 8:'楼栋号', 9:'单元号', 10:'门牌号'}
+        for i in range(len(address)):
+            ner_label = ner_labels[i]
+            word_id = word_ids[i]
+            idx = torch.where(attention_mask[i]>0)
+            ner_label = ner_label[idx][1:-1]
+            word_id = word_id[idx][1:-1]
+            
+            addr_dict = {}
+            addr_dict = dict.fromkeys(tmp.values(),'无')
+            for j in range(1,11):
+                idx = torch.where(ner_label == j)
+                addr_dict[tmp[j]] = ''.join(tokenizer.decode(word_id[idx]).replace(' ',''))
+
+            output.append(deepcopy(addr_dict))
+            
+        return output 
+
+    # House Info Extraction
+    def house_info(self, address):
+        tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh')
+        encoded_input = tokenizer(address, return_tensors='pt', padding='max_length',
+                                truncation=True,  # 超过最大长度截断
+                                max_length=60,
+                                add_special_tokens=True).to(self.device)
+        word_ids = encoded_input['input_ids']
+        attention_mask = encoded_input['attention_mask']
+
+        length = len(word_ids)
+        node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device)
+        in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device)
+        edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device)
+
+        logits = self.ner_model(**encoded_input,
+                                    node_position_ids = node_position_ids,
+                                    spatial_pos = spatial_pos, 
+                                    in_degree = in_degree, 
+                                    out_degree = out_degree,
+                                    edge_type_matrix = edge_type_matrix,
+                                    edge_input = edge_input,)[0]
+
+        ner_labels = torch.argmax(logits, dim=-1)
+        if len(address) == 1:
+            ner_labels = ner_labels.unsqueeze(0)
+        output = []
+        for i in range(len(address)):
+            ner_label = ner_labels[i]
+            word_id = word_ids[i]
+            idx = torch.where(attention_mask[i]>0)
+            ner_label = ner_label[idx][1:-1]
+            word_id = word_id[idx][1:-1]
+
+            building = []
+            unit = []
+            room = []
+            for j in range(len(ner_label)):
+                if ner_label[j] == 8:
+                    building.append(word_id[j])
+                elif ner_label[j] == 9:
+                    unit.append(word_id[j])
+                elif ner_label[j] == 10:
+                    room.append(word_id[j])
+
+            output.append({'楼栋':tokenizer.decode(building).replace(' ',''), '单元':tokenizer.decode(unit).replace(' ',''),  
+                '门牌号': tokenizer.decode(room).replace(' ','')})
+        return output
+                
+
+    # Address Completion
+    def addr_complet(self, address):
+        tokenizer = BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh')
+        encoded_input = tokenizer(address, return_tensors='pt', padding='max_length',
+                                truncation=True,  # 超过最大长度截断
+                                max_length=60,
+                                add_special_tokens=True).to(self.device)
+        word_ids = encoded_input['input_ids']
+        attention_mask = encoded_input['attention_mask']
+
+        length = len(word_ids)
+        node_position_ids = torch.tensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        spatial_pos = torch.LongTensor(np.zeros((length, 1, 1), dtype=np.int64)).to(self.device)
+        in_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        out_degree = torch.LongTensor(np.ones((length, 1), dtype=np.int64)).to(self.device)
+        edge_type_matrix = torch.LongTensor(8*np.ones((length, 1, 1), dtype=np.int64)).to(self.device)
+        edge_input = torch.LongTensor(8*np.ones((length, 1, 1, 1), dtype=np.int64)).to(self.device)
+
+        logits = self.ner_model(**encoded_input,
+                                    node_position_ids = node_position_ids,
+                                    spatial_pos = spatial_pos, 
+                                    in_degree = in_degree, 
+                                    out_degree = out_degree,
+                                    edge_type_matrix = edge_type_matrix,
+                                    edge_input = edge_input,)[0]
+
+        ner_labels = torch.argmax(logits, dim=-1)
+        if len(address) == 1:
+            ner_labels = ner_labels.unsqueeze(0)
+        if isinstance(address, list):
+            address = address[0]
+
+        # HTC result
+        g2ptl_model = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True)
+        g2ptl_model.eval()
+        g2ptl_output = g2ptl_model(**encoded_input)
+        htc_layer_out = g2ptl_output.htc_layer_out
+        arr = g2ptl_model.get_htc_code(htc_layer_out)
+        htc_pred = '{:02d}{:02d}{:02d}{:01d}{:02d}'.format(arr[0], arr[1], arr[2], arr[3], arr[4])
+        with open('remap_code_2_chn_with_all_htc.pkl', 'rb') as fr:
+            remap_code_2_chn = pickle.loads(fr.read())
+        
+        try:
+            htc_list = remap_code_2_chn[htc_pred][-1]
+        except:
+            return address
+
+        # revise address level of four city
+        if htc_list[0] in ['北京','上海','重庆','天津']:
+            htc_list = htc_list[1:]
+            htc_list.append('')
+        
+        idx = torch.where(attention_mask>0)
+        ner_label = ner_labels[idx][1:-1].cpu().numpy().tolist()
+        word_id = word_ids[idx][1:-1]
+        
+        for i in range(1,5):
+            # judge the lacked address unit
+            if i not in ner_label:
+                if i == 1:
+                    address = htc_list[0] + address
+                    ner_label = [1] * len(htc_list[0]) + ner_label
+                else :
+                    # find the insert position
+                    idx = 0
+                    for j in range(len(ner_label)):
+                        if ner_label[j] > i:
+                            idx = j
+                            break
+                    address = address[:idx] + htc_list[i-1] + address[idx:]
+                    ner_label = ner_label[:idx] + [i] * len(htc_list[i-1]) + ner_label[idx:]
+
+        return address
+
+    # Geo-locating from text to geospatial
+    def geolocate(self, address):
+        g2ptl_model = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True)
+        encoded_input = tokenizer(address, return_tensors='pt')
+
+        g2ptl_model.eval()
+        output = g2ptl_model(**encoded_input)
+        geo_labels = torch.argmax(output.gc_layer_out, dim=-1)
+        output = [s2_label_dict_remap[int(i)] for i in geo_labels]
+
+        return 's2网格化结果：' + ''.join(output)
+
+    # Pick-up Estimation Time of Arrival
+    def pickup_ETA(self, address):
+        print('Users can get the address embeddings using model.encode(address) and feed them to your own ETA model.')
+    
+    # Pick-up and Delivery Route Prediction
+    def route_predict(self, route_data):
+        print('Users can get the address embeddings using model.encode(address) and feed them to your own Route Prediction model.')
+
+    # Address embeddings
+    def encode(self, address):
+        tokenizer = AutoTokenizer.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True)
+        g2ptl_model = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True)
+        encoded_input = tokenizer(address, return_tensors='pt', padding='max_length',
+                                truncation=True,  # 超过最大长度截断
+                                max_length=60,
+                                add_special_tokens=True)
+        g2ptl_model.eval()
+        output = g2ptl_model(**encoded_input)
+
+        return output.final_hidden_state
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                xavier_uniform_(p)
+
+    def generate_square_subsequent_mask(self, sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask  # [sz,sz]
+
+    def save_weights(self, path):
+        torch.save(self.state_dict(), path)
+
+    def load_weights(self, path):
+        self.load_state_dict(torch.load(path, map_location=torch.device('cpu')), False)
+
+    def set_pretrained_weights(self, path):
+        pre_train_weights = torch.load(path, map_location=torch.device('cpu'))
+        new_weights = dict()
+
+        for layer in self.state_dict().keys():
+            if layer == 'embedding.position_ids':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_ids']
+            elif layer == 'embedding.word_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.word_embeddings.weight']
+            elif layer == 'embedding.position_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.position_embeddings.weight']
+            elif layer == 'embedding.token_type_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.token_type_embeddings.weight']
+            elif layer == 'embedding.task_type_embeddings.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.task_type_embeddings.weight']
+            elif layer == 'embedding.LayerNorm.weight':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.weight']
+            elif layer == 'embedding.LayerNorm.bias':
+                new_weights[layer] = pre_train_weights['ernie_model.embeddings.LayerNorm.bias']
+            elif 'stellar_model' in layer:
+                new_weights[layer] = pre_train_weights[layer.replace('stellar_model', 'ernie_model')]
+            elif layer in pre_train_weights.keys():
+                new_weights[layer] = pre_train_weights[layer]
+            else:
+                new_weights[layer] = self.state_dict()[layer]
+
+        self.load_state_dict(new_weights)

ner_model.py

@@ -0,0 +1,88 @@
+#! python3
+# -*- encoding: utf-8 -*-
+
+import torch
+import torch.nn as nn
+from typing import Optional
+from transformers import AutoModel
+from torch.nn.init import xavier_uniform_
+
+def cal_ner_acc(y, y_hat):
+    if len(y) == 0:
+        return 0, 1
+    y,y_hat = y.cpu().numpy(), y_hat.cpu().numpy()
+    
+    acc_cnt, len_cnt = 0, 0
+    for i in range(len(y)):
+        if y[i] <= 7 and y_hat[i] <= 7:
+            len_cnt += 1
+            if y[i] == y_hat[i]:
+                acc_cnt += 1
+    
+    return acc_cnt, len_cnt
+
+
+class NER_model(nn.Module):
+    def __init__(self, vocab_size):
+        super(NER_model, self).__init__()
+        
+        while True:
+            try:
+                self.g2ptl = AutoModel.from_pretrained('Cainiao-AI/G2PTL', trust_remote_code=True)
+                break
+            except:
+                continue
+        """
+        Ner head
+        """
+        # print('model loaded.')
+        self.dropout = nn.Dropout(p = 0.1, inplace = False)
+        self.linear1 = nn.Linear(in_features=768, out_features=128, bias=True)
+        self.linear2 = nn.Linear(in_features=128, out_features=vocab_size, bias=True)
+        # self.classifier = nn.Linear(in_features=768, out_features=vocab_size, bias=True)
+        # self.cls = ErnieForMaskedLM.from_pretrained('nghuyong/ernie-3.0-base-zh').cls
+        #self._reset_parameters()
+
+    def forward(self,
+                input_ids,
+                attention_mask,
+                token_type_ids,
+                node_position_ids,spatial_pos, in_degree, out_degree, edge_type_matrix, edge_input,
+                prov_city_mask: Optional[torch.Tensor] = None,
+                sequence_len=6,
+                labels: Optional[torch.Tensor] = None
+                ):
+        output= self.g2ptl(input_ids, attention_mask, token_type_ids, node_position_ids, spatial_pos,
+                                    in_degree,
+                                    out_degree,
+                                    edge_type_matrix,
+                                    edge_input )
+        
+        pooler_output_embedding = output.final_hidden_state
+        sequence_output = pooler_output_embedding.squeeze()
+        # Input的是Bert输出的token sequence的embedding，而不是pooler的embedding
+        sequence_output = self.dropout(sequence_output)
+        linear_out = self.linear1(sequence_output)
+        logits = self.linear2(self.dropout(linear_out))
+        # logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        return [logits, loss]
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                xavier_uniform_(p)
+
+    def save_weights(self, path):
+        torch.save(self.state_dict(), path)
+
+    def load_weights(self, path):
+        self.load_state_dict(torch.load(path, map_location=torch.device('cpu')), False)
+        
+        
+   

pytorch_model.bin

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

remap_code_2_chn.pkl

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

remap_code_2_chn_with_all_htc.pkl

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

requirements.txt

@@ -0,0 +1,30 @@
+numpy
+pandas
+tensorboard
+tqdm==4.64.1
+transformers==4.25.1
+utils
+datasets
+oss2
+fairseq 
+tensorboardX
+rouge
+matplotlib
+seaborn
+SentencePiece
+ujson
+eas_prediction
+openpyxl
+s2sphere
+s2cell
+tensorboard
+onnx
+onnxsim
+
+# lightseq
+# onnxruntime
+# tqdm
+# torch==1.13.1
+# transformers==4.27.4
+# datasets
+# fairseq

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "special_tokens_map_file": null,
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

utils.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+# @File    : utils.py
+# @Author  : 刘建林(霜旻)
+# @Email   : [email protected]
+# @Time    : 2022/10/27 下午8:52
+"""
+import operator
+import pickle
+import numpy as np
+import pandas as pd
+
+s2_label_dict = {
+    '0': 0,
+    '1': 1,
+    '2': 2,
+    '3': 3,
+    '4': 4,
+    '5': 5,
+    '6': 6,
+    '7': 7,
+    '8': 8,
+    '9': 9,
+    'a': 10,
+    'b': 11,
+    'c': 12,
+    'd': 13,
+    'e': 14,
+    'f': 15
+}
+s2_label_decode_dict = {v: k for k, v in s2_label_dict.items()}
+
+s2_weights = [0.025, 0.025, 0.025,
+              0.025, 0.025, 0.025,
+              0.025, 0.025, 0.025,
+              0.0325, 0.0325, 0.0325,
+              0.035, 0.035, 0.035,
+              0.0375, 0.0375, 0.0375,
+              0.04, 0.04, 0.04,
+              0.0425, 0.0425, 0.0425,
+              0.045, 0.045, 0.0475,
+              0.025, 0.025, 0.025,
+              0.0, 0.0, 0.0]
+
+def generate_s2_index(s2_label):
+    result = [0 for _ in range(33)]
+    for i, char_ in enumerate(s2_label):
+        result[i] = s2_label_dict[char_]
+    return result
+
+
+def decode_s2(x):
+    result = []
+    for i in x:
+        result.append(s2_label_decode_dict[i])
+    return ''.join(result)
+
+
+def sample_csv2pkl(csv_path, pkl_path):
+    # df = pd.read_csv('/Users/liujianlin/odps_clt_release_64/bin/addr6node_small1.csv', sep='^', encoding="utf_8_sig")
+    df = pd.read_csv(csv_path, sep='^', encoding="utf_8_sig")
+    # print(df)
+    data = []
+    for index, row in df.iterrows():
+        node_s = []
+        label = []
+        node1 = [row['node_t1'], row['poi_address_mask1'], row['node1'], generate_s2_index(row['node1'])]
+        node2 = [row['node_t2'], row['poi_address_mask2'], row['node2'], generate_s2_index(row['node2'])]
+        node3 = [row['node_t3'], row['poi_address_mask3'], row['node3'], generate_s2_index(row['node3'])]
+        node4 = [row['node_t4'], row['poi_address_mask4'], row['node4'], generate_s2_index(row['node4'])]
+        node5 = [row['node_t5'], row['poi_address_mask5'], row['node5'], generate_s2_index(row['node5'])]
+        node6 = [row['node_t6'], row['poi_address_mask6'], row['node6'], generate_s2_index(row['node6'])]
+        label.extend(node1[3])
+        label.extend(node2[3])
+        label.extend(node3[3])
+        label.extend(node4[3])
+        label.extend(node5[3])
+        label.extend(node6[3])
+        node1.append(label)
+        node2.append(label)
+        node3.append(label)
+        node4.append(label)
+        node5.append(label)
+        node6.append(label)
+        node_s.append(node1)
+        node_s.append(node2)
+        node_s.append(node3)
+        node_s.append(node4)
+        node_s.append(node5)
+        node_s.append(node6)
+        data.append(node_s)
+    # print(data)
+
+    with open(pkl_path,'wb') as f:
+        pickle.dump(data,f)
+
+
+def calculate_multi_s2_acc(predicted_s2, y):
+    acc_cnt = np.array([0, 0, 0, 0, 0, 0, 0])
+    y = y.view(-1, 33).tolist()
+    predicted = predicted_s2.view(-1, 33).tolist()
+    # print(y.shape, predicted.shape)
+    for index, s2 in enumerate(y):
+        for c, i in enumerate(range(12, 33, 3)):
+            y_l10 = y[index][12:i+3]
+            p_l10 = predicted[index][12:i+3]
+            # print(y_l10, p_l10, operator.eq(y_l10, p_l10))
+            if operator.eq(y_l10, p_l10):
+                acc_cnt[c] += 1
+        # print('==='*20)
+    # print(acc_cnt)
+    return acc_cnt
+
+def calculate_multi_s2_acc_batch(predicted_s2, y, sequence_len = 6):
+    acc_cnt = np.array([0, 0, 0, 0, 0, 0, 0])
+    y = y.view(-1, sequence_len, 33).tolist()
+    predicted = predicted_s2.view(-1, sequence_len, 33).tolist()
+    # print(y.shape, predicted.shape)
+    batch_size = len(y)
+    for batch_i in range(batch_size):
+        for index, s2 in enumerate(y[batch_i]):
+            for c, i in enumerate(range(12, 33, 3)):
+                y_l10 = y[batch_i][index][12:i+3]
+                p_l10 = predicted[batch_i][index][12:i+3]
+                # print(y_l10, p_l10, operator.eq(y_l10, p_l10))
+                if operator.eq(y_l10, p_l10):
+                    acc_cnt[c] += 1
+        # print('==='*20)
+    # print(acc_cnt)
+    return acc_cnt
+
+
+
+def calculate_alias_acc(predicted, y):
+    tp, fp, fn, tn = 0, 0, 0, 0
+    acc = 0
+    for index, label in enumerate(y):
+        if int(label) == int(predicted[index]):
+            acc += 1
+        if int(label) == 1:
+            fn += 1
+            if int(predicted[index]) == 1:
+                tp += 1
+    if fn == 0:
+        precision = 0
+    else:
+        precision = tp / fn * 100
+    return tp, fn, acc
+
+
+def calculate_aoi_acc(predicted, y):
+    tp, fp, fn, tn = 0, 0, 0, 0
+    acc = 0
+    for index, label in enumerate(y):
+        if int(label) == int(predicted[index]):
+            acc += 1
+        if int(label) == 0:
+            fn += 1
+            if int(predicted[index]) == 0:
+                tp += 1
+    if fn == 0:
+        precision = 0
+    else:
+        precision = tp / fn * 100
+    return tp, fn, acc