Merge branch 'master-github' into master-merge-github-0412

modelscope/cli/download.py

@@ -7,7 +7,7 @@ from modelscope.hub.snapshot_download import snapshot_download
 
 
 def subparser_func(args):
-    """ Fuction which will be called for a specific sub parser.
+    """ Function which will be called for a specific sub parser.
     """
     return DownloadCMD(args)
 

modelscope/cli/modelcard.py

@@ -18,7 +18,7 @@ template_path = os.path.join(curren_path, 'template')
 
 
 def subparser_func(args):
-    """ Fuction which will be called for a specific sub parser.
+    """ Function which will be called for a specific sub parser.
     """
     return ModelCardCMD(args)
 

modelscope/cli/pipeline.py

@@ -13,7 +13,7 @@ template_path = os.path.join(curren_path, 'template')
 
 
 def subparser_func(args):
-    """ Fuction which will be called for a specific sub parser.
+    """ Function which will be called for a specific sub parser.
     """
     return PipelineCMD(args)
 

modelscope/cli/plugins.py

@@ -9,7 +9,7 @@ plugins_manager = PluginsManager()
 
 
 def subparser_func(args):
-    """ Fuction which will be called for a specific sub parser.
+    """ Function which will be called for a specific sub parser.
     """
     return PluginsCMD(args)
 

modelscope/cli/template/template.tpl

@@ -24,7 +24,7 @@ class ${model_name}(TorchModel):
     def init_model(self, **kwargs):
         """Provide default implementation based on TorchModel and user can reimplement it.
             include init model and load ckpt from the model_dir, maybe include preprocessor
-            if nothing to do, then return lambdx x: x
+            if nothing to do, then return lambda x: x
         """
         return lambda x: x
 
@@ -41,7 +41,7 @@ class ${preprocessor_name}(Preprocessor):
 
     def init_preprocessor(self, **kwarg):
         """ Provide default implementation based on preprocess_cfg and user can reimplement it.
-            if nothing to do, then return lambdx x: x
+            if nothing to do, then return lambda x: x
         """
         return lambda x: x
 

modelscope/exporters/nlp/model_for_token_classification_exporter.py

@@ -89,7 +89,7 @@ class ModelForSequenceClassificationExporter(TorchModelExporter):
             outputs_origin = list(numpify_tensor_nested(outputs_origin))
 
         outputs_origin = [outputs_origin[0]
-                          ]  # keeo `predictions`, drop other outputs
+                          ]  # keep `predictions`, drop other outputs
 
         np_dummy_inputs = numpify_tensor_nested(dummy_inputs)
         np_dummy_inputs['label_mask'] = np_dummy_inputs['label_mask'].astype(

modelscope/metainfo.py

@@ -170,6 +170,7 @@ class Models(object):
     plug_mental = 'plug-mental'
     doc2bot = 'doc2bot'
     peer = 'peer'
+    llama = 'llama'
 
     # audio models
     sambert_hifigan = 'sambert-hifigan'
@@ -230,6 +231,7 @@ class Heads(object):
     fill_mask = 'fill-mask'
     bert_mlm = 'bert-mlm'
     roberta_mlm = 'roberta-mlm'
+    xlm_roberta_mlm = 'xlm-roberta-mlm'
     # token cls
     token_classification = 'token-classification'
     # extraction

modelscope/models/audio/kws/farfield/fsmn_sele_v2.py

@@ -18,7 +18,7 @@ class FSMNUnit(nn.Module):
         Args:
             dimlinear:              input / output dimension
             dimproj:                fsmn input / output dimension
-            lorder:                 left ofder
+            lorder:                 left order
             rorder:                 right order
         """
         super(FSMNUnit, self).__init__()

modelscope/models/audio/kws/nearfield/fsmn.py

@@ -435,7 +435,7 @@ class FSMN(nn.Module):
         """
         Args:
             input (torch.Tensor): Input tensor (B, T, D)
-            in_cache(torhc.Tensor): (B, D, C), C is the accumulated cache size
+            in_cache(torch.Tensor): (B, D, C), C is the accumulated cache size
         """
 
         # print("FSMN forward!!!!")

modelscope/models/audio/kws/nearfield/model.py

@@ -39,8 +39,8 @@ class FSMNDecorator(TorchModel):
             model_dir (str): the model path.
             cmvn_file (str): cmvn file
             backbone (dict): params related to backbone
-            input_dim (int): input dimention of network
-            output_dim (int): output dimention of network
+            input_dim (int): input dimension of network
+            output_dim (int): output dimension of network
             training (bool): training or inference mode
         """
         super().__init__(model_dir, *args, **kwargs)
@@ -108,7 +108,7 @@ class FSMNDecorator(TorchModel):
 class KWSModel(nn.Module):
     """Our model consists of four parts:
     1. global_cmvn: Optional, (idim, idim)
-    2. preprocessing: feature dimention projection, (idim, hdim)
+    2. preprocessing: feature dimension projection, (idim, hdim)
     3. backbone: backbone or feature extractor of the whole network, (hdim, hdim)
     4. classifier: output layer or classifier of KWS model, (hdim, odim)
     5. activation:
@@ -133,7 +133,7 @@ class KWSModel(nn.Module):
             odim (int): output dimension of network
             hdim (int): hidden dimension of network
             global_cmvn (nn.Module): cmvn for input feature, (idim, idim)
-            preprocessing (nn.Module): feature dimention projection, (idim, hdim)
+            preprocessing (nn.Module): feature dimension projection, (idim, hdim)
             backbone (nn.Module): backbone or feature extractor of the whole network, (hdim, hdim)
             classifier (nn.Module): output layer or classifier of KWS model, (hdim, odim)
             activation (nn.Module): nn.Identity for training, nn.Sigmoid for inference

modelscope/models/cv/action_recognition/temporal_patch_shift_transformer.py

@@ -871,7 +871,7 @@ class SwinTransformer2D_TPS(nn.Module):
 
         Args:
             logger (logging.Logger): The logger used to print
-                debugging infomation.
+                debugging information.
         """
         checkpoint = torch.load(self.pretrained, map_location='cpu')
         state_dict = checkpoint['model']

modelscope/models/cv/body_3d_keypoints/cannonical_pose/body_3d_pose.py

@@ -181,7 +181,7 @@ class BodyKeypointsDetection3D(TorchModel):
                 "camera_pose": Tensor, [1, NUM_FRAME, OUT_NUM_JOINTS, OUT_3D_FEATURE_DIM],
                     3D human pose keypoints in camera frame.
                 "camera_traj": Tensor, [1, NUM_FRAME, 1, 3],
-                    root keypoints coordinates in camere frame.
+                    root keypoints coordinates in camera frame.
         """
         inputs_2d = input['inputs_2d']
         pose2d_rr = input['pose2d_rr']

modelscope/models/cv/body_3d_keypoints/hdformer/directed_graph.py

@@ -46,14 +46,14 @@ class DiGraph():
         super().__init__()
         self.num_nodes = len(skeleton.parents())
         self.directed_edges_hop1 = [
-            (parrent, child)
-            for child, parrent in enumerate(skeleton.parents()) if parrent >= 0
+            (parent, child) for child, parent in enumerate(skeleton.parents())
+            if parent >= 0
         ]
         self.directed_edges_hop2 = [(0, 1, 2), (0, 4, 5), (0, 7, 8), (1, 2, 3),
                                     (4, 5, 6), (7, 8, 9),
                                     (7, 8, 11), (7, 8, 14), (8, 9, 10),
                                     (8, 11, 12), (8, 14, 15), (11, 12, 13),
-                                    (14, 15, 16)]  # (parrent, child)
+                                    (14, 15, 16)]  # (parent, child)
         self.directed_edges_hop3 = [(0, 1, 2, 3), (0, 4, 5, 6), (0, 7, 8, 9),
                                     (7, 8, 9, 10), (7, 8, 11, 12),
                                     (7, 8, 14, 15), (8, 11, 12, 13),
@@ -112,8 +112,8 @@ class Graph():
         # edge is a list of [child, parent] paris
         self.num_node = len(skeleton.parents())
         self_link = [(i, i) for i in range(self.num_node)]
-        neighbor_link = [(child, parrent)
-                         for child, parrent in enumerate(skeleton.parents())]
+        neighbor_link = [(child, parent)
+                         for child, parent in enumerate(skeleton.parents())]
         self.self_link = self_link
         self.neighbor_link = neighbor_link
         self.edge = self_link + neighbor_link

modelscope/models/cv/cartoon/facelib/facer.py

@@ -94,7 +94,7 @@ class FaceAna():
         sorted_bboxes = [bboxes[x] for x in picked]
         return np.array(sorted_bboxes)
 
-    def judge_boxs(self, previuous_bboxs, now_bboxs):
+    def judge_boxs(self, previous_bboxs, now_bboxs):
 
         def iou(rec1, rec2):
 
@@ -116,17 +116,16 @@ class FaceAna():
 
             return intersect / (sum_area - intersect)
 
-        if previuous_bboxs is None:
+        if previous_bboxs is None:
             return now_bboxs
 
         result = []
 
         for i in range(now_bboxs.shape[0]):
             contain = False
-            for j in range(previuous_bboxs.shape[0]):
-                if iou(now_bboxs[i], previuous_bboxs[j]) > self.iou_thres:
-                    result.append(
-                        self.smooth(now_bboxs[i], previuous_bboxs[j]))
+            for j in range(previous_bboxs.shape[0]):
+                if iou(now_bboxs[i], previous_bboxs[j]) > self.iou_thres:
+                    result.append(self.smooth(now_bboxs[i], previous_bboxs[j]))
                     contain = True
                     break
             if not contain:

modelscope/models/nlp/T5/backbone.py

@@ -24,6 +24,8 @@ import torch
 from torch import nn
 from torch.utils.checkpoint import checkpoint
 from transformers.activations import ACT2FN
+from transformers.modeling_outputs import \
+    BaseModelOutputWithPastAndCrossAttentions
 from transformers.modeling_utils import (PreTrainedModel,
                                          find_pruneable_heads_and_indices,
                                          prune_linear_layer)
@@ -1184,7 +1186,7 @@ class T5Stack(T5PreTrainedModel):
                 all_attentions,
                 all_cross_attentions,
             ] if v is not None)
-        return AttentionBackboneModelOutput(
+        return BaseModelOutputWithPastAndCrossAttentions(
             last_hidden_state=hidden_states,
             past_key_values=present_key_value_states,
             hidden_states=all_hidden_states,

modelscope/models/nlp/__init__.py

@@ -72,6 +72,7 @@ if TYPE_CHECKING:
                        DocumentGroundedDialogRetrievalModel,
                        DocumentGroundedDialogRerankModel)
     from .xlm_roberta import XLMRobertaConfig, XLMRobertaModel
+    from .llama import LlamaForTextGeneration, LlamaConfig, LlamaModel, LlamaTokenizer, LlamaTokenizerFast
 
 else:
     _import_structure = {
@@ -155,6 +156,10 @@ else:
             'DocumentGroundedDialogRerankModel'
         ],
         'xlm_roberta': ['XLMRobertaConfig', 'XLMRobertaModel'],
+        'llama': [
+            'LlamaForTextGeneration', 'LlamaConfig', 'LlamaModel',
+            'LlamaTokenizer', 'LlamaTokenizerFast'
+        ],
     }
 
     import sys

modelscope/models/nlp/gpt3/distributed_gpt3.py

@@ -955,7 +955,6 @@ class DistributedGPT3(TorchModel):
                  megatron_cfg=None,
                  **kwargs):
         super().__init__(model_dir, *args, **kwargs)
-
         init_megatron_util(megatron_cfg, model_dir, rank=rank)
 
         self.config = GPT3Config.from_pretrained(model_dir)
@@ -981,7 +980,8 @@ class DistributedGPT3(TorchModel):
         load_model = pre_load(ckpt_rank, model_dir, tag=path_load_tag)
         load_model = split_state_dict(load_model, model, tensor_ws // ckpt_ws)
 
-        self.dist_model.load_state_dict(load_model)
+        self.dist_model.load_state_dict(
+            load_model, strict=kwargs.get('strict', True))
 
         self.inference_params = None
 

modelscope/models/nlp/heads/fill_mask_head.py

@@ -21,7 +21,7 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from torch.nn import CrossEntropyLoss
-from transformers.activations import ACT2FN
+from transformers.activations import ACT2FN, gelu
 
 from modelscope.metainfo import Heads
 from modelscope.models.base import TorchHead
@@ -71,6 +71,51 @@ class BertFillMaskHead(TorchHead):
         return masked_lm_loss
 
 
+@HEADS.register_module(Tasks.fill_mask, module_name=Heads.xlm_roberta_mlm)
+class XlmRobertaMaskHead(TorchHead):
+    _keys_to_ignore_on_load_missing = [
+        r'lm_head.decoder.weight', 'lm_head.decoder.bias'
+    ]
+
+    def __init__(self,
+                 hidden_size=1024,
+                 hidden_act='gelu',
+                 layer_norm_eps=1e-05,
+                 vocab_size=274701,
+                 **kwargs):
+        super().__init__(
+            hidden_size=hidden_size,
+            hidden_act=hidden_act,
+            layer_norm_eps=layer_norm_eps,
+            vocab_size=vocab_size)
+        self.lm_head = XLMRobertaLMHead(self.config)
+
+    def forward(self,
+                inputs: ModelOutputBase,
+                attention_mask=None,
+                labels=None,
+                **kwargs):
+        logits = self.lm_head(inputs.last_hidden_state)
+        loss = None
+        if labels is not None:
+            loss = self.compute_loss(logits, labels)
+        return AttentionFillMaskModelOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=inputs.hidden_states,
+            attentions=inputs.attentions,
+        )
+
+    def compute_loss(self, logits: torch.Tensor, labels) -> torch.Tensor:
+        loss_fct = CrossEntropyLoss()
+        masked_lm_loss = loss_fct(
+            logits.view(-1, self.config.vocab_size), labels.view(-1))
+        return masked_lm_loss
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+
 class BertPredictionHeadTransform(nn.Module):
 
     def __init__(self, config):
@@ -121,3 +166,35 @@ class BertOnlyMLMHead(nn.Module):
     def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
         prediction_scores = self.predictions(sequence_output)
         return prediction_scores
+
+
+class XLMRobertaLMHead(nn.Module):
+    """Roberta Head for masked language modeling."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.layer_norm = nn.LayerNorm(
+            config.hidden_size, eps=config.layer_norm_eps)
+
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        self.decoder.bias = self.bias
+
+    def forward(self, features, **kwargs):
+        x = self.dense(features)
+        x = gelu(x)
+        x = self.layer_norm(x)
+
+        # project back to size of vocabulary with bias
+        x = self.decoder(x)
+
+        return x
+
+    def _tie_weights(self):
+        # To tie those two weights if they get disconnected (on TPU or when the bias is resized)
+        # For accelerate compatibility and to not break backward compatibility
+        if self.decoder.bias.device.type == 'meta':
+            self.decoder.bias = self.bias
+        else:
+            self.bias = self.decoder.bias

modelscope/models/nlp/llama/__init__.py

@@ -0,0 +1,29 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from modelscope.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+    from .configuration import LlamaConfig
+    from .text_generation import LlamaForTextGeneration
+    from .backbone import LlamaModel
+    from .tokenization import LlamaTokenizer
+    from .tokenization_fast import LlamaTokenizerFast
+else:
+    _import_structure = {
+        'configuration': ['LlamaConfig'],
+        'text_generation': ['LlamaForTextGeneration'],
+        'backbone': ['LlamaModel'],
+        'tokenization': ['LlamaTokenizer'],
+        'tokenization_fast': ['LlamaTokenizerFast'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/models/nlp/llama/backbone.py

@@ -0,0 +1,687 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch LLaMA model."""
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_utils import PreTrainedModel
+
+from modelscope.metainfo import Models
+from modelscope.models import Model, TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.outputs import AttentionBackboneModelOutput
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+from .configuration import LlamaConfig
+
+logger = get_logger(__name__)
+
+_CONFIG_FOR_DOC = 'LlamaConfig'
+
+
+# This file is mainly copied from the llama code of transformers
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(input_ids_shape: torch.Size,
+                      dtype: torch.dtype,
+                      device: torch.device,
+                      past_key_values_length: int = 0):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len),
+                      torch.tensor(torch.finfo(dtype).min, device=device),
+                      device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat(
+            [
+                torch.zeros(
+                    tgt_len,
+                    past_key_values_length,  # noqa
+                    dtype=dtype,
+                    device=device),
+                mask
+            ],
+            dim=-1)  # noqa
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len,
+                                         tgt_len + past_key_values_length)
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor,
+                 dtype: torch.dtype,
+                 tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len,
+                                                  src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(
+        inverted_mask.to(torch.bool),
+        torch.finfo(dtype).min)
+
+
+class LlamaRMSNorm(nn.Module):
+
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        variance = hidden_states.to(torch.float32).pow(2).mean(
+            -1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance
+                                                    + self.variance_epsilon)
+
+        # convert into half-precision if necessary
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+
+        return self.weight * hidden_states
+
+
+class LlamaRotaryEmbedding(torch.nn.Module):
+
+    def __init__(self,
+                 dim,
+                 max_position_embeddings=2048,
+                 base=10000,
+                 device=None):
+        super().__init__()
+        inv_freq = 1.0 / (
+            base**(torch.arange(0, dim, 2).float().to(device) / dim))
+        self.register_buffer('inv_freq', inv_freq)
+
+        # Build here to make `torch.jit.trace` work.
+        self.max_seq_len_cached = max_position_embeddings
+        t = torch.arange(
+            self.max_seq_len_cached,
+            device=self.inv_freq.device,
+            dtype=self.inv_freq.dtype)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer(
+            'cos_cached', emb.cos()[None, None, :, :], persistent=False)
+        self.register_buffer(
+            'sin_cached', emb.sin()[None, None, :, :], persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+            t = torch.arange(
+                self.max_seq_len_cached,
+                device=x.device,
+                dtype=self.inv_freq.dtype)
+            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+            # Different from paper, but it uses a different permutation in order to obtain the same calculation
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.register_buffer(
+                'cos_cached', emb.cos()[None, None, :, :], persistent=False)
+            self.register_buffer(
+                'sin_cached', emb.sin()[None, None, :, :], persistent=False)
+        return (
+            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+        )
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    gather_indices = position_ids[:, None, :, None]  # [bs, 1, seq_len, 1]
+    gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3])
+    cos = torch.gather(
+        cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
+    sin = torch.gather(
+        sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class LlamaMLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+    ):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.act_fn = ACT2FN[hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LlamaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.max_position_embeddings = config.max_position_embeddings
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f'hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}'
+                f' and `num_heads`: {self.num_heads}).')
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        self.rotary_emb = LlamaRotaryEmbedding(
+            self.head_dim,
+            max_position_embeddings=self.max_position_embeddings)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads,
+                           self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
+               Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states).view(
+            bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(
+            bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(
+            bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin, position_ids)
+        # [bsz, nh, t, hd]
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(
+            2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f'Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is'
+                f' {attn_weights.size()}')
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f'Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}'
+                )
+            attn_weights = attn_weights + attention_mask
+            attn_weights = torch.max(
+                attn_weights,
+                torch.tensor(torch.finfo(attn_weights.dtype).min))
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(
+            attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'
+                f' {attn_output.size()}')
+
+        attn_output = attn_output.transpose(1, 2)
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LlamaDecoderLayer(nn.Module):
+
+    def __init__(self, config: LlamaConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = LlamaAttention(config=config)
+        self.mlp = LlamaMLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+        )
+        self.input_layernorm = LlamaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = LlamaRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor,
+                                                 torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            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`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, )
+
+        if output_attentions:
+            outputs += (self_attn_weights, )
+
+        if use_cache:
+            outputs += (present_key_value, )
+
+        return outputs
+
+
+class LlamaPreTrainedModel(TorchModel, PreTrainedModel):
+    config_class = LlamaConfig
+    base_model_prefix = 'model'
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['LlamaDecoderLayer']
+    _keys_to_ignore_on_load_unexpected = [r'decoder\.version']
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config.name_or_path, **kwargs)
+        super(Model, self).__init__(config)
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, LlamaModel):
+            module.gradient_checkpointing = value
+
+    @classmethod
+    def _instantiate(cls, **kwargs):
+        """Instantiate the model.
+
+        Args:
+            kwargs: Input args.
+                    model_dir: The model dir used to load the checkpoint and the label information.
+                    num_labels: An optional arg to tell the model how many classes to initialize.
+                                    Method will call utils.parse_label_mapping if num_labels not supplied.
+                                    If num_labels is not found, the model will use the default setting (2 classes).
+
+        Returns:
+            The loaded model, which is initialized by transformers.PreTrainedModel.from_pretrained
+        """
+
+        model_dir = kwargs.pop('model_dir', None)
+        if model_dir is None:
+            config = LlamaConfig(**kwargs)
+            model = cls(config)
+        else:
+            model = super(Model, cls).from_pretrained(
+                pretrained_model_name_or_path=model_dir, **kwargs)
+        model.model_dir = model_dir
+        return model
+
+
+@MODELS.register_module(Tasks.backbone, module_name=Models.llama)
+class LlamaModel(LlamaPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]
+
+    Args:
+        config: LlamaConfig
+    """
+
+    def __init__(self, config: LlamaConfig, **kwargs):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size,
+                                         self.padding_idx)
+        self.layers = nn.ModuleList([
+            LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)
+        ])
+        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape,
+                                        inputs_embeds, past_key_values_length):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype,
+                tgt_len=input_shape[-1]).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else
+                expanded_attn_mask + combined_attention_mask)
+
+        return combined_attention_mask
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[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, AttentionBackboneModelOutput]:
+        r"""
+        Args:
+            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+                Indices of input sequence tokens in the vocabulary.
+                Padding will be ignored by default should you provide it.
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                [What are input IDs?](../glossary#input-ids)
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+
+                Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+                [`PreTrainedTokenizer.__call__`] for details.
+
+                If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+                `past_key_values`).
+
+                If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+                and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+                information on the default strategy.
+
+                - 1 indicates the head is **not masked**,
+                - 0 indicates the head is **masked**.
+            position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Indices of positions of each input sequence tokens in the position embeddings.
+                Selected in the range `[0, config.n_positions - 1]`.
+
+                [What are position IDs?](../glossary#position-ids)
+            past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed
+                or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`,
+                with each tuple having 2 tensors of shape
+                `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
+                `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
+
+                Contains pre-computed hidden-states (key and values in the self-attention blocks and in the
+                cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential 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)`.
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            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 (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+                tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+                more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        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)
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                'You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time'
+            )
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError(
+                'You have to specify either decoder_input_ids or decoder_inputs_embeds'
+            )
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device)
+            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+        else:
+            position_ids = position_ids.view(-1, seq_length).long()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        # embed positions
+        if attention_mask is None:
+            attention_mask = torch.ones((batch_size, seq_length_with_past),
+                                        dtype=torch.bool,
+                                        device=inputs_embeds.device)
+        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask, (batch_size, seq_length), inputs_embeds,
+            past_key_values_length)
+
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...'
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states, )
+
+            past_key_value = past_key_values[
+                idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (
+                    layer_outputs[2 if output_attentions else 1], )
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1], )
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states, )
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(
+                v for v in
+                [hidden_states, next_cache, all_hidden_states, all_self_attns]
+                if v is not None)
+        return AttentionBackboneModelOutput(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )

modelscope/models/nlp/llama/configuration.py

@@ -0,0 +1,101 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" LLaMA model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+# This file is mainly copied from the llama code of transformers
+class LlamaConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the LLaMA-7B.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`LlamaModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings(`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+    """
+    model_type = 'llama'
+
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        hidden_act='silu',
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

modelscope/models/nlp/llama/convert_llama_weights_to_hf.py

@@ -0,0 +1,310 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import gc
+import math
+import os
+import shutil
+
+import json
+import torch
+
+from .configuration import LlamaConfig
+from .text_generation import LlamaForTextGeneration
+
+# This file is mainly copied from the llama code of transformers
+INTERMEDIATE_SIZE_MAP = {
+    '7B': 11008,
+    '13B': 13824,
+    '30B': 17920,
+    '65B': 22016,
+}
+NUM_SHARDS = {
+    '7B': 1,
+    '13B': 2,
+    '30B': 4,
+    '65B': 8,
+}
+
+
+def compute_intermediate_size(n):
+    return int(math.ceil(n * 8 / 3) + 255) // 256 * 256
+
+
+def read_json(path):
+    with open(path, 'r') as f:
+        return json.load(f)
+
+
+def write_json(text, path):
+    with open(path, 'w') as f:
+        json.dump(text, f)
+
+
+def write_model(model_path, input_base_path, model_size):
+    os.makedirs(model_path, exist_ok=True)
+    tmp_model_path = os.path.join(model_path, 'tmp')
+    os.makedirs(tmp_model_path, exist_ok=True)
+
+    params = read_json(os.path.join(input_base_path, 'params.json'))
+    num_shards = NUM_SHARDS[model_size]
+    n_layers = params['n_layers']
+    n_heads = params['n_heads']
+    n_heads_per_shard = n_heads // num_shards
+    dim = params['dim']
+    dims_per_head = dim // n_heads
+    base = 10000.0
+    inv_freq = 1.0 / (
+        base**(torch.arange(0, dims_per_head, 2).float() / dims_per_head))
+
+    # permute for sliced rotary
+    def permute(w):
+        return w.view(n_heads, dim // n_heads // 2, 2,
+                      dim).transpose(1, 2).reshape(dim, dim)
+
+    print(f'Fetching all parameters from the checkpoint at {input_base_path}.')
+    # Load weights
+    if model_size == '7B':
+        # Not shared
+        # (The sharded implementation would also work, but this is simpler.)
+        loaded = torch.load(
+            os.path.join(input_base_path, 'consolidated.00.pth'),
+            map_location='cpu')
+    else:
+        # Sharded
+        loaded = [
+            torch.load(
+                os.path.join(input_base_path, f'consolidated.{i:02d}.pth'),
+                map_location='cpu') for i in range(num_shards)
+        ]
+    param_count = 0
+    index_dict = {'weight_map': {}}
+    for layer_i in range(n_layers):
+        filename = f'pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin'
+        if model_size == '7B':
+            # Unsharded
+            state_dict = {
+                f'model.layers.{layer_i}.self_attn.q_proj.weight':
+                permute(loaded[f'layers.{layer_i}.attention.wq.weight']),
+                f'model.layers.{layer_i}.self_attn.k_proj.weight':
+                permute(loaded[f'layers.{layer_i}.attention.wk.weight']),
+                f'model.layers.{layer_i}.self_attn.v_proj.weight':
+                loaded[f'layers.{layer_i}.attention.wv.weight'],
+                f'model.layers.{layer_i}.self_attn.o_proj.weight':
+                loaded[f'layers.{layer_i}.attention.wo.weight'],
+                f'model.layers.{layer_i}.mlp.gate_proj.weight':
+                loaded[f'layers.{layer_i}.feed_forward.w1.weight'],
+                f'model.layers.{layer_i}.mlp.down_proj.weight':
+                loaded[f'layers.{layer_i}.feed_forward.w2.weight'],
+                f'model.layers.{layer_i}.mlp.up_proj.weight':
+                loaded[f'layers.{layer_i}.feed_forward.w3.weight'],
+                f'model.layers.{layer_i}.input_layernorm.weight':
+                loaded[f'layers.{layer_i}.attention_norm.weight'],
+                f'model.layers.{layer_i}.post_attention_layernorm.weight':
+                loaded[f'layers.{layer_i}.ffn_norm.weight'],
+            }
+        else:
+            # Sharded
+            # Note that in the 13B checkpoint, not cloning the two following weights will result in the checkpoint
+            # becoming 37GB instead of 26GB for some reason.
+            state_dict = {
+                f'model.layers.{layer_i}.input_layernorm.weight':
+                loaded[0][f'layers.{layer_i}.attention_norm.weight'].clone(),
+                f'model.layers.{layer_i}.post_attention_layernorm.weight':
+                loaded[0][f'layers.{layer_i}.ffn_norm.weight'].clone(),
+            }
+            state_dict[
+                f'model.layers.{layer_i}.self_attn.q_proj.weight'] = permute(
+                    torch.cat(
+                        [
+                            loaded[i]
+                            [f'layers.{layer_i}.attention.wq.weight'].view(
+                                n_heads_per_shard, dims_per_head, dim)
+                            for i in range(num_shards)
+                        ],
+                        dim=0,
+                    ).reshape(dim, dim))
+            state_dict[
+                f'model.layers.{layer_i}.self_attn.k_proj.weight'] = permute(
+                    torch.cat(
+                        [
+                            loaded[i]
+                            [f'layers.{layer_i}.attention.wk.weight'].view(
+                                n_heads_per_shard, dims_per_head, dim)
+                            for i in range(num_shards)
+                        ],
+                        dim=0,
+                    ).reshape(dim, dim))
+            state_dict[
+                f'model.layers.{layer_i}.self_attn.v_proj.weight'] = torch.cat(
+                    [
+                        loaded[i]
+                        [f'layers.{layer_i}.attention.wv.weight'].view(
+                            n_heads_per_shard, dims_per_head, dim)
+                        for i in range(num_shards)
+                    ],
+                    dim=0,
+                ).reshape(dim, dim)  # noqa
+
+            state_dict[
+                f'model.layers.{layer_i}.self_attn.o_proj.weight'] = torch.cat(
+                    [
+                        loaded[i][f'layers.{layer_i}.attention.wo.weight']
+                        for i in range(num_shards)
+                    ],
+                    dim=1)
+            state_dict[
+                f'model.layers.{layer_i}.mlp.gate_proj.weight'] = torch.cat(
+                    [
+                        loaded[i][f'layers.{layer_i}.feed_forward.w1.weight']
+                        for i in range(num_shards)
+                    ],
+                    dim=0)
+            state_dict[
+                f'model.layers.{layer_i}.mlp.down_proj.weight'] = torch.cat(
+                    [
+                        loaded[i][f'layers.{layer_i}.feed_forward.w2.weight']
+                        for i in range(num_shards)
+                    ],
+                    dim=1)
+            state_dict[
+                f'model.layers.{layer_i}.mlp.up_proj.weight'] = torch.cat(
+                    [
+                        loaded[i][f'layers.{layer_i}.feed_forward.w3.weight']
+                        for i in range(num_shards)
+                    ],
+                    dim=0)
+
+        state_dict[
+            f'model.layers.{layer_i}.self_attn.rotary_emb.inv_freq'] = inv_freq
+        for k, v in state_dict.items():
+            index_dict['weight_map'][k] = filename
+            param_count += v.numel()
+        torch.save(state_dict, os.path.join(tmp_model_path, filename))
+
+    filename = f'pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin'
+    if model_size == '7B':
+        # Unsharded
+        state_dict = {
+            'model.embed_tokens.weight': loaded['tok_embeddings.weight'],
+            'model.norm.weight': loaded['norm.weight'],
+            'lm_head.weight': loaded['output.weight'],
+        }
+    else:
+        state_dict = {
+            'model.norm.weight':
+            loaded[0]['norm.weight'],
+            'model.embed_tokens.weight':
+            torch.cat([
+                loaded[i]['tok_embeddings.weight'] for i in range(num_shards)
+            ],
+                      dim=1),  # noqa
+            'lm_head.weight':
+            torch.cat([loaded[i]['output.weight'] for i in range(num_shards)],
+                      dim=0),
+        }
+
+    for k, v in state_dict.items():
+        index_dict['weight_map'][k] = filename
+        param_count += v.numel()
+    torch.save(state_dict, os.path.join(tmp_model_path, filename))
+
+    # Write configs
+    index_dict['metadata'] = {'total_size': param_count * 2}
+    write_json(index_dict,
+               os.path.join(tmp_model_path, 'pytorch_model.bin.index.json'))
+
+    config = LlamaConfig(
+        hidden_size=dim,
+        intermediate_size=compute_intermediate_size(dim),
+        num_attention_heads=params['n_heads'],
+        num_hidden_layers=params['n_layers'],
+        rms_norm_eps=params['norm_eps'],
+    )
+    config.save_pretrained(tmp_model_path)
+
+    # Make space so we can load the model properly now.
+    del state_dict
+    del loaded
+    gc.collect()
+
+    print('Loading the checkpoint in a Llama model.')
+    model = LlamaForTextGeneration.from_pretrained(
+        tmp_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    # Avoid saving this as part of the config.
+    del model.config._name_or_path
+
+    print('Saving in the Transformers format.')
+    model.save_pretrained(model_path)
+    shutil.rmtree(tmp_model_path)
+
+
+def write_tokenizer(tokenizer_path, input_tokenizer_path):
+    print(f'Fetching the tokenizer from {input_tokenizer_path}.')
+    os.makedirs(tokenizer_path, exist_ok=True)
+    write_json({}, os.path.join(tokenizer_path, 'special_tokens_map.json'))
+    write_json(
+        {
+            'bos_token': '',
+            'eos_token': '',
+            'model_max_length': int(1e30),
+            'tokenizer_class': 'LlamaTokenizer',
+            'unk_token': '',
+        },
+        os.path.join(tokenizer_path, 'tokenizer_config.json'),
+    )
+    shutil.copyfile(input_tokenizer_path,
+                    os.path.join(tokenizer_path, 'tokenizer.model'))
+
+
+def main():
+    """
+    Sample usage:
+
+    ```
+    python src/transformers/models/llama/convert_llama_weights_to_hf.py \
+        --input_dir /path/to/downloaded/llama/weights --model_size 7B --output_dir /output/path
+    ```
+    """
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--input_dir',
+        help=
+        'Location of LLaMA weights, which contains tokenizer.model and model folders',
+    )
+    parser.add_argument(
+        '--model_size',
+        choices=['7B', '13B', '30B', '65B', 'tokenizer_only'],
+    )
+    parser.add_argument(
+        '--output_dir',
+        help='Location to write HF model and tokenizer',
+    )
+    args = parser.parse_args()
+    if args.model_size != 'tokenizer_only':
+        write_model(
+            model_path=args.output_dir,
+            input_base_path=os.path.join(args.input_dir, args.model_size),
+            model_size=args.model_size,
+        )
+    write_tokenizer(
+        tokenizer_path=args.output_dir,
+        input_tokenizer_path=os.path.join(args.input_dir, 'tokenizer.model'),
+    )
+
+
+if __name__ == '__main__':
+    main()

modelscope/models/nlp/llama/text_generation.py

@@ -0,0 +1,177 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from modelscope.metainfo import Models
+from modelscope.models.base import Tensor, TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.outputs import AttentionTextGenerationModelOutput
+from modelscope.utils.constant import Tasks
+from .backbone import LlamaModel, LlamaPreTrainedModel
+
+
+# This file is mainly copied from the llama code of transformers
+@MODELS.register_module(Tasks.text_generation, module_name=Models.llama)
+class LlamaForTextGeneration(LlamaPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r'lm_head.weight']
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+
+        self.lm_head = nn.Linear(
+            config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, AttentionTextGenerationModelOutput]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+
+        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
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits, ) + outputs[1:]
+            return (loss, ) + output if loss is not None else output
+
+        return AttentionTextGenerationModelOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(self,
+                                      input_ids,
+                                      past_key_values=None,
+                                      attention_mask=None,
+                                      inputs_embeds=None,
+                                      **kwargs):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        position_ids = kwargs.get('position_ids', None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {'inputs_embeds': inputs_embeds}
+        else:
+            model_inputs = {'input_ids': input_ids}
+
+        model_inputs.update({
+            'position_ids': position_ids,
+            'past_key_values': past_key_values,
+            'use_cache': kwargs.get('use_cache'),
+            'attention_mask': attention_mask,
+        })
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (tuple(
+                past_state.index_select(0, beam_idx)
+                for past_state in layer_past), )
+        return reordered_past
+
+    def generate(self, inputs: Dict[str, Tensor],
+                 **kwargs) -> Dict[str, Tensor]:
+        return super().generate(**inputs, **kwargs)

modelscope/models/nlp/llama/tokenization.py

@@ -0,0 +1,272 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for LLaMA."""
+import os
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+
+from modelscope.utils.logger import get_logger
+
+# This file is mainly copied from the llama code of transformers
+logger = get_logger(__name__)
+
+VOCAB_FILES_NAMES = {'vocab_file': 'tokenizer.model'}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    'vocab_file': {
+        'hf-internal-testing/llama-tokenizer':
+        'https://huggingface.co/hf-internal-testing/llama-tokenizer/resolve/main/tokenizer.model',
+    },
+    'tokenizer_file': {
+        'hf-internal-testing/llama-tokenizer':
+        'https://huggingface.co/hf-internal-testing/llama-tokenizer/resolve/main/tokenizer_config.json',
+    },
+}
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    'hf-internal-testing/llama-tokenizer': 2048,
+}
+
+
+class LlamaTokenizer(PreTrainedTokenizer):
+    """
+    Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ['input_ids', 'attention_mask']
+
+    def __init__(
+        self,
+        vocab_file,
+        unk_token='<unk>',
+        bos_token='<s>',
+        eos_token='</s>',
+        pad_token=None,
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        add_bos_token=True,
+        add_eos_token=False,
+        clean_up_tokenization_spaces=False,
+        **kwargs,
+    ):
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+        bos_token = AddedToken(
+            bos_token, lstrip=False, rstrip=False) if isinstance(
+                bos_token, str) else bos_token
+        eos_token = AddedToken(
+            eos_token, lstrip=False, rstrip=False) if isinstance(
+                eos_token, str) else eos_token
+        unk_token = AddedToken(
+            unk_token, lstrip=False, rstrip=False) if isinstance(
+                unk_token, str) else unk_token
+        pad_token = AddedToken(
+            pad_token, lstrip=False, rstrip=False) if isinstance(
+                pad_token, str) else pad_token
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            sp_model_kwargs=self.sp_model_kwargs,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+        self.vocab_file = vocab_file
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state['sp_model'] = None
+        return state
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(self.vocab_file)
+
+    @property
+    def vocab_size(self):
+        """Returns vocab size"""
+        return self.sp_model.get_piece_size()
+
+    def get_vocab(self):
+        """Returns vocab as a dict"""
+        vocab = {
+            self.convert_ids_to_tokens(i): i
+            for i in range(self.vocab_size)
+        }
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text):
+        """Returns a tokenized string."""
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        token = self.sp_model.IdToPiece(index)
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        current_sub_tokens = []
+        out_string = ''
+        prev_is_special = False
+        for i, token in enumerate(tokens):
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                if not prev_is_special and i != 0:
+                    out_string += ' '
+                out_string += self.sp_model.decode(current_sub_tokens) + token
+                prev_is_special = True
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+                prev_is_special = False
+        out_string += self.sp_model.decode(current_sub_tokens)
+        return out_string
+
+    def save_vocabulary(self,
+                        save_directory,
+                        filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if not os.path.isdir(save_directory):
+            logger.error(
+                f'Vocabulary path ({save_directory}) should be a directory')
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + '-' if filename_prefix else '')
+            + VOCAB_FILES_NAMES['vocab_file'])
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(
+                out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, 'wb') as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file, )
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = bos_token_id + token_ids_0 + eos_token_id
+
+        if token_ids_1 is not None:
+            output = output + bos_token_id + token_ids_1 + eos_token_id
+
+        return output
+
+    def get_special_tokens_mask(
+            self,
+            token_ids_0: List[int],
+            token_ids_1: Optional[List[int]] = None,
+            already_has_special_tokens: bool = False) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0,
+                token_ids_1=token_ids_1,
+                already_has_special_tokens=True)
+
+        bos_token_id = [1] if self.add_bos_token else []
+        eos_token_id = [1] if self.add_eos_token else []
+
+        if token_ids_1 is None:
+            return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
+        return (bos_token_id +  # noqa
+                ([0] * len(token_ids_0)) + eos_token_id + bos_token_id  # noqa
+                +  # noqa
+                ([0] * len(token_ids_1)) + eos_token_id)  # noqa
+
+    def create_token_type_ids_from_sequences(
+            self,
+            token_ids_0: List[int],
+            token_ids_1: Optional[List[int]] = None) -> List[int]:
+        """
+        Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
+        sequence pair mask has the following format:
+
+        ```
+        0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+        | first sequence    | second sequence |
+        ```
+
+        if token_ids_1 is None, only returns the first portion of the mask (0s).
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of ids.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+        """
+        sep = [self.sep_token_id]
+        cls = [self.cls_token_id]
+
+        if token_ids_1 is None:
+            return len(cls + token_ids_0 + sep) * [0]
+        return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1
+                                                        + sep) * [1]

modelscope/models/nlp/llama/tokenization_fast.py

@@ -0,0 +1,127 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from shutil import copyfile
+from typing import Optional, Tuple
+
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+from transformers.utils import is_sentencepiece_available
+from transformers.utils.versions import require_version
+
+from modelscope.utils.logger import get_logger
+
+# This file is mainly copied from the llama code of transformers
+require_version('tokenizers>=0.13.3')
+
+if is_sentencepiece_available():
+    from .tokenization import LlamaTokenizer
+else:
+    LlamaTokenizer = None
+
+logger = get_logger(__name__)
+VOCAB_FILES_NAMES = {
+    'vocab_file': 'tokenizer.model',
+    'tokenizer_file': 'tokenizer.json'
+}
+
+
+class LlamaTokenizerFast(PreTrainedTokenizerFast):
+    """
+    Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    This uses notably ByteFallback and no normalization.
+
+    ```
+    from transformers import LlamaTokenizerFast
+
+    tokenizer = LlaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
+    tokenizer.encode("Hello this is a test")
+    >>> [1, 15043, 445, 338, 263, 1243]
+    ```
+
+    This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
+    refer to this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .model extension) that
+            contains the vocabulary necessary to instantiate a tokenizer.
+        tokenizer_file (`str`):
+            [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that
+            contains everything needed to load the tokenizer.
+
+        clean_up_tokenization_spaces (`str`, *optional*, defaults to `False`):
+            Wether to cleanup spaces after decoding, cleanup consists in removing potential artifacts like extra
+            spaces.
+
+        bos_token (`str`, *optional*, defaults to `"<s>"`):
+            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+
+        eos_token (`str`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+
+        unk_token (`str`, *optional*, defaults to `"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    slow_tokenizer_class = LlamaTokenizer
+    padding_side = 'left'
+
+    def __init__(
+        self,
+        vocab_file=None,
+        tokenizer_file=None,
+        clean_up_tokenization_spaces=False,
+        unk_token='<unk>',
+        bos_token='<s>',
+        eos_token='</s>',
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            tokenizer_file=tokenizer_file,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            **kwargs,
+        )
+
+        self.vocab_file = vocab_file
+        self.can_save_slow_tokenizer = False if not self.vocab_file else True
+
+    def save_vocabulary(self,
+                        save_directory: str,
+                        filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not self.can_save_slow_tokenizer:
+            raise ValueError(
+                'Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '
+                'tokenizer.')
+
+        if not os.path.isdir(save_directory):
+            logger.error(
+                f'Vocabulary path ({save_directory}) should be a directory')
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + '-' if filename_prefix else '')
+            + VOCAB_FILES_NAMES['vocab_file'])
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+
+        return (out_vocab_file, )

modelscope/models/nlp/task_models/task_model.py

@@ -484,6 +484,16 @@ class EncoderModel(TorchModel):
         self.build_encoder(backbone_cfg)
         if head_cfg.type is not None:
             self.build_head(head_cfg)
+        self.post_init()
+
+    def post_init(self):
+        try:
+            head_keys_to_ignore_on_load_missing = getattr(
+                self.head, '_keys_to_ignore_on_load_missing')
+            for i in head_keys_to_ignore_on_load_missing:
+                self._keys_to_ignore_on_load_missing.append('head.' + i)
+        except Exception:
+            logger.info('head has no _keys_to_ignore_on_load_missing')
 
     def __repr__(self):
         # only log backbone and head name

modelscope/msdatasets/ms_dataset.py

@@ -32,11 +32,6 @@ from modelscope.utils.constant import (DEFAULT_DATASET_NAMESPACE,
 from modelscope.utils.import_utils import is_tf_available, is_torch_available
 from modelscope.utils.logger import get_logger
 
-try:
-    from tensorflow.data import Dataset as TfDataset
-except Exception as e:
-    print(e)
-
 logger = get_logger()
 
 

modelscope/outputs/nlp_outputs.py

@@ -318,6 +318,7 @@ class AttentionTextClassificationModelOutput(TextClassificationModelOutput):
     """
     attentions: Tensor = None
     hidden_states: Tensor = None
+    past_key_values: Tensor = None
 
 
 @dataclass
@@ -351,6 +352,21 @@ class TextGenerationModelOutput(ModelOutputBase):
     loss: Tensor = None
 
 
+@dataclass
+class AttentionTextGenerationModelOutput(TextGenerationModelOutput):
+    """The output class for text generation of attention based models.
+
+    Args:
+        logits (`Tensor`): The logits output of the model. loss (`Tensor`,
+        *optional*) The loss of the model, available when training.
+        hidden_states (`Tensor`, *optional*) Hidden-states of the model at the
+        output of each layer plus the optional initial embedding outputs.
+    """
+    attentions: Tensor = None
+    hidden_states: Tensor = None
+    past_key_values: Tensor = None
+
+
 @dataclass
 class TokenGeneratorOutput(ModelOutputBase):
     """

modelscope/pipelines/base.py

@@ -339,15 +339,18 @@ class Pipeline(ABC):
                         check_input_type(input_type[k], input[k])
             else:
                 raise ValueError(f'invalid input_type definition {input_type}')
-        else:
+        elif not getattr(self, '_input_has_warned', False):
             logger.warning(f'task {task_name} input definition is missing')
+            self._input_has_warned = True
 
     def _check_output(self, input):
         # this attribute is dynamically attached by registry
         # when cls is registered in registry using task name
         task_name = self.group_key
         if task_name not in TASK_OUTPUTS:
-            logger.warning(f'task {task_name} output keys are missing')
+            if not getattr(self, '_output_has_warned', False):
+                logger.warning(f'task {task_name} output keys are missing')
+                self._output_has_warned = True
             return
         output_keys = TASK_OUTPUTS[task_name]
         missing_keys = []

modelscope/preprocessors/nlp/transformers_tokenizer.py

@@ -88,6 +88,11 @@ class NLPTokenizer:
             tokenizer = XLMRobertaTokenizerFast if self.use_fast else XLMRobertaTokenizer
             return tokenizer.from_pretrained(
                 model_dir) if model_dir is not None else tokenizer()
+        elif model_type == Models.llama:
+            from modelscope.models.nlp import LlamaTokenizer, LlamaTokenizerFast
+            tokenizer = LlamaTokenizerFast if self.use_fast else LlamaTokenizer
+            return tokenizer.from_pretrained(
+                model_dir) if model_dir is not None else tokenizer()
 
         assert model_dir is not None
         return AutoTokenizer.from_pretrained(model_dir, use_fast=self.use_fast)

modelscope/utils/checkpoint.py

@@ -512,6 +512,38 @@ def load_task_model_checkpoint(model_to_load,
 
         return retrieved_modules
 
+    def _tie_or_clone_weights(output_embeddings,
+                              input_embeddings,
+                              torchscript=False):
+        if torchscript:
+            output_embeddings.weight = nn.Parameter(
+                input_embeddings.weight.clone())
+        else:
+            output_embeddings.weight = input_embeddings.weight
+
+        if getattr(output_embeddings, 'bias', None) is not None:
+            output_embeddings.bias.data = nn.functional.pad(
+                output_embeddings.bias.data,
+                (
+                    0,
+                    output_embeddings.weight.shape[0]
+                    - output_embeddings.bias.shape[0],
+                ),
+                'constant',
+                0,
+            )
+
+        if hasattr(output_embeddings, 'out_features') and hasattr(
+                input_embeddings, 'num_embeddings'):
+            output_embeddings.out_features = input_embeddings.num_embeddings
+
+    def tie_weights(model, tie_word_embeddings=False):
+        if tie_word_embeddings:
+            output_embeddings = model.head.get_output_embeddings()
+            if output_embeddings is not None:
+                input_embeddings = model.encoder.get_input_embeddings()
+                _tie_or_clone_weights(output_embeddings, input_embeddings)
+
     # TODO Sharded ckpt
     ckpt_file = os.path.join(model_local_dir, ModelFile.TORCH_MODEL_BIN_FILE)
     state_dict = torch.load(ckpt_file, map_location='cpu')
@@ -526,6 +558,9 @@ def load_task_model_checkpoint(model_to_load,
         _fast_init=True,
     )
 
+    if getattr(kwargs.get('head'), 'tie_word_embeddings', False):
+        tie_weights(model_to_load, kwargs.get('head').tie_word_embeddings)
+
     return {
         'model': model_to_load,
         'missing_keys': missing_keys,

modelscope/utils/test_utils.py

@@ -365,8 +365,11 @@ class DistributedTestCase(unittest.TestCase):
               **kwargs):
         from .torch_utils import _find_free_port
         ip = socket.gethostbyname(socket.gethostname())
-        dist_start_cmd = '%s -m torch.distributed.launch --nproc_per_node=%d --master_addr=\'%s\' --master_port=%s' % (
-            sys.executable, num_gpus, ip, _find_free_port())
+        if 'dist_start_cmd' in kwargs:
+            dist_start_cmd = kwargs.pop('dist_start_cmd')
+        else:
+            dist_start_cmd = '%s -m torch.distributed.launch --nproc_per_node=%d ' \
+                             '--master_addr=\'%s\' --master_port=%s' % (sys.executable, num_gpus, ip, _find_free_port())
 
         return self._start(
             dist_start_cmd=dist_start_cmd,

tests/trainers/test_finetune_gpt3_smoke_test.py

@@ -0,0 +1,99 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import shutil
+import tempfile
+import unittest
+
+import torch
+
+from modelscope.metainfo import Trainers
+from modelscope.msdatasets import MsDataset
+from modelscope.trainers import build_trainer
+from modelscope.utils.hub import Config, read_config, snapshot_download
+from modelscope.utils.test_utils import DistributedTestCase, test_level
+
+
+@unittest.skipIf(not torch.cuda.is_available()
+                 or torch.cuda.device_count() <= 1, 'distributed unittest')
+class TestFinetuneGPT3Smoke(DistributedTestCase):
+
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+
+        self.tmp_dir = tempfile.TemporaryDirectory().name
+        if not os.path.exists(self.tmp_dir):
+            os.makedirs(self.tmp_dir)
+
+        self.model_dir = snapshot_download(
+            'damo/nlp_gpt3_text-generation_1.3B')
+        config: Config = read_config(
+            os.path.join(self.model_dir, 'configuration.json'))
+        config.megatron.world_size = 2
+        config.megatron.tensor_model_parallel_size = 2
+        config.dump(os.path.join(self.model_dir, 'configuration.json'))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+        super().tearDown()
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_multi_finetune_portry(self):
+        dist_start_cmd = 'torchrun --nproc_per_node 2'
+        self.start(finetune_poetry, num_gpus=2, dist_start_cmd=dist_start_cmd)
+
+    # TODO: add gpt3 trainer predict unittest
+
+
+def finetune_poetry(work_dir='./gpt3_poetry'):
+    dataset_dict = MsDataset.load('chinese-poetry-collection')
+    train_dataset = dataset_dict['train'].remap_columns({
+        'text1': 'src_txt'
+    }).select(range(20))
+    eval_dataset = dataset_dict['test'].remap_columns({
+        'text1': 'src_txt'
+    }).select(range(20))
+    max_epochs = 2
+    tmp_dir = './gpt3_poetry'
+
+    num_warmup_steps = 100
+
+    def noam_lambda(current_step: int):
+        current_step += 1
+        return min(current_step**(-0.5),
+                   current_step * num_warmup_steps**(-1.5))
+
+    def cfg_modify_fn(cfg):
+        cfg.train.lr_scheduler = {
+            'type': 'LambdaLR',
+            'lr_lambda': noam_lambda,
+            'options': {
+                'by_epoch': False
+            }
+        }
+        cfg.train.optimizer = {'type': 'AdamW', 'lr': 3e-4}
+        cfg.train.dataloader = {'batch_size_per_gpu': 2, 'workers_per_gpu': 1}
+        cfg.train.hooks.append({'type': 'MegatronHook'})
+        cfg.evaluation.dataloader = {
+            'batch_size_per_gpu': 2,
+            'workers_per_gpu': 1
+        }
+        cfg.evaluation.metrics = 'ppl'
+        cfg.num_hidden_layers = 1
+        cfg.model.strict = False
+        return cfg
+
+    kwargs = dict(
+        model='damo/nlp_gpt3_text-generation_1.3B',
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        max_epochs=max_epochs,
+        work_dir=tmp_dir,
+        cfg_modify_fn=cfg_modify_fn)
+
+    # Construct trainer and train
+    trainer = build_trainer(name=Trainers.gpt3_trainer, default_args=kwargs)
+    trainer.train()
+
+
+if __name__ == '__main__':
+    unittest.main()