diff --git a/.gitignore b/.gitignore index 0385322..06828d7 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,8 @@ /.idea *.xml /generate_dataset/build_labels/instance_funcs/llm_instance_datadir +.idea/workspace.xml +/original_data +/data +/.idea +/utils/__pycache__ diff --git a/.idea/workspace.xml b/.idea/workspace.xml index 96fcc12..70faa67 100644 --- a/.idea/workspace.xml +++ b/.idea/workspace.xml @@ -6,10 +6,9 @@ - - - - + + + - + - @@ -190,6 +191,9 @@ + + + @@ -212,12 +216,12 @@ file://$PROJECT_DIR$/train/self_train.py - 164 + 168 file://$PROJECT_DIR$/train/self_train.py - 56 + 60 @@ -232,29 +236,34 @@ file://$PROJECT_DIR$/utils/dataset.py - 248 + 268 file://$PROJECT_DIR$/train/self_train.py - 132 + 136 file://$PROJECT_DIR$/utils/dataset.py - 155 + 164 file://$PROJECT_DIR$/utils/dataset.py - 198 + 216 file://$PROJECT_DIR$/utils/dataset.py - 273 + 290 + + file://$PROJECT_DIR$/main.py + 210 + @@ -264,12 +273,13 @@ + + - - - + + \ No newline at end of file diff --git a/main.py b/main.py index 74956a7..5bd3d70 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,18 @@ import argparse -from collections import defaultdict -from typing import Union, Tuple +from datetime import datetime import os -os.environ["CUDA_VISIBLE_DEVICES"] = "6" + +os.environ["CUDA_VISIBLE_DEVICES"] = "3" + + +from train.ray_train import tune_hyperparameters_ray + + +from datasets import load_dataset + +from test_model import test_model + import torch import yaml @@ -15,7 +24,6 @@ from utils.data_preprocess import read_dataset, select_dataset, preprocess_dataset, split_dataset, \ PreliminaryDataset, read_unlabeled_dataset import utils.tokenization -from module.MT5 import MT5ForConditionalGeneration from transformers import AutoModelForSeq2SeqLM, AutoTokenizer # from transformers import AutoModel @@ -32,13 +40,13 @@ def args_parse(): parser.add_argument('--config', type=str, default='config.yaml', help='config file, 只使用单层的嵌套') - parser.add_argument("--dataset", type=str, default="zcl", + parser.add_argument("--dataset", type=str, default="topv2", choices=["ours", "topv2", "zcl", "zcl_mixed"]) parser.add_argument("--train_dataset_dir", type=str, default="./data/dev", help="train dataset dir") - parser.add_argument("--unlabel_dataset_dir", type=str, default="./data/dev", + parser.add_argument("--unlabel_dataset_path", type=str, default="./data/dev", help="unlabel dataset dir") parser.add_argument("--test_dataset_dir", type=str, default="./data/dev", @@ -47,12 +55,15 @@ def args_parse(): parser.add_argument("--task", type=str, default="self-train", choices=["preliminary", "self-train"], help="省得分文件了") + parser.add_argument("--close_selftrain", type=bool, default=False, + help="如果当前选项为True且task为self-train，则只训练基础模型，不进行自训练") + parser.add_argument("--model_dir", type=str, - default="/data/lbq/models/mt5-base-trained-final-500+500-2-7_again",#"/data/lbq/models/mt5-base-trained-final-500+500-2-7_again", + default="/data/pretrained_models/t5-base",#"/data/pretrained_models/t5-base", #,"/home/lzx/T5-base/model3/mt5-base-trained-final-500+500-2-7_again" help="model dir") - parser.add_argument("--save_dir", type=str, default="/data/lbq/models/mt5-base", + parser.add_argument("--save_dir", type=str, default="/home/lbq/data/semantic_parsing_711/output", help="save dir") parser.add_argument("--experiment_name", type=str, default="Default", choices=["Default"], # 这个比如10样例、100样例等 @@ -61,22 +72,25 @@ def args_parse(): parser.add_argument("--optimizer", type=str, default="Adam", help="optimizer") - parser.add_argument("--lr", type=float, default=3e-5, + parser.add_argument("--lr", type=float, default=1e-5, help="learning rate") + parser.add_argument("--sf_lr", type=float, default=1e-5, + help="learning rate for self-train") + parser.add_argument("--criterion", type=str, default="CrossEntropyLoss", help="criterion") - parser.add_argument("--device", type=str, default="cpu", + parser.add_argument("--device", type=str, default="cuda", help="device") - parser.add_argument("--epoch", type=int, default=3, + parser.add_argument("--epoch", type=int, default=300, help="epoch") - parser.add_argument("--batch_size", type=int, default=1, + parser.add_argument("--batch_size", type=int, default=128, help="batch size") - parser.add_argument("--max_length", type=int, default=512, + parser.add_argument("--max_length", type=int, default=128, help="max length") parser.add_argument("--operator_num", type=int, default=5, @@ -97,13 +111,13 @@ def args_parse(): parser.add_argument("--seed", type=int, default=192, help="random seed") - parser.add_argument("--selftrain_iteration", type=int, default=3, + parser.add_argument("--selftrain_iteration", type=int, default=10, help="self train的迭代次数") parser.add_argument("--selftrain_topk", type=int, default=5, help="self train的topk") - parser.add_argument("--given_model", type=bool, default=True, + parser.add_argument("--given_model", type=bool, default=False, help="是否给定模型，如果是的话就直接训self train") args = parser.parse_args() @@ -143,10 +157,12 @@ def get_dataset(tokenizer, args) -> dict: # 非预实验的时候，有个拆分或者按某个实验标准进行分割的过程，不过这个放在别处也行 for key in dataset: + print('key', key) dataset[key] = tokenizer_dataset(tokenizer, preprocess_dataset(dataset[key])) # self train需要正常训练测试+一个无标签数据集，所以需要额外一个读入或者无标签的读入是从训练集or验证集里拆出来的 - dataset["unlabeled"] = read_unlabeled_dataset(args.unlabel_dataset_dir) + if not args.close_selftrain: + dataset["unlabeled"] = read_unlabeled_dataset(args.unlabel_dataset_path) # 然后self train还有个保存和读入topk的环节，但这个应该也是边训边存 # dataset["unlabeled"] = dataset["unlabeled"]["input_ids"] @@ -171,27 +187,108 @@ def get_criterion(args): raise ValueError(f"Unknown criterion: {args.criterion}") +def get_dataset_path(): + tasks = ['event', 'reminder', 'weather'] + dataset_types = ['top', 'our_data', 'our_add_top'] + exp_settings = ['SPIS25', 'SPIS50', 'full'] + + gpus = 1 # torch count_devices是因为预先生成安全一些 + batch_small = round(8 / gpus) + batch_middle = round(64 / gpus) + batch_large = round(128 / gpus) + + data_dir = r'./original_data' + data_dir = os.path.abspath(data_dir) + + for task in tasks: + for dataset_type in dataset_types: + for exp_setting in exp_settings: + if dataset_type == 'our_data': + data_path = os.path.join(data_dir, task, dataset_type) + else: + data_path = os.path.join(data_dir, task, dataset_type, exp_setting) + + if dataset_type == 'top': + batchsize = batch_small + elif exp_setting == 'full': + batchsize = batch_large + else: + batchsize = batch_middle + + unlabeled_path = os.path.join(data_dir, f"{task}_unlabel_train.tsv") + + yield data_path, batchsize, unlabeled_path + def main(): args = args_parse() - model = MT5ForConditionalGeneration.from_pretrained(args.model_dir) - model.to(args.device) - tokenizer = AutoTokenizer.from_pretrained(args.model_dir) + for path, batchsize, unlabeled_path in get_dataset_path(): + args.train_dataset_dir = path + args.unlabel_dataset_path = unlabeled_path + args.test_dataset_dir = path - dataset = get_dataset(tokenizer, args) + args.batch_size = batchsize - optimizer = get_optimizer(args.optimizer, model, args) - criterion = get_criterion(args) + args.close_selftrain = True + # fixme: args.output (或者就不改了，每次测一下直接覆盖) + # fixme: 开不开self train的都要来一遍 + # batchsize + model = AutoModelForSeq2SeqLM.from_pretrained(args.model_dir) + model.to(args.device) + tokenizer = AutoTokenizer.from_pretrained(args.model_dir) - if args.task == "preliminary": - for op in dataset: - print("算子是", op) - train_model_preliminary(model, optimizer, dataset[op], args) + dataset = get_dataset(tokenizer, args) - elif args.task == "self-train": - train_model_self_train(model, tokenizer, optimizer, dataset, args) + optimizer = get_optimizer(args.optimizer, model, args) + criterion = get_criterion(args) + + if args.task == "preliminary": + for op in dataset: + print("算子是", op) + train_model_preliminary(model, optimizer, dataset[op], args) + + elif args.task == "self-train": + # train_model_self_train(model, tokenizer, optimizer, dataset, args) + tune_hyperparameters_ray(tokenizer, dataset, args, args.save_dir, optimizer, model) + + # break + dataset = preprocess_dataset(load_dataset(path)) + # acc, f1 = test_model(model, tokenizer, dataset, id_sign='main', device='cuda') + # 保存指标到文件 + # _save_metrics_to_file(acc, f1, args, path) + # print((path, acc, f1)) + # break + + +def _save_metrics_to_file(acc, f1, args, dataset_path): + """ + 将acc和f1指标存入文件 + :param acc: 准确率 + :param f1: F1分数 + :param args: 命令行参数+ + :param dataset_path: 当前数据集路径（用于区分不同实验） + """ + # 确保保存目录存在 + os.makedirs(args.save_dir, exist_ok=True) + + # 生成带时间戳的文件名（避免重复） + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + filename = f"metrics_{args.experiment_name}_{timestamp}.txt" + file_path = os.path.join(args.save_dir, filename) + + # 写入指标内容（包含关键上下文信息） + with open(file_path, "a", encoding="utf-8") as f: + f.write(f"===== 实验指标记录 =====\n") + f.write(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") + f.write(f"实验名称: {args.experiment_name}\n") + f.write(f"数据集路径: {dataset_path}\n") + f.write(f"批次大小: {args.batch_size}\n") + f.write(f"任务类型: {args.task}\n") + f.write(f"准确率(Accuracy): {acc:.4f}\n") + f.write(f"F1分数: {f1:.4f}\n") + f.write(f"========================\n\n") if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/module/.MT5.py.swp b/module/.MT5.py.swp deleted file mode 100644 index 5ab1660..0000000 Binary files a/module/.MT5.py.swp and /dev/null differ diff --git a/module/MT5.py b/module/MT5.py deleted file mode 100644 index 879eff5..0000000 --- a/module/MT5.py +++ /dev/null @@ -1,2663 +0,0 @@ -# coding=utf-8 -# Copyright 2020 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" PyTorch mT5 model.""" - -import copy -import math -import os -import warnings -from typing import List, Optional, Tuple, Union - -import torch -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss - -from transformers.activations import ACT2FN -from transformers.modeling_outputs import ( - BaseModelOutput, - BaseModelOutputWithPastAndCrossAttentions, - Seq2SeqLMOutput, - Seq2SeqModelOutput, - Seq2SeqQuestionAnsweringModelOutput, - Seq2SeqSequenceClassifierOutput, -) -from transformers.modeling_utils import PreTrainedModel -from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer -from transformers.utils import ( - DUMMY_INPUTS, - DUMMY_MASK, - add_start_docstrings, - add_start_docstrings_to_model_forward, - is_torch_fx_proxy, - logging, - replace_return_docstrings, -) -from transformers.utils.model_parallel_utils import assert_device_map, get_device_map -from .configuration_mt5 import MT5Config -from transformers import AutoTokenizer - - -logger = logging.get_logger(__name__) - -_CONFIG_FOR_DOC = "MT5Config" -_CHECKPOINT_FOR_DOC = "mt5-small" - - -PARALLELIZE_DOCSTRING = r""" - This is an experimental feature and is a subject to change at a moment's notice. - - Uses a device map to distribute attention modules of the model across several devices. If no device map is given, - it will evenly distribute blocks across all devices. - - Args: - device_map (`Dict[int, list]`, optional, defaults to None): - A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always - automatically mapped to the first device (for esoteric reasons). That means that the first device should - have fewer attention modules mapped to it than other devices. For reference, the mt5 models have the - following number of attention modules: - - - mt5-small: 6 - - mt5-base: 12 - - mt5-large: 24 - - mt5-xl: 24 - - mt5-xxl: 24 - - Example: - - ```python - # Here is an example of a device map on a machine with 4 GPUs using mt5-xl, which has a total of 24 attention modules: - model = MT5ForConditionalGeneration.from_pretrained("mt5-xl") - device_map = { - 0: [0, 1, 2], - 1: [3, 4, 5, 6, 7, 8, 9], - 2: [10, 11, 12, 13, 14, 15, 16], - 3: [17, 18, 19, 20, 21, 22, 23], - } - model.parallelize(device_map) - ``` -""" -DEPARALLELIZE_DOCSTRING = r""" - Moves the model to cpu from a model parallel state. - - Example: - - ```python - # On a 4 GPU machine with mt5-xl: - model = MT5ForConditionalGeneration.from_pretrained("Mt5-xl") - device_map = { - 0: [0, 1, 2], - 1: [3, 4, 5, 6, 7, 8, 9], - 2: [10, 11, 12, 13, 14, 15, 16], - 3: [17, 18, 19, 20, 21, 22, 23], - } - model.parallelize(device_map) # Splits the model across several devices - model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache() - ``` -""" - - -# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->MT5 -class MT5LayerNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - """ - Construct a layernorm module in the MT5 style. No bias and no subtraction of mean. - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - # MT5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean - # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus varience is calculated - # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for - # half-precision inputs is done in fp32 - - 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 - - -# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->MT5 -class MT5DenseActDense(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - self.wi = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.dropout = nn.Dropout(config.dropout_rate) - self.act = ACT2FN[config.dense_act_fn] - - def forward(self, hidden_states): - hidden_states = self.wi(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.dropout(hidden_states) - if ( - isinstance(self.wo.weight, torch.Tensor) - and hidden_states.dtype != self.wo.weight.dtype - and self.wo.weight.dtype != torch.int8 - ): - hidden_states = hidden_states.to(self.wo.weight.dtype) - hidden_states = self.wo(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->MT5 -class MT5DenseGatedActDense(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.dropout = nn.Dropout(config.dropout_rate) - self.act = ACT2FN[config.dense_act_fn] - - def forward(self, hidden_states): - hidden_gelu = self.act(self.wi_0(hidden_states)) - hidden_linear = self.wi_1(hidden_states) - hidden_states = hidden_gelu * hidden_linear - hidden_states = self.dropout(hidden_states) - - # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. - # See https://github.com/huggingface/transformers/issues/20287 - # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` - if ( - isinstance(self.wo.weight, torch.Tensor) - and hidden_states.dtype != self.wo.weight.dtype - and self.wo.weight.dtype != torch.int8 - ): - hidden_states = hidden_states.to(self.wo.weight.dtype) - - hidden_states = self.wo(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->MT5 -class MT5LayerFF(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - if config.is_gated_act: - self.DenseReluDense = MT5DenseGatedActDense(config) - else: - self.DenseReluDense = MT5DenseActDense(config) - - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward(self, hidden_states): - forwarded_states = self.layer_norm(hidden_states) - forwarded_states = self.DenseReluDense(forwarded_states) - hidden_states = hidden_states + self.dropout(forwarded_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->MT5 -class MT5Attention(nn.Module): - def __init__(self, config: MT5Config, has_relative_attention_bias=False): - super().__init__() - self.is_decoder = config.is_decoder - self.has_relative_attention_bias = has_relative_attention_bias - self.relative_attention_num_buckets = config.relative_attention_num_buckets - self.relative_attention_max_distance = config.relative_attention_max_distance - self.d_model = config.d_model - self.key_value_proj_dim = config.d_kv - self.n_heads = config.num_heads - self.dropout = config.dropout_rate - self.inner_dim = self.n_heads * self.key_value_proj_dim - - # Mesh TensorFlow initialization to avoid scaling before softmax - self.q = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.k = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.v = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.o = nn.Linear(self.inner_dim, self.d_model, bias=False) - - if self.has_relative_attention_bias: - self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) - self.pruned_heads = set() - self.gradient_checkpointing = False - - def prune_heads(self, heads): - if len(heads) == 0: - return - heads, index = find_pruneable_heads_and_indices( - heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads - ) - # Prune linear layers - self.q = prune_linear_layer(self.q, index) - self.k = prune_linear_layer(self.k, index) - self.v = prune_linear_layer(self.v, index) - self.o = prune_linear_layer(self.o, index, dim=1) - # Update hyper params - self.n_heads = self.n_heads - len(heads) - self.inner_dim = self.key_value_proj_dim * self.n_heads - self.pruned_heads = self.pruned_heads.union(heads) - - @staticmethod - def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): - """ - Adapted from Mesh Tensorflow: - https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 - - Translate relative position to a bucket number for relative attention. The relative position is defined as - memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to - position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for - small absolute relative_position and larger buckets for larger absolute relative_positions. All relative - positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. - This should allow for more graceful generalization to longer sequences than the model has been trained on - - Args: - relative_position: an int32 Tensor - bidirectional: a boolean - whether the attention is bidirectional - num_buckets: an integer - max_distance: an integer - - Returns: - a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) - """ - relative_buckets = 0 - if bidirectional: - num_buckets //= 2 - relative_buckets += (relative_position > 0).to(torch.long) * num_buckets - relative_position = torch.abs(relative_position) - else: - relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) - # now relative_position is in the range [0, inf) - - # half of the buckets are for exact increments in positions - max_exact = num_buckets // 2 - is_small = relative_position < max_exact - - # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance - relative_position_if_large = max_exact + ( - torch.log(relative_position.float() / max_exact) - / math.log(max_distance / max_exact) - * (num_buckets - max_exact) - ).to(torch.long) - relative_position_if_large = torch.min( - relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) - ) - - relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) - return relative_buckets - - def compute_bias(self, query_length, key_length, device=None): - """Compute binned relative position bias""" - if device is None: - device = self.relative_attention_bias.weight.device - context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] - memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] - relative_position = memory_position - context_position # shape (query_length, key_length) - relative_position_bucket = self._relative_position_bucket( - relative_position, # shape (query_length, key_length) - bidirectional=(not self.is_decoder), - num_buckets=self.relative_attention_num_buckets, - max_distance=self.relative_attention_max_distance, - ) - values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) - values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) - return values - - def forward( - self, - hidden_states, - mask=None, - key_value_states=None, - position_bias=None, - past_key_value=None, - layer_head_mask=None, - query_length=None, - use_cache=False, - output_attentions=False, - ): - """ - Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). - """ - # Input is (batch_size, seq_length, dim) - # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) - # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head) - batch_size, seq_length = hidden_states.shape[:2] - - real_seq_length = seq_length - - if past_key_value is not None: - if len(past_key_value) != 2: - raise ValueError( - f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" - ) - real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length - - key_length = real_seq_length if key_value_states is None else key_value_states.shape[1] - - def shape(states): - """projection""" - return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) - - def unshape(states): - """reshape""" - return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) - - def project(hidden_states, proj_layer, key_value_states, past_key_value): - """projects hidden states correctly to key/query states""" - if key_value_states is None: - # self-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(hidden_states)) - elif past_key_value is None: - # cross-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(key_value_states)) - - if past_key_value is not None: - if key_value_states is None: - # self-attn - # (batch_size, n_heads, key_length, dim_per_head) - hidden_states = torch.cat([past_key_value, hidden_states], dim=2) - elif past_key_value.shape[2] != key_value_states.shape[1]: - # checking that the `sequence_length` of the `past_key_value` is the same as - # the provided `key_value_states` to support prefix tuning - # cross-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(key_value_states)) - else: - # cross-attn - hidden_states = past_key_value - return hidden_states - - # get query states - query_states = shape(self.q(hidden_states)) # (batch_size, n_heads, seq_length, dim_per_head) - - # get key/value states - key_states = project( - hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None - ) - value_states = project( - hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None - ) - - # compute scores - scores = torch.matmul( - query_states, key_states.transpose(3, 2) - ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 - - if position_bias is None: - if not self.has_relative_attention_bias: - position_bias = torch.zeros( - (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype - ) - if self.gradient_checkpointing and self.training: - position_bias.requires_grad = True - else: - position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device) - - # if key and values are already calculated - # we want only the last query position bias - if past_key_value is not None: - position_bias = position_bias[:, :, -hidden_states.size(1) :, :] - - if mask is not None: - position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length) - - if self.pruned_heads: - mask = torch.ones(position_bias.shape[1]) - mask[list(self.pruned_heads)] = 0 - position_bias_masked = position_bias[:, mask.bool()] - else: - position_bias_masked = position_bias - - scores += position_bias_masked - attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as( - scores - ) # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.dropout( - attn_weights, p=self.dropout, training=self.training - ) # (batch_size, n_heads, seq_length, key_length) - - # Mask heads if we want to - if layer_head_mask is not None: - attn_weights = attn_weights * layer_head_mask - - attn_output = unshape(torch.matmul(attn_weights, value_states)) # (batch_size, seq_length, dim) - attn_output = self.o(attn_output) - - present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None - outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) - - if output_attentions: - outputs = outputs + (attn_weights,) - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->MT5 -class MT5LayerSelfAttention(nn.Module): - def __init__(self, config, has_relative_attention_bias=False): - super().__init__() - self.SelfAttention = MT5Attention(config, has_relative_attention_bias=has_relative_attention_bias) - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward( - self, - hidden_states, - attention_mask=None, - position_bias=None, - layer_head_mask=None, - past_key_value=None, - use_cache=False, - output_attentions=False, - ): - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( - normed_hidden_states, - mask=attention_mask, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states = hidden_states + self.dropout(attention_output[0]) - outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->MT5 -class MT5LayerCrossAttention(nn.Module): - def __init__(self, config): - super().__init__() - self.EncDecAttention = MT5Attention(config, has_relative_attention_bias=False) - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward( - self, - hidden_states, - key_value_states, - attention_mask=None, - position_bias=None, - layer_head_mask=None, - past_key_value=None, - use_cache=False, - query_length=None, - output_attentions=False, - ): - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( - normed_hidden_states, - mask=attention_mask, - key_value_states=key_value_states, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - query_length=query_length, - output_attentions=output_attentions, - ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5Block with T5->MT5 -class MT5Block(nn.Module): - def __init__(self, config, has_relative_attention_bias=False): - super().__init__() - self.is_decoder = config.is_decoder - self.layer = nn.ModuleList() - self.layer.append(MT5LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias)) - if self.is_decoder: - self.layer.append(MT5LayerCrossAttention(config)) - - self.layer.append(MT5LayerFF(config)) - - def forward( - self, - hidden_states, - attention_mask=None, - position_bias=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - encoder_decoder_position_bias=None, - layer_head_mask=None, - cross_attn_layer_head_mask=None, - past_key_value=None, - use_cache=False, - output_attentions=False, - return_dict=True, - ): - if past_key_value is not None: - if not self.is_decoder: - logger.warning("`past_key_values` is passed to the encoder. Please make sure this is intended.") - expected_num_past_key_values = 2 if encoder_hidden_states is None else 4 - - if len(past_key_value) != expected_num_past_key_values: - raise ValueError( - f"There should be {expected_num_past_key_values} past states. " - f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}" - f"Got {len(past_key_value)} past key / value states" - ) - - self_attn_past_key_value = past_key_value[:2] - cross_attn_past_key_value = past_key_value[2:] - else: - self_attn_past_key_value, cross_attn_past_key_value = None, None - - self_attention_outputs = self.layer[0]( - hidden_states, - attention_mask=attention_mask, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=self_attn_past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states, present_key_value_state = self_attention_outputs[:2] - attention_outputs = self_attention_outputs[2:] # Keep self-attention outputs and relative position weights - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - do_cross_attention = self.is_decoder and encoder_hidden_states is not None - if do_cross_attention: - # the actual query length is unknown for cross attention - # if using past key value states. Need to inject it here - if present_key_value_state is not None: - query_length = present_key_value_state[0].shape[2] - else: - query_length = None - - cross_attention_outputs = self.layer[1]( - hidden_states, - key_value_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - position_bias=encoder_decoder_position_bias, - layer_head_mask=cross_attn_layer_head_mask, - past_key_value=cross_attn_past_key_value, - query_length=query_length, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states = cross_attention_outputs[0] - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - # Combine self attn and cross attn key value states - if present_key_value_state is not None: - present_key_value_state = present_key_value_state + cross_attention_outputs[1] - - # Keep cross-attention outputs and relative position weights - attention_outputs = attention_outputs + cross_attention_outputs[2:] - - # Apply Feed Forward layer - hidden_states = self.layer[-1](hidden_states) - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - outputs = (hidden_states,) - - if use_cache: - outputs = outputs + (present_key_value_state,) + attention_outputs - else: - outputs = outputs + attention_outputs - - return outputs # hidden-states, present_key_value_states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) - - -def load_tf_weights_in_mt5(model, config, tf_checkpoint_path): - """Load tf checkpoints in a pytorch model.""" - try: - import re - - import numpy as np - import tensorflow as tf - except ImportError: - logger.error( - "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " - "https://www.tensorflow.org/install/ for installation instructions." - ) - raise - tf_path = os.path.abspath(tf_checkpoint_path) - logger.info(f"Converting TensorFlow checkpoint from {tf_path}") - # Load weights from TF model - init_vars = tf.train.list_variables(tf_path) - names = [] - tf_weights = {} - for name, shape in init_vars: - logger.info(f"Loading TF weight {name} with shape {shape}") - array = tf.train.load_variable(tf_path, name) - names.append(name) - tf_weights[name] = array - - for txt_name in names: - name = txt_name.split("/") - # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v - # which are not required for using pretrained model - if any( - n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] - for n in name - ): - logger.info(f"Skipping {'/'.join(name)}") - tf_weights.pop(txt_name, None) - continue - if "_slot_" in name[-1]: - logger.info(f"Skipping {'/'.join(name)}") - tf_weights.pop(txt_name, None) - continue - pointer = model - array = tf_weights[txt_name] - - for m_name in name: - if re.fullmatch(r"[A-Za-z]+_\d+", m_name): - scope_names = re.split(r"_(\d+)", m_name) - else: - scope_names = [m_name] - if scope_names[0] in ["kernel", "scale", "embedding"]: - pointer = getattr(pointer, "weight") - elif scope_names[0] == "self_attention": - pointer = getattr(pointer, "layer") - pointer = pointer[0] - elif scope_names[0] == "enc_dec_attention": - pointer = getattr(pointer, "layer") - pointer = pointer[1] - elif scope_names[0] == "dense_relu_dense": - pointer = getattr(pointer, "layer") - pointer = pointer[2] - elif scope_names[0] == "rms_norm": - if hasattr(pointer, "layer_norm"): - pointer = getattr(pointer, "layer_norm") - elif hasattr(pointer, "final_layer_norm"): - pointer = getattr(pointer, "final_layer_norm") - elif scope_names[0] == "scale": - pointer = getattr(pointer, "weight") - elif scope_names[0] == "output_bias" or scope_names[0] == "beta": - pointer = getattr(pointer, "bias") - elif scope_names[0] == "squad": - pointer = getattr(pointer, "classifier") - elif scope_names[0] == "decoder" and name[1] == "logits": - continue - elif scope_names[0] == "logits": - pointer = getattr(pointer, "lm_head") - elif scope_names[0] == "wi" and len(scope_names) > 1 and scope_names[1].isdigit(): - pointer = getattr(pointer, f"wi_{scope_names[1]}") - continue - else: - try: - pointer = getattr(pointer, scope_names[0]) - except AttributeError: - logger.info(f"Skipping {'/'.join(name)}") - continue - if len(scope_names) >= 2: - num = int(scope_names[1]) - pointer = pointer[num] - if scope_names[0] not in ["kernel", "scale", "embedding"]: - pointer = getattr(pointer, "weight") - if scope_names[0] != "embedding": - logger.info(f"Transposing numpy weight of shape {array.shape} for {name}") - array = np.transpose(array) - try: - assert ( - pointer.shape == array.shape - ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched" - except AssertionError as e: - e.args += (pointer.shape, array.shape) - raise - logger.info(f"Initialize PyTorch weight {name}") - pointer.data = torch.from_numpy(array.astype(np.float32)) - tf_weights.pop(txt_name, None) - - logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}.") - return model - - -# Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->MT5 -class MT5ClassificationHead(nn.Module): - """Head for sentence-level classification tasks.""" - - def __init__(self, config: MT5Config): - super().__init__() - self.dense = nn.Linear(config.d_model, config.d_model) - self.dropout = nn.Dropout(p=config.classifier_dropout) - self.out_proj = nn.Linear(config.d_model, config.num_labels) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.dropout(hidden_states) - hidden_states = self.dense(hidden_states) - hidden_states = torch.tanh(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.out_proj(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel with T5->MT5, t5->mt5 -class MT5PreTrainedModel(PreTrainedModel): - """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. - """ - - config_class = MT5Config - load_tf_weights = load_tf_weights_in_mt5 - base_model_prefix = "transformer" - is_parallelizable = True - supports_gradient_checkpointing = True - _no_split_modules = ["MT5Block"] - _keep_in_fp32_modules = ["wo"] - - @property - def dummy_inputs(self): - input_ids = torch.tensor(DUMMY_INPUTS) - input_mask = torch.tensor(DUMMY_MASK) - dummy_inputs = { - "decoder_input_ids": input_ids, - "input_ids": input_ids, - "decoder_attention_mask": input_mask, - } - return dummy_inputs - - def _init_weights(self, module): - """Initialize the weights""" - factor = self.config.initializer_factor # Used for testing weights initialization - if isinstance(module, MT5LayerNorm): - module.weight.data.fill_(factor * 1.0) - elif isinstance( - module, - (MT5Model, MT5ForConditionalGeneration, MT5EncoderModel, MT5ForQuestionAnswering), - ): - # Mesh TensorFlow embeddings initialization - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624 - module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) - if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: - module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0) - if hasattr(module, "qa_outputs"): - module.qa_outputs.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - module.qa_outputs.bias.data.zero_() - elif isinstance(module, MT5ClassificationHead): - module.dense.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.dense, "bias") and module.dense.bias is not None: - module.dense.bias.data.zero_() - module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: - module.out_proj.bias.data.zero_() - elif isinstance(module, MT5DenseActDense): - # Mesh TensorFlow FF initialization - # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56 - # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89 - module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi, "bias") and module.wi.bias is not None: - module.wi.bias.data.zero_() - module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) - if hasattr(module.wo, "bias") and module.wo.bias is not None: - module.wo.bias.data.zero_() - elif isinstance(module, MT5DenseGatedActDense): - module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: - module.wi_0.bias.data.zero_() - module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: - module.wi_1.bias.data.zero_() - module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) - if hasattr(module.wo, "bias") and module.wo.bias is not None: - module.wo.bias.data.zero_() - elif isinstance(module, MT5Attention): - # Mesh TensorFlow attention initialization to avoid scaling before softmax - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136 - d_model = self.config.d_model - key_value_proj_dim = self.config.d_kv - n_heads = self.config.num_heads - module.q.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5)) - module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) - module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) - module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) - if module.has_relative_attention_bias: - module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5)) - - def _shift_right(self, input_ids): - decoder_start_token_id = self.config.decoder_start_token_id - pad_token_id = self.config.pad_token_id - - if decoder_start_token_id is None: - raise ValueError( - "self.model.config.decoder_start_token_id has to be defined. In MT5 it is usually set to the pad_token_id. " - "See MT5 docs for more information." - ) - - # shift inputs to the right - if is_torch_fx_proxy(input_ids): - # Item assignment is not supported natively for proxies. - shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id) - shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) - else: - shifted_input_ids = input_ids.new_zeros(input_ids.shape) - shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() - shifted_input_ids[..., 0] = decoder_start_token_id - - if pad_token_id is None: - raise ValueError("self.model.config.pad_token_id has to be defined.") - # replace possible -100 values in labels by `pad_token_id` - shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) - - return shifted_input_ids - - -# Copied from transformers.models.t5.modeling_t5.T5Stack with T5->MT5 -class MT5Stack(MT5PreTrainedModel): - def __init__(self, config, embed_tokens=None): - super().__init__(config) - - self.embed_tokens = embed_tokens - self.is_decoder = config.is_decoder - - self.block = nn.ModuleList( - [MT5Block(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)] - ) - self.final_layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - # Initialize weights and apply final processing - self.post_init() - # Model parallel - self.model_parallel = False - self.device_map = None - self.gradient_checkpointing = False - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - def parallelize(self, device_map=None): - warnings.warn( - "`MT5Stack.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" - " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," - " 'block.1': 1, ...}", - FutureWarning, - ) - # Check validity of device_map - self.device_map = ( - get_device_map(len(self.block), range(torch.cuda.device_count())) if device_map is None else device_map - ) - assert_device_map(self.device_map, len(self.block)) - self.model_parallel = True - self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys())) - self.last_device = "cuda:" + str(max(self.device_map.keys())) - # Load onto devices - for k, v in self.device_map.items(): - for layer in v: - cuda_device = "cuda:" + str(k) - self.block[layer] = self.block[layer].to(cuda_device) - - # Set embed_tokens to first layer - self.embed_tokens = self.embed_tokens.to(self.first_device) - # Set final layer norm to last device - self.final_layer_norm = self.final_layer_norm.to(self.last_device) - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.model_parallel = False - self.device_map = None - self.first_device = "cpu" - self.last_device = "cpu" - for i in range(len(self.block)): - self.block[i] = self.block[i].to("cpu") - self.embed_tokens = self.embed_tokens.to("cpu") - self.final_layer_norm = self.final_layer_norm.to("cpu") - torch.cuda.empty_cache() - - def get_input_embeddings(self): - return self.embed_tokens - - def set_input_embeddings(self, new_embeddings): - self.embed_tokens = new_embeddings - - def forward( - self, - input_ids=None, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - inputs_embeds=None, - head_mask=None, - cross_attn_head_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - # Model parallel - if self.model_parallel: - torch.cuda.set_device(self.first_device) - self.embed_tokens = self.embed_tokens.to(self.first_device) - use_cache = use_cache if use_cache is not None else self.config.use_cache - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is not None and inputs_embeds is not None: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError( - f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" - ) - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") - - if inputs_embeds is None: - if self.embed_tokens is None: - raise ValueError("You have to initialize the model with valid token embeddings") - inputs_embeds = self.embed_tokens(input_ids) - # print(input_ids) - batch_size, seq_length = input_shape - - # required mask seq length can be calculated via length of past - mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length - - if use_cache is True: - if not self.is_decoder: - raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") - - if attention_mask is None: - attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) - if self.is_decoder and encoder_attention_mask is None and encoder_hidden_states is not None: - encoder_seq_length = encoder_hidden_states.shape[1] - encoder_attention_mask = torch.ones( - batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long - ) - - # initialize past_key_values with `None` if past does not exist - if past_key_values is None: - past_key_values = [None] * len(self.block) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape) - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.is_decoder and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = None - - 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 - - # Prepare head mask if needed - head_mask = self.get_head_mask(head_mask, self.config.num_layers) - cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers) - present_key_value_states = () if use_cache else None - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - all_cross_attentions = () if (output_attentions and self.is_decoder) else None - position_bias = None - encoder_decoder_position_bias = None - - hidden_states = self.dropout(inputs_embeds) - - for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)): - layer_head_mask = head_mask[i] - cross_attn_layer_head_mask = cross_attn_head_mask[i] - # Model parallel - if self.model_parallel: - torch.cuda.set_device(hidden_states.device) - # Ensure that attention_mask is always on the same device as hidden_states - if attention_mask is not None: - attention_mask = attention_mask.to(hidden_states.device) - if position_bias is not None: - position_bias = position_bias.to(hidden_states.device) - if encoder_hidden_states is not None: - encoder_hidden_states = encoder_hidden_states.to(hidden_states.device) - if encoder_extended_attention_mask is not None: - encoder_extended_attention_mask = encoder_extended_attention_mask.to(hidden_states.device) - if encoder_decoder_position_bias is not None: - encoder_decoder_position_bias = encoder_decoder_position_bias.to(hidden_states.device) - if layer_head_mask is not None: - layer_head_mask = layer_head_mask.to(hidden_states.device) - if cross_attn_layer_head_mask is not None: - cross_attn_layer_head_mask = cross_attn_layer_head_mask.to(hidden_states.device) - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer_module.forward, - hidden_states, - extended_attention_mask, - position_bias, - encoder_hidden_states, - encoder_extended_attention_mask, - encoder_decoder_position_bias, - layer_head_mask, - cross_attn_layer_head_mask, - None, # past_key_value is always None with gradient checkpointing - use_cache, - output_attentions, - ) - - else: - layer_outputs = layer_module( - hidden_states, - attention_mask=extended_attention_mask, - position_bias=position_bias, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - encoder_decoder_position_bias=encoder_decoder_position_bias, - layer_head_mask=layer_head_mask, - cross_attn_layer_head_mask=cross_attn_layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - - # layer_outputs is a tuple with: - # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) - if use_cache is False: - layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:] - - hidden_states, present_key_value_state = layer_outputs[:2] - - # We share the position biases between the layers - the first layer store them - # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), - # (cross-attention position bias), (cross-attention weights) - position_bias = layer_outputs[2] - if self.is_decoder and encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[4 if output_attentions else 3] - # append next layer key value states - if use_cache: - present_key_value_states = present_key_value_states + (present_key_value_state,) - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[3],) - if self.is_decoder: - all_cross_attentions = all_cross_attentions + (layer_outputs[5],) - - # Model Parallel: If it's the last layer for that device, put things on the next device - if self.model_parallel: - for k, v in self.device_map.items(): - if i == v[-1] and "cuda:" + str(k) != self.last_device: - hidden_states = hidden_states.to("cuda:" + str(k + 1)) - - hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.dropout(hidden_states) - - # Add last layer - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - present_key_value_states, - all_hidden_states, - all_attentions, - all_cross_attentions, - ] - if v is not None - ) - # print("all_hidden_states", all_hidden_states.shape) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=present_key_value_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - cross_attentions=all_cross_attentions, - ) - - -MT5_START_DOCSTRING = r""" - - The MT5 model was proposed in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text - Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan - Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. It's an encoder decoder transformer pre-trained in a - text-to-text denoising generative setting. - - This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the - library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads - etc.) - - This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. - Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage - and behavior. - - Parameters: - config ([`MT5Config`]): Model configuration class with all the parameters of the model. - Initializing with a config file does not load the weights associated with the model, only the - configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. -""" - -MT5_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you - should be able to pad the inputs on both the right and the left. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for detail. - - [What are input IDs?](../glossary#input-ids) - - To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). - attention_mask (`torch.FloatTensor` 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) - decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): - Indices of decoder input sequence tokens in the vocabulary. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for details. - - [What are decoder input IDs?](../glossary#decoder-input-ids) - - MT5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` - is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). - - To know more on how to prepare `decoder_input_ids` for pretraining take a look at [MT5 - Training](./mt5#training). - decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): - Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also - be used by default. - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules in the encoder. Mask values selected in `[0, - 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, - 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in - `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): - Tuple consists of (`last_hidden_state`, `optional`: *hidden_states*, `optional`: *attentions*) - `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden states at - the output of the last layer of the encoder. Used in the cross-attention of the decoder. - past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - 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. - decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): - Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded - representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be - input (see `past_key_values`). This is useful if you want more control over how to convert - `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. - - If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value - of `inputs_embeds`. - - 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. -""" - -MT5_ENCODER_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you - should be able to pad the inputs on both the right and the left. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for detail. - - To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). - attention_mask (`torch.FloatTensor` 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) - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - 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. - 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. -""" - -# Warning message for FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask -__HEAD_MASK_WARNING_MSG = """ -The input argument `head_mask` was split into two arguments `head_mask` and `decoder_head_mask`. Currently, -`decoder_head_mask` is set to copy `head_mask`, but this feature is deprecated and will be removed in future versions. -If you do not want to use any `decoder_head_mask` now, please set `decoder_head_mask = torch.ones(num_layers, -num_heads)`. -""" - - -@add_start_docstrings( - "The bare MT5 Model transformer outputting raw hidden-states without any specific head on top.", - MT5_START_DOCSTRING, -) -class MT5Model(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5Model, AutoTokenizer - - >>> model = MT5Model.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> summary = "Weiter Verhandlung in Syrien." - >>> inputs = tokenizer(article, return_tensors="pt") - >>> labels = tokenizer(text_target=summary, return_tensors="pt") - - >>> outputs = model(input_ids=inputs["input_ids"], decoder_input_ids=labels["input_ids"]) - >>> hidden_states = outputs.last_hidden_state - ```""" - model_type = "mt5" - config_class = MT5Config - _keys_to_ignore_on_load_missing = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5Model.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - self.decoder = MT5Stack(decoder_config, self.shared) - - # Initialize weights and apply final processing - self.post_init() - - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5Model.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" - " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'encoder.block.0':" - " 0, 'encoder.block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.decoder.parallelize(self.device_map) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5Model.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.decoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.decoder = self.decoder.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5Model.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_decoder - def get_decoder(self): - return self.decoder - - # Copied from transformers.models.t5.modeling_t5.T5Model._prune_heads - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.layer[layer].attention.prune_heads(heads) - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5Model.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - inputs_embeds: Optional[torch.Tensor] = None, - decoder_inputs_embeds: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], Seq2SeqModelOutput]: - r""" - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MT5Model - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5Model.from_pretrained("mt5-small") - - >>> input_ids = tokenizer( - ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 - - >>> # preprocess: Prepend decoder_input_ids with start token which is pad token for MT5Model. - >>> # This is not needed for torch's MT5ForConditionalGeneration as it does this internally using labels arg. - >>> decoder_input_ids = model._shift_right(decoder_input_ids) - - >>> # forward pass - >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) - >>> last_hidden_states = outputs.last_hidden_state - ```""" - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - hidden_states = hidden_states.to(self.decoder.first_device) - if decoder_input_ids is not None: - decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) - if attention_mask is not None: - attention_mask = attention_mask.to(self.decoder.first_device) - if decoder_attention_mask is not None: - decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - if not return_dict: - return decoder_outputs + encoder_outputs - - return Seq2SeqModelOutput( - last_hidden_state=decoder_outputs.last_hidden_state, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) - - -import torch.nn.functional as F -class PointAttention(nn.Module): - def __init__(self): - super().__init__() - self.hidden_size = 768 - # self.linear_layer = nn.Linear(250778, 1) - - - def forward(self,hidden1_states,hidden2_states,lm_logits,tags = [[0,0,1,2,-1,3,4,5,-1]]): - taglist = [] - finalBiasList = [] - sentence_len = hidden1_states.shape[1] - output_sentence_len = hidden2_states.shape[1] - for tag in tags: - vect,bias = merge_weights(tag,sentence_len) - # print('sentence_len', sentence_len) - taglist.append(vect) - biasList = [] - for i in range(output_sentence_len): - biasList.append(bias) - finalBiasList.append(biasList) - merge = torch.tensor(taglist, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) # 8 512[实则300] 512[指明需要的行] - finalBias = torch.tensor(finalBiasList, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) - # 8 512[实则300] 728 - - # print(hidden1_states.shape) #[8, 512, 768] - # 尝试给merge扩维 - TopK_Shape = hidden1_states.shape[0] - Shape = merge.shape[0] - k = TopK_Shape/Shape - if(k>1): - merge1 = merge.repeat_interleave(int(k),dim=0) - finalBias1 = finalBias.repeat_interleave(int(k),dim=0) - else: - merge1 = merge - finalBias1 = finalBias - # merge = merge.repeat_interleave(,dim=0) - attn_weights = torch.matmul(merge1,hidden1_states) # hidden1_states1 8 512【300】 728 - - attn_weights1 = torch.matmul(hidden2_states, attn_weights.transpose(1, 2))+finalBias1 # hidden2_states 8 512 768 attn_weights 8 512 768 attn_weights1 8 512 512 - score = torch.cat([lm_logits,attn_weights1], dim=2) - # probabilities = F.softmax(score, dim=-1) - # lm_logits[:, :, -30:] = attn_weights1 - return score - - -def get_vect(i,j,sentence_len=512): - vect = [] - for t in range(sentence_len): - vect.append(0) - t = i - count = j-i - while(t', '', '',''} - -def isdigit(s): - if s.replace('.', '', 1).isdigit() or (s.find(".")!=-1 and s.replace('.', '', 1) == ""): - return True - else: - return False -def label_words(original_sentence,lst): - i = 0 - indices = [] - num = 0 - # print(lst) - while i < len(lst): - if isdigit(lst[i]): - # 当前元素是数字字符串，检查下一个元素 - number_str = lst[i] - indices.append(num) - i += 1 - while i < len(lst) and isdigit(lst[i]): - # 下一个元素也是数字字符串，进行拼接 - number_str += lst[i] - indices.append(num) - i += 1 - num+=1 - elif lst[i] in special_tokens: # 特殊字符检查 - indices.append(-1) - num+=1 - i += 1 - else: - indices.append(num) - num+=1 - i += 1 - # print(indices) - return indices - -# 为列表中的每个元素分配序号 -# 中文没考虑数字 -# def label_words(original_sentence,word_list): -# indices = [] -# num = 0 -# for index, word in enumerate(word_list): -# if word in special_tokens: # 特殊字符检查 -# indices.append(-1) -# num += 1 -# else: -# indices.append(index-num) - -# return indices - -# 英文拼接 -# def label_words(original_sentence,tokenized_words): -# labels = [] -# label = 0 -# # word_index = 0 -# current_word = "" -# lastword = "" -# # 判断是不是最后一个句号 -# period = original_sentence.count(".") -# period_num = 1 -# # print(tokenized_words) -# for word in tokenized_words: -# # 如果当前单词是非字母（如空格或特殊符号），单独分配一个标签 -# pmFlag = False -# if not word.isalpha() and not word.isdigit() and word.find("'") == -1 and not is_valid_time_format(word): -# # .不能丢。如果是句子中的最后一个，则... 如果不是句子中的最后一个。则... -# if word == "." and period_num < period: -# period_num += 1 -# elif word == "." and period_num == period: -# labels.append(label+1) -# continue -# else: -# labels.append(-1) -# lastword = word -# continue - -# # 累加拆分的单词片段 -# current_word += word - -# if lastword.isdigit() and original_sentence.find(lastword + word) != -1 and (word.lower() == "pm" or word.lower() == "am"): -# pmFlag = True -# # 检查当前累加的单词片段是否与原句中的单词相匹配 -# # print(original_sentence.split()) - -# if current_word in original_sentence.split() or pmFlag == True: -# if pmFlag == True: -# label+=1 - -# if current_word in original_sentence.split(): -# current_word = "" - -# labels.append(label) -# # print(current_word) -# else: -# labels.append(label) - -# # 如果当前累加的片段已经形成了完整的单词，则增加标签值 -# if current_word == "": -# label += 1 -# lastword = word -# # return changeModel(labels) -# return labels - -def label_tokenized_words_corrected(tokenizer,token_list): - # tokenizer = AutoTokenizer.from_pretrained("/home/lzx/T5-base/tokenizer1/") - original_sentences = [tokenizer.decode(t[:-1]) for t in token_list] #token_list[0]是length, 512 - lsts = [] - for tokens in token_list: - decoded_words = [tokenizer.decode(token_id) for token_id in tokens] #可能需要括号[token_id]，这个维度有点奇怪 - lsts.append(decoded_words) - # print(lsts) - # 遍历分词后的单词 - llables = [] - for i, tokenized_words in enumerate(lsts): - original_sentence = original_sentences[i] - labels = label_words(original_sentence,tokenized_words) - llables.append(labels) - return llables - -@add_start_docstrings("""MT5 Model with a `language modeling` head on top.""", MT5_START_DOCSTRING) -class MT5ForConditionalGeneration(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5ForConditionalGeneration, AutoTokenizer - - >>> model = MT5ForConditionalGeneration.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> summary = "Weiter Verhandlung in Syrien." - >>> inputs = tokenizer(article, text_target=summary, return_tensors="pt") - - >>> outputs = model(**inputs) - >>> loss = outputs.loss - ```""" - - model_type = "mt5" - config_class = MT5Config - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.model_dim = config.d_model - - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - # print("decoder_config", decoder_config) - self.decoder = MT5Stack(decoder_config, self.shared) - - # -512 - self.lm_head = nn.Linear(config.d_model, config.vocab_size-128, bias=False) #-512 - # Initialize weights and apply final processing - self.post_init() - self.pointNet = PointAttention() - # /home/lzx/T5-base-lora/tokenizer2/ - # /home/lzx/T5-base/model_cl_multi/mt5-base-trained-final-save - self.tokenizer = AutoTokenizer.from_pretrained("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again") #("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again") # ("/home/lzx/T5-base/model3/mt5-base-trained-final-500+500-2-7_again")# - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5ForConditionalGeneration.parallelize` is deprecated and will be removed in v5 of Transformers, you" - " should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also" - " provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance" - " {'encoder.block.0': 0, 'encoder.block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.decoder.parallelize(self.device_map) - self.lm_head = self.lm_head.to(self.decoder.first_device) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.decoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.decoder = self.decoder.to("cpu") - self.lm_head = self.lm_head.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.set_output_embeddings - def set_output_embeddings(self, new_embeddings): - self.lm_head = new_embeddings - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_output_embeddings - def get_output_embeddings(self): - return self.lm_head - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_decoder - def get_decoder(self): - return self.decoder - - def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding: - """ - Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`. - - Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method. - - Arguments: - new_num_tokens (`int`, *optional*): - The number of new tokens in the embedding matrix. Increasing the size will add newly initialized - vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just - returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything. - - Return: - `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model. - """ - model_embeds = self._resize_token_embeddings(new_num_tokens) - if new_num_tokens is None: - return model_embeds - - # Update base model and current model config - self.config.vocab_size = new_num_tokens - self.vocab_size = new_num_tokens - - # Tie weights again if needed - self.tie_weights() - - return model_embeds - - def _resize_token_embeddings(self, new_num_tokens): - old_embeddings = self.get_input_embeddings() - new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens) - self.set_input_embeddings(new_embeddings) - - # if word embeddings are not tied, make sure that lm head is resized as well - if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings: - old_lm_head = self.get_output_embeddings() - #512 - new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens-128)#-128) - self.set_output_embeddings(new_lm_head) - - return self.get_input_embeddings() - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - labels1: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., - config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for - labels in `[0, ..., config.vocab_size]` - - Returns: - - Examples: - - ```python - >>> from transformers import AutoTokenizer, MT5ForConditionalGeneration - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5ForConditionalGeneration.from_pretrained("mt5-small") - - >>> # training - >>> input_ids = tokenizer("The walks in park", return_tensors="pt").input_ids - >>> labels = tokenizer(" cute dog the ", return_tensors="pt").input_ids - >>> outputs = model(input_ids=input_ids, labels=labels) - >>> loss = outputs.loss - >>> logits = outputs.logits - - >>> # inference - >>> input_ids = tokenizer( - ... "summarize: studies have shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> outputs = model.generate(input_ids) - >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) - >>> # studies have shown that owning a dog is good for you. - ```""" - # input_ids = input_ids[0] # 这里只是权宜之计 - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - self.tags = label_tokenized_words_corrected(self.tokenizer,input_ids) # [batchsize,sequence] - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - # Convert encoder inputs in embeddings if needed - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - - if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: - # get decoder inputs from shifting lm labels to the right - decoder_input_ids = self._shift_right(labels) - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - hidden_states = hidden_states.to(self.decoder.first_device) - if decoder_input_ids is not None: - decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) - if attention_mask is not None: - attention_mask = attention_mask.to(self.decoder.first_device) - if decoder_attention_mask is not None: - decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions="True", - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = decoder_outputs[0] - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.encoder.first_device) - self.lm_head = self.lm_head.to(self.encoder.first_device) - sequence_output = sequence_output.to(self.lm_head.weight.device) - - if self.config.tie_word_embeddings: - # Rescale output before projecting on vocab - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 - sequence_output = sequence_output * (self.model_dim**-0.5) - - lm_logits1 = self.lm_head(sequence_output) - # print(hidden_states.shape) - # print(sequence_output.shape) - # print(lm_logits.shape) - # cross_attentions = decoder_outputs["cross_attentions"][-1] - # cross_attentions = torch.mean(cross_attentions, dim=1) - # print(cross_attentions.shape) - # print(cross_attentions.shape) - # lm_logits = torch.cat([lm_logits, cross_attentions], dim=2) - # print(lm_logits.shape) - # eps = 1e-7 - # lm_logits = torch.log(lm_logits + eps) - lm_logits = self.pointNet(hidden_states,sequence_output,lm_logits1,self.tags) - # print(labels) - # print(lm_logits) - - loss = None - if labels is not None: - # # 创建NLLLoss对象 - # loss_fct = nn.NLLLoss() - loss_fct = CrossEntropyLoss(ignore_index=-100) - # move labels to correct device to enable PP - labels = labels.to(lm_logits.device) - # print(labels.shape) - # print(lm_logits.view(-1, lm_logits.size(-1)).shape) - loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666 - - if not return_dict: - output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs - return ((loss,) + output) if loss is not None else output - - return Seq2SeqLMOutput( - loss=loss, - logits=lm_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_inputs_for_generation - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - attention_mask=None, - head_mask=None, - decoder_head_mask=None, - decoder_attention_mask=None, - cross_attn_head_mask=None, - use_cache=None, - encoder_outputs=None, - **kwargs, - ): - # cut decoder_input_ids if past_key_values is used - if past_key_values is not None: - past_length = past_key_values[0][0].shape[2] - - # Some generation methods already pass only the last input ID - if input_ids.shape[1] > past_length: - remove_prefix_length = past_length - else: - # Default to old behavior: keep only final ID - remove_prefix_length = input_ids.shape[1] - 1 - - input_ids = input_ids[:, remove_prefix_length:] - - return { - "decoder_input_ids": input_ids, - "past_key_values": past_key_values, - "encoder_outputs": encoder_outputs, - "attention_mask": attention_mask, - "head_mask": head_mask, - "decoder_head_mask": decoder_head_mask, - "decoder_attention_mask": decoder_attention_mask, - "cross_attn_head_mask": cross_attn_head_mask, - "use_cache": use_cache, - } - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_decoder_input_ids_from_labels - def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): - return self._shift_right(labels) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration._reorder_cache - def _reorder_cache(self, past_key_values, beam_idx): - # if decoder past is not included in output - # speedy decoding is disabled and no need to reorder - if past_key_values is None: - logger.warning("You might want to consider setting `use_cache=True` to speed up decoding") - return past_key_values - - reordered_decoder_past = () - for layer_past_states in past_key_values: - # get the correct batch idx from layer past batch dim - # batch dim of `past` is at 2nd position - reordered_layer_past_states = () - for layer_past_state in layer_past_states: - # need to set correct `past` for each of the four key / value states - reordered_layer_past_states = reordered_layer_past_states + ( - layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)), - ) - - if reordered_layer_past_states[0].shape != layer_past_states[0].shape: - raise ValueError( - f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched" - ) - if len(reordered_layer_past_states) != len(layer_past_states): - raise ValueError( - f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched" - ) - - reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,) - return reordered_decoder_past - - -@add_start_docstrings( - "The bare MT5 Model transformer outputting encoder's raw hidden-states without any specific head on top.", - MT5_START_DOCSTRING, -) -class MT5EncoderModel(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5EncoderModel, AutoTokenizer - - >>> model = MT5EncoderModel.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> input_ids = tokenizer(article, return_tensors="pt").input_ids - >>> outputs = model(input_ids) - >>> hidden_state = outputs.last_hidden_state - ```""" - - model_type = "mt5" - config_class = MT5Config - _tied_weights_keys = ["encoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - # Initialize weights and apply final processing - self.post_init() - - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5EncoderModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load" - " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," - " 'block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel._prune_heads - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.block[layer].layer[0].SelfAttention.prune_heads(heads) - - @add_start_docstrings_to_model_forward(MT5_ENCODER_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=BaseModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]: - r""" - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MT5EncoderModel - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5EncoderModel.from_pretrained("mt5-small") - >>> input_ids = tokenizer( - ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> outputs = model(input_ids=input_ids) - >>> last_hidden_states = outputs.last_hidden_state - ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - return encoder_outputs - - -@add_start_docstrings( - """ - MT5 model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE - tasks. - """, - MT5_START_DOCSTRING, -) -class MT5ForSequenceClassification(MT5PreTrainedModel): - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.transformer = MT5Model(config) - self.classification_head = MT5ClassificationHead(config) - - # Initialize weights and apply final processing - self.post_init() - - self.model_parallel = False - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.forward - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - decoder_head_mask: Optional[torch.Tensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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, Seq2SeqSequenceClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - Returns: - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if labels is not None: - use_cache = False - - if input_ids is None and inputs_embeds is not None: - raise NotImplementedError( - f"Passing input embeddings is currently not supported for {self.__class__.__name__}" - ) - - # Copied from models.bart.modeling_bart.BartModel.forward different to other models, T5 automatically creates - # decoder_input_ids from input_ids if no decoder_input_ids are provided - if decoder_input_ids is None and decoder_inputs_embeds is None: - if input_ids is None: - raise ValueError( - "If no `decoder_input_ids` or `decoder_inputs_embeds` are " - "passed, `input_ids` cannot be `None`. Please pass either " - "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." - ) - decoder_input_ids = self._shift_right(input_ids) - - outputs = self.transformer( - input_ids, - attention_mask=attention_mask, - decoder_input_ids=decoder_input_ids, - decoder_attention_mask=decoder_attention_mask, - head_mask=head_mask, - decoder_head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - encoder_outputs=encoder_outputs, - inputs_embeds=inputs_embeds, - decoder_inputs_embeds=decoder_inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - sequence_output = outputs[0] - - eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) - - if len(torch.unique_consecutive(eos_mask.sum(1))) > 1: - raise ValueError("All examples must have the same number of tokens.") - batch_size, _, hidden_size = sequence_output.shape - sentence_representation = sequence_output[eos_mask, :].view(batch_size, -1, hidden_size)[:, -1, :] - logits = self.classification_head(sentence_representation) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.config.num_labels == 1: - self.config.problem_type = "regression" - elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.config.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output - - return Seq2SeqSequenceClassifierOutput( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - decoder_hidden_states=outputs.decoder_hidden_states, - decoder_attentions=outputs.decoder_attentions, - cross_attentions=outputs.cross_attentions, - encoder_last_hidden_state=outputs.encoder_last_hidden_state, - encoder_hidden_states=outputs.encoder_hidden_states, - encoder_attentions=outputs.encoder_attentions, - ) - - -@add_start_docstrings( - """ - MT5 Model with a span classification head on top for extractive question-answering tasks like SQuAD (linear layers - on top of the hidden-states output to compute `span start logits` and `span end logits`). - """, - MT5_START_DOCSTRING, -) -class MT5ForQuestionAnswering(MT5PreTrainedModel): - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.model_dim = config.d_model - - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - self.decoder = MT5Stack(decoder_config, self.shared) - - self.num_labels = config.num_labels - self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) - - # Initialize weights and apply final processing - self.post_init() - - self.model_parallel = False - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_decoder - def get_decoder(self): - return self.decoder - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqQuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.forward - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - start_positions: Optional[torch.LongTensor] = None, - end_positions: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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[torch.FloatTensor], Seq2SeqQuestionAnsweringModelOutput]: - r""" - start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the start of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence - are not taken into account for computing the loss. - end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the end of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence - are not taken into account for computing the loss. - Returns: - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - use_cache = use_cache if use_cache is not None else self.config.use_cache - if start_positions is not None and end_positions is not None: - use_cache = False - - # Copied from models.bart.modeling_bart.BartModel.forward - # different to other models, T5 automatically creates decoder_input_ids from - # input_ids if no decoder_input_ids are provided - if decoder_input_ids is None and decoder_inputs_embeds is None: - if input_ids is None: - raise ValueError( - "If no `decoder_input_ids` or `decoder_inputs_embeds` are " - "passed, `input_ids` cannot be `None`. Please pass either " - "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." - ) - decoder_input_ids = self._shift_right(input_ids) - - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=None, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = decoder_outputs[0] - - logits = self.qa_outputs(sequence_output) - start_logits, end_logits = logits.split(1, dim=-1) - start_logits = start_logits.squeeze(-1).contiguous() - end_logits = end_logits.squeeze(-1).contiguous() - - total_loss = None - if start_positions is not None and end_positions is not None: - # If we are on multi-GPU, split add a dimension - if len(start_positions.size()) > 1: - start_positions = start_positions.squeeze(-1).to(start_logits.device) - if len(end_positions.size()) > 1: - end_positions = end_positions.squeeze(-1).to(end_logits.device) - # sometimes the start/end positions are outside our model inputs, we ignore these terms - ignored_index = start_logits.size(1) - start_positions = start_positions.clamp(0, ignored_index) - end_positions = end_positions.clamp(0, ignored_index) - - loss_fct = CrossEntropyLoss(ignore_index=ignored_index) - start_loss = loss_fct(start_logits, start_positions) - end_loss = loss_fct(end_logits, end_positions) - total_loss = (start_loss + end_loss) / 2 - - if not return_dict: - output = (start_logits, end_logits) + decoder_outputs[1:] + encoder_outputs - return ((total_loss,) + output) if total_loss is not None else output - - return Seq2SeqQuestionAnsweringModelOutput( - loss=total_loss, - start_logits=start_logits, - end_logits=end_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) diff --git a/module/MT5SelfTrain.py b/module/MT5SelfTrain.py deleted file mode 100644 index 4eb8709..0000000 --- a/module/MT5SelfTrain.py +++ /dev/null @@ -1,2662 +0,0 @@ -# coding=utf-8 -# Copyright 2020 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" PyTorch mT5 model.""" - -import copy -import math -import os -import warnings -from typing import List, Optional, Tuple, Union - -import torch -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss - -from transformers.activations import ACT2FN -from transformers.modeling_outputs import ( - BaseModelOutput, - BaseModelOutputWithPastAndCrossAttentions, - Seq2SeqLMOutput, - Seq2SeqModelOutput, - Seq2SeqQuestionAnsweringModelOutput, - Seq2SeqSequenceClassifierOutput, -) -from transformers.modeling_utils import PreTrainedModel -from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer -from transformers.utils import ( - DUMMY_INPUTS, - DUMMY_MASK, - add_start_docstrings, - add_start_docstrings_to_model_forward, - is_torch_fx_proxy, - logging, - replace_return_docstrings, -) -from transformers.utils.model_parallel_utils import assert_device_map, get_device_map -from .configuration_mt5 import MT5Config -from transformers import AutoTokenizer - - -logger = logging.get_logger(__name__) - -_CONFIG_FOR_DOC = "MT5Config" -_CHECKPOINT_FOR_DOC = "mt5-small" - - -PARALLELIZE_DOCSTRING = r""" - This is an experimental feature and is a subject to change at a moment's notice. - - Uses a device map to distribute attention modules of the model across several devices. If no device map is given, - it will evenly distribute blocks across all devices. - - Args: - device_map (`Dict[int, list]`, optional, defaults to None): - A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always - automatically mapped to the first device (for esoteric reasons). That means that the first device should - have fewer attention modules mapped to it than other devices. For reference, the mt5 models have the - following number of attention modules: - - - mt5-small: 6 - - mt5-base: 12 - - mt5-large: 24 - - mt5-xl: 24 - - mt5-xxl: 24 - - Example: - - ```python - # Here is an example of a device map on a machine with 4 GPUs using mt5-xl, which has a total of 24 attention modules: - model = MT5ForConditionalGeneration.from_pretrained("mt5-xl") - device_map = { - 0: [0, 1, 2], - 1: [3, 4, 5, 6, 7, 8, 9], - 2: [10, 11, 12, 13, 14, 15, 16], - 3: [17, 18, 19, 20, 21, 22, 23], - } - model.parallelize(device_map) - ``` -""" -DEPARALLELIZE_DOCSTRING = r""" - Moves the model to cpu from a model parallel state. - - Example: - - ```python - # On a 4 GPU machine with mt5-xl: - model = MT5ForConditionalGeneration.from_pretrained("Mt5-xl") - device_map = { - 0: [0, 1, 2], - 1: [3, 4, 5, 6, 7, 8, 9], - 2: [10, 11, 12, 13, 14, 15, 16], - 3: [17, 18, 19, 20, 21, 22, 23], - } - model.parallelize(device_map) # Splits the model across several devices - model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache() - ``` -""" - - -# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->MT5 -class MT5LayerNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - """ - Construct a layernorm module in the MT5 style. No bias and no subtraction of mean. - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - # MT5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean - # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus varience is calculated - # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for - # half-precision inputs is done in fp32 - - 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 - - -# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->MT5 -class MT5DenseActDense(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - self.wi = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.dropout = nn.Dropout(config.dropout_rate) - self.act = ACT2FN[config.dense_act_fn] - - def forward(self, hidden_states): - hidden_states = self.wi(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.dropout(hidden_states) - if ( - isinstance(self.wo.weight, torch.Tensor) - and hidden_states.dtype != self.wo.weight.dtype - and self.wo.weight.dtype != torch.int8 - ): - hidden_states = hidden_states.to(self.wo.weight.dtype) - hidden_states = self.wo(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->MT5 -class MT5DenseGatedActDense(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.dropout = nn.Dropout(config.dropout_rate) - self.act = ACT2FN[config.dense_act_fn] - - def forward(self, hidden_states): - hidden_gelu = self.act(self.wi_0(hidden_states)) - hidden_linear = self.wi_1(hidden_states) - hidden_states = hidden_gelu * hidden_linear - hidden_states = self.dropout(hidden_states) - - # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. - # See https://github.com/huggingface/transformers/issues/20287 - # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` - if ( - isinstance(self.wo.weight, torch.Tensor) - and hidden_states.dtype != self.wo.weight.dtype - and self.wo.weight.dtype != torch.int8 - ): - hidden_states = hidden_states.to(self.wo.weight.dtype) - - hidden_states = self.wo(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->MT5 -class MT5LayerFF(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - if config.is_gated_act: - self.DenseReluDense = MT5DenseGatedActDense(config) - else: - self.DenseReluDense = MT5DenseActDense(config) - - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward(self, hidden_states): - forwarded_states = self.layer_norm(hidden_states) - forwarded_states = self.DenseReluDense(forwarded_states) - hidden_states = hidden_states + self.dropout(forwarded_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->MT5 -class MT5Attention(nn.Module): - def __init__(self, config: MT5Config, has_relative_attention_bias=False): - super().__init__() - self.is_decoder = config.is_decoder - self.has_relative_attention_bias = has_relative_attention_bias - self.relative_attention_num_buckets = config.relative_attention_num_buckets - self.relative_attention_max_distance = config.relative_attention_max_distance - self.d_model = config.d_model - self.key_value_proj_dim = config.d_kv - self.n_heads = config.num_heads - self.dropout = config.dropout_rate - self.inner_dim = self.n_heads * self.key_value_proj_dim - - # Mesh TensorFlow initialization to avoid scaling before softmax - self.q = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.k = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.v = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.o = nn.Linear(self.inner_dim, self.d_model, bias=False) - - if self.has_relative_attention_bias: - self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) - self.pruned_heads = set() - self.gradient_checkpointing = False - - def prune_heads(self, heads): - if len(heads) == 0: - return - heads, index = find_pruneable_heads_and_indices( - heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads - ) - # Prune linear layers - self.q = prune_linear_layer(self.q, index) - self.k = prune_linear_layer(self.k, index) - self.v = prune_linear_layer(self.v, index) - self.o = prune_linear_layer(self.o, index, dim=1) - # Update hyper params - self.n_heads = self.n_heads - len(heads) - self.inner_dim = self.key_value_proj_dim * self.n_heads - self.pruned_heads = self.pruned_heads.union(heads) - - @staticmethod - def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): - """ - Adapted from Mesh Tensorflow: - https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 - - Translate relative position to a bucket number for relative attention. The relative position is defined as - memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to - position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for - small absolute relative_position and larger buckets for larger absolute relative_positions. All relative - positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. - This should allow for more graceful generalization to longer sequences than the model has been trained on - - Args: - relative_position: an int32 Tensor - bidirectional: a boolean - whether the attention is bidirectional - num_buckets: an integer - max_distance: an integer - - Returns: - a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) - """ - relative_buckets = 0 - if bidirectional: - num_buckets //= 2 - relative_buckets += (relative_position > 0).to(torch.long) * num_buckets - relative_position = torch.abs(relative_position) - else: - relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) - # now relative_position is in the range [0, inf) - - # half of the buckets are for exact increments in positions - max_exact = num_buckets // 2 - is_small = relative_position < max_exact - - # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance - relative_position_if_large = max_exact + ( - torch.log(relative_position.float() / max_exact) - / math.log(max_distance / max_exact) - * (num_buckets - max_exact) - ).to(torch.long) - relative_position_if_large = torch.min( - relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) - ) - - relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) - return relative_buckets - - def compute_bias(self, query_length, key_length, device=None): - """Compute binned relative position bias""" - if device is None: - device = self.relative_attention_bias.weight.device - context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] - memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] - relative_position = memory_position - context_position # shape (query_length, key_length) - relative_position_bucket = self._relative_position_bucket( - relative_position, # shape (query_length, key_length) - bidirectional=(not self.is_decoder), - num_buckets=self.relative_attention_num_buckets, - max_distance=self.relative_attention_max_distance, - ) - values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) - values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) - return values - - def forward( - self, - hidden_states, - mask=None, - key_value_states=None, - position_bias=None, - past_key_value=None, - layer_head_mask=None, - query_length=None, - use_cache=False, - output_attentions=False, - ): - """ - Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). - """ - # Input is (batch_size, seq_length, dim) - # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) - # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head) - batch_size, seq_length = hidden_states.shape[:2] - - real_seq_length = seq_length - - if past_key_value is not None: - if len(past_key_value) != 2: - raise ValueError( - f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" - ) - real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length - - key_length = real_seq_length if key_value_states is None else key_value_states.shape[1] - - def shape(states): - """projection""" - return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) - - def unshape(states): - """reshape""" - return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) - - def project(hidden_states, proj_layer, key_value_states, past_key_value): - """projects hidden states correctly to key/query states""" - if key_value_states is None: - # self-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(hidden_states)) - elif past_key_value is None: - # cross-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(key_value_states)) - - if past_key_value is not None: - if key_value_states is None: - # self-attn - # (batch_size, n_heads, key_length, dim_per_head) - hidden_states = torch.cat([past_key_value, hidden_states], dim=2) - elif past_key_value.shape[2] != key_value_states.shape[1]: - # checking that the `sequence_length` of the `past_key_value` is the same as - # the provided `key_value_states` to support prefix tuning - # cross-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(key_value_states)) - else: - # cross-attn - hidden_states = past_key_value - return hidden_states - - # get query states - query_states = shape(self.q(hidden_states)) # (batch_size, n_heads, seq_length, dim_per_head) - - # get key/value states - key_states = project( - hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None - ) - value_states = project( - hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None - ) - - # compute scores - scores = torch.matmul( - query_states, key_states.transpose(3, 2) - ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 - - if position_bias is None: - if not self.has_relative_attention_bias: - position_bias = torch.zeros( - (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype - ) - if self.gradient_checkpointing and self.training: - position_bias.requires_grad = True - else: - position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device) - - # if key and values are already calculated - # we want only the last query position bias - if past_key_value is not None: - position_bias = position_bias[:, :, -hidden_states.size(1) :, :] - - if mask is not None: - position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length) - - if self.pruned_heads: - mask = torch.ones(position_bias.shape[1]) - mask[list(self.pruned_heads)] = 0 - position_bias_masked = position_bias[:, mask.bool()] - else: - position_bias_masked = position_bias - - scores += position_bias_masked - attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as( - scores - ) # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.dropout( - attn_weights, p=self.dropout, training=self.training - ) # (batch_size, n_heads, seq_length, key_length) - - # Mask heads if we want to - if layer_head_mask is not None: - attn_weights = attn_weights * layer_head_mask - - attn_output = unshape(torch.matmul(attn_weights, value_states)) # (batch_size, seq_length, dim) - attn_output = self.o(attn_output) - - present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None - outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) - - if output_attentions: - outputs = outputs + (attn_weights,) - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->MT5 -class MT5LayerSelfAttention(nn.Module): - def __init__(self, config, has_relative_attention_bias=False): - super().__init__() - self.SelfAttention = MT5Attention(config, has_relative_attention_bias=has_relative_attention_bias) - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward( - self, - hidden_states, - attention_mask=None, - position_bias=None, - layer_head_mask=None, - past_key_value=None, - use_cache=False, - output_attentions=False, - ): - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( - normed_hidden_states, - mask=attention_mask, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states = hidden_states + self.dropout(attention_output[0]) - outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->MT5 -class MT5LayerCrossAttention(nn.Module): - def __init__(self, config): - super().__init__() - self.EncDecAttention = MT5Attention(config, has_relative_attention_bias=False) - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward( - self, - hidden_states, - key_value_states, - attention_mask=None, - position_bias=None, - layer_head_mask=None, - past_key_value=None, - use_cache=False, - query_length=None, - output_attentions=False, - ): - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( - normed_hidden_states, - mask=attention_mask, - key_value_states=key_value_states, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - query_length=query_length, - output_attentions=output_attentions, - ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5Block with T5->MT5 -class MT5Block(nn.Module): - def __init__(self, config, has_relative_attention_bias=False): - super().__init__() - self.is_decoder = config.is_decoder - self.layer = nn.ModuleList() - self.layer.append(MT5LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias)) - if self.is_decoder: - self.layer.append(MT5LayerCrossAttention(config)) - - self.layer.append(MT5LayerFF(config)) - - def forward( - self, - hidden_states, - attention_mask=None, - position_bias=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - encoder_decoder_position_bias=None, - layer_head_mask=None, - cross_attn_layer_head_mask=None, - past_key_value=None, - use_cache=False, - output_attentions=False, - return_dict=True, - ): - if past_key_value is not None: - if not self.is_decoder: - logger.warning("`past_key_values` is passed to the encoder. Please make sure this is intended.") - expected_num_past_key_values = 2 if encoder_hidden_states is None else 4 - - if len(past_key_value) != expected_num_past_key_values: - raise ValueError( - f"There should be {expected_num_past_key_values} past states. " - f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}" - f"Got {len(past_key_value)} past key / value states" - ) - - self_attn_past_key_value = past_key_value[:2] - cross_attn_past_key_value = past_key_value[2:] - else: - self_attn_past_key_value, cross_attn_past_key_value = None, None - - self_attention_outputs = self.layer[0]( - hidden_states, - attention_mask=attention_mask, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=self_attn_past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states, present_key_value_state = self_attention_outputs[:2] - attention_outputs = self_attention_outputs[2:] # Keep self-attention outputs and relative position weights - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - do_cross_attention = self.is_decoder and encoder_hidden_states is not None - if do_cross_attention: - # the actual query length is unknown for cross attention - # if using past key value states. Need to inject it here - if present_key_value_state is not None: - query_length = present_key_value_state[0].shape[2] - else: - query_length = None - - cross_attention_outputs = self.layer[1]( - hidden_states, - key_value_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - position_bias=encoder_decoder_position_bias, - layer_head_mask=cross_attn_layer_head_mask, - past_key_value=cross_attn_past_key_value, - query_length=query_length, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states = cross_attention_outputs[0] - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - # Combine self attn and cross attn key value states - if present_key_value_state is not None: - present_key_value_state = present_key_value_state + cross_attention_outputs[1] - - # Keep cross-attention outputs and relative position weights - attention_outputs = attention_outputs + cross_attention_outputs[2:] - - # Apply Feed Forward layer - hidden_states = self.layer[-1](hidden_states) - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - outputs = (hidden_states,) - - if use_cache: - outputs = outputs + (present_key_value_state,) + attention_outputs - else: - outputs = outputs + attention_outputs - - return outputs # hidden-states, present_key_value_states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) - - -def load_tf_weights_in_mt5(model, config, tf_checkpoint_path): - """Load tf checkpoints in a pytorch model.""" - try: - import re - - import numpy as np - import tensorflow as tf - except ImportError: - logger.error( - "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " - "https://www.tensorflow.org/install/ for installation instructions." - ) - raise - tf_path = os.path.abspath(tf_checkpoint_path) - logger.info(f"Converting TensorFlow checkpoint from {tf_path}") - # Load weights from TF model - init_vars = tf.train.list_variables(tf_path) - names = [] - tf_weights = {} - for name, shape in init_vars: - logger.info(f"Loading TF weight {name} with shape {shape}") - array = tf.train.load_variable(tf_path, name) - names.append(name) - tf_weights[name] = array - - for txt_name in names: - name = txt_name.split("/") - # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v - # which are not required for using pretrained model - if any( - n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] - for n in name - ): - logger.info(f"Skipping {'/'.join(name)}") - tf_weights.pop(txt_name, None) - continue - if "_slot_" in name[-1]: - logger.info(f"Skipping {'/'.join(name)}") - tf_weights.pop(txt_name, None) - continue - pointer = model - array = tf_weights[txt_name] - - for m_name in name: - if re.fullmatch(r"[A-Za-z]+_\d+", m_name): - scope_names = re.split(r"_(\d+)", m_name) - else: - scope_names = [m_name] - if scope_names[0] in ["kernel", "scale", "embedding"]: - pointer = getattr(pointer, "weight") - elif scope_names[0] == "self_attention": - pointer = getattr(pointer, "layer") - pointer = pointer[0] - elif scope_names[0] == "enc_dec_attention": - pointer = getattr(pointer, "layer") - pointer = pointer[1] - elif scope_names[0] == "dense_relu_dense": - pointer = getattr(pointer, "layer") - pointer = pointer[2] - elif scope_names[0] == "rms_norm": - if hasattr(pointer, "layer_norm"): - pointer = getattr(pointer, "layer_norm") - elif hasattr(pointer, "final_layer_norm"): - pointer = getattr(pointer, "final_layer_norm") - elif scope_names[0] == "scale": - pointer = getattr(pointer, "weight") - elif scope_names[0] == "output_bias" or scope_names[0] == "beta": - pointer = getattr(pointer, "bias") - elif scope_names[0] == "squad": - pointer = getattr(pointer, "classifier") - elif scope_names[0] == "decoder" and name[1] == "logits": - continue - elif scope_names[0] == "logits": - pointer = getattr(pointer, "lm_head") - elif scope_names[0] == "wi" and len(scope_names) > 1 and scope_names[1].isdigit(): - pointer = getattr(pointer, f"wi_{scope_names[1]}") - continue - else: - try: - pointer = getattr(pointer, scope_names[0]) - except AttributeError: - logger.info(f"Skipping {'/'.join(name)}") - continue - if len(scope_names) >= 2: - num = int(scope_names[1]) - pointer = pointer[num] - if scope_names[0] not in ["kernel", "scale", "embedding"]: - pointer = getattr(pointer, "weight") - if scope_names[0] != "embedding": - logger.info(f"Transposing numpy weight of shape {array.shape} for {name}") - array = np.transpose(array) - try: - assert ( - pointer.shape == array.shape - ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched" - except AssertionError as e: - e.args += (pointer.shape, array.shape) - raise - logger.info(f"Initialize PyTorch weight {name}") - pointer.data = torch.from_numpy(array.astype(np.float32)) - tf_weights.pop(txt_name, None) - - logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}.") - return model - - -# Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->MT5 -class MT5ClassificationHead(nn.Module): - """Head for sentence-level classification tasks.""" - - def __init__(self, config: MT5Config): - super().__init__() - self.dense = nn.Linear(config.d_model, config.d_model) - self.dropout = nn.Dropout(p=config.classifier_dropout) - self.out_proj = nn.Linear(config.d_model, config.num_labels) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.dropout(hidden_states) - hidden_states = self.dense(hidden_states) - hidden_states = torch.tanh(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.out_proj(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel with T5->MT5, t5->mt5 -class MT5PreTrainedModel(PreTrainedModel): - """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. - """ - - config_class = MT5Config - load_tf_weights = load_tf_weights_in_mt5 - base_model_prefix = "transformer" - is_parallelizable = True - supports_gradient_checkpointing = True - _no_split_modules = ["MT5Block"] - _keep_in_fp32_modules = ["wo"] - - @property - def dummy_inputs(self): - input_ids = torch.tensor(DUMMY_INPUTS) - input_mask = torch.tensor(DUMMY_MASK) - dummy_inputs = { - "decoder_input_ids": input_ids, - "input_ids": input_ids, - "decoder_attention_mask": input_mask, - } - return dummy_inputs - - def _init_weights(self, module): - """Initialize the weights""" - factor = self.config.initializer_factor # Used for testing weights initialization - if isinstance(module, MT5LayerNorm): - module.weight.data.fill_(factor * 1.0) - elif isinstance( - module, - (MT5Model, MT5ForConditionalGeneration, MT5EncoderModel, MT5ForQuestionAnswering), - ): - # Mesh TensorFlow embeddings initialization - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624 - module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) - if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: - module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0) - if hasattr(module, "qa_outputs"): - module.qa_outputs.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - module.qa_outputs.bias.data.zero_() - elif isinstance(module, MT5ClassificationHead): - module.dense.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.dense, "bias") and module.dense.bias is not None: - module.dense.bias.data.zero_() - module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: - module.out_proj.bias.data.zero_() - elif isinstance(module, MT5DenseActDense): - # Mesh TensorFlow FF initialization - # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56 - # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89 - module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi, "bias") and module.wi.bias is not None: - module.wi.bias.data.zero_() - module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) - if hasattr(module.wo, "bias") and module.wo.bias is not None: - module.wo.bias.data.zero_() - elif isinstance(module, MT5DenseGatedActDense): - module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: - module.wi_0.bias.data.zero_() - module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: - module.wi_1.bias.data.zero_() - module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) - if hasattr(module.wo, "bias") and module.wo.bias is not None: - module.wo.bias.data.zero_() - elif isinstance(module, MT5Attention): - # Mesh TensorFlow attention initialization to avoid scaling before softmax - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136 - d_model = self.config.d_model - key_value_proj_dim = self.config.d_kv - n_heads = self.config.num_heads - module.q.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5)) - module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) - module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) - module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) - if module.has_relative_attention_bias: - module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5)) - - def _shift_right(self, input_ids): - decoder_start_token_id = self.config.decoder_start_token_id - pad_token_id = self.config.pad_token_id - - if decoder_start_token_id is None: - raise ValueError( - "self.model.config.decoder_start_token_id has to be defined. In MT5 it is usually set to the pad_token_id. " - "See MT5 docs for more information." - ) - - # shift inputs to the right - if is_torch_fx_proxy(input_ids): - # Item assignment is not supported natively for proxies. - shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id) - shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) - else: - shifted_input_ids = input_ids.new_zeros(input_ids.shape) - shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() - shifted_input_ids[..., 0] = decoder_start_token_id - - if pad_token_id is None: - raise ValueError("self.model.config.pad_token_id has to be defined.") - # replace possible -100 values in labels by `pad_token_id` - shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) - - return shifted_input_ids - - -# Copied from transformers.models.t5.modeling_t5.T5Stack with T5->MT5 -class MT5Stack(MT5PreTrainedModel): - def __init__(self, config, embed_tokens=None): - super().__init__(config) - - self.embed_tokens = embed_tokens - self.is_decoder = config.is_decoder - - self.block = nn.ModuleList( - [MT5Block(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)] - ) - self.final_layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - # Initialize weights and apply final processing - self.post_init() - # Model parallel - self.model_parallel = False - self.device_map = None - self.gradient_checkpointing = False - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - def parallelize(self, device_map=None): - warnings.warn( - "`MT5Stack.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" - " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," - " 'block.1': 1, ...}", - FutureWarning, - ) - # Check validity of device_map - self.device_map = ( - get_device_map(len(self.block), range(torch.cuda.device_count())) if device_map is None else device_map - ) - assert_device_map(self.device_map, len(self.block)) - self.model_parallel = True - self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys())) - self.last_device = "cuda:" + str(max(self.device_map.keys())) - # Load onto devices - for k, v in self.device_map.items(): - for layer in v: - cuda_device = "cuda:" + str(k) - self.block[layer] = self.block[layer].to(cuda_device) - - # Set embed_tokens to first layer - self.embed_tokens = self.embed_tokens.to(self.first_device) - # Set final layer norm to last device - self.final_layer_norm = self.final_layer_norm.to(self.last_device) - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.model_parallel = False - self.device_map = None - self.first_device = "cpu" - self.last_device = "cpu" - for i in range(len(self.block)): - self.block[i] = self.block[i].to("cpu") - self.embed_tokens = self.embed_tokens.to("cpu") - self.final_layer_norm = self.final_layer_norm.to("cpu") - torch.cuda.empty_cache() - - def get_input_embeddings(self): - return self.embed_tokens - - def set_input_embeddings(self, new_embeddings): - self.embed_tokens = new_embeddings - - def forward( - self, - input_ids=None, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - inputs_embeds=None, - head_mask=None, - cross_attn_head_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - # Model parallel - if self.model_parallel: - torch.cuda.set_device(self.first_device) - self.embed_tokens = self.embed_tokens.to(self.first_device) - use_cache = use_cache if use_cache is not None else self.config.use_cache - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is not None and inputs_embeds is not None: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError( - f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" - ) - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") - - if inputs_embeds is None: - if self.embed_tokens is None: - raise ValueError("You have to initialize the model with valid token embeddings") - inputs_embeds = self.embed_tokens(input_ids) - # print(input_ids) - batch_size, seq_length = input_shape - - # required mask seq length can be calculated via length of past - mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length - - if use_cache is True: - if not self.is_decoder: - raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") - - if attention_mask is None: - attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) - if self.is_decoder and encoder_attention_mask is None and encoder_hidden_states is not None: - encoder_seq_length = encoder_hidden_states.shape[1] - encoder_attention_mask = torch.ones( - batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long - ) - - # initialize past_key_values with `None` if past does not exist - if past_key_values is None: - past_key_values = [None] * len(self.block) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape) - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.is_decoder and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = None - - 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 - - # Prepare head mask if needed - head_mask = self.get_head_mask(head_mask, self.config.num_layers) - cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers) - present_key_value_states = () if use_cache else None - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - all_cross_attentions = () if (output_attentions and self.is_decoder) else None - position_bias = None - encoder_decoder_position_bias = None - - hidden_states = self.dropout(inputs_embeds) - - for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)): - layer_head_mask = head_mask[i] - cross_attn_layer_head_mask = cross_attn_head_mask[i] - # Model parallel - if self.model_parallel: - torch.cuda.set_device(hidden_states.device) - # Ensure that attention_mask is always on the same device as hidden_states - if attention_mask is not None: - attention_mask = attention_mask.to(hidden_states.device) - if position_bias is not None: - position_bias = position_bias.to(hidden_states.device) - if encoder_hidden_states is not None: - encoder_hidden_states = encoder_hidden_states.to(hidden_states.device) - if encoder_extended_attention_mask is not None: - encoder_extended_attention_mask = encoder_extended_attention_mask.to(hidden_states.device) - if encoder_decoder_position_bias is not None: - encoder_decoder_position_bias = encoder_decoder_position_bias.to(hidden_states.device) - if layer_head_mask is not None: - layer_head_mask = layer_head_mask.to(hidden_states.device) - if cross_attn_layer_head_mask is not None: - cross_attn_layer_head_mask = cross_attn_layer_head_mask.to(hidden_states.device) - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer_module.forward, - hidden_states, - extended_attention_mask, - position_bias, - encoder_hidden_states, - encoder_extended_attention_mask, - encoder_decoder_position_bias, - layer_head_mask, - cross_attn_layer_head_mask, - None, # past_key_value is always None with gradient checkpointing - use_cache, - output_attentions, - ) - - else: - layer_outputs = layer_module( - hidden_states, - attention_mask=extended_attention_mask, - position_bias=position_bias, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - encoder_decoder_position_bias=encoder_decoder_position_bias, - layer_head_mask=layer_head_mask, - cross_attn_layer_head_mask=cross_attn_layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - - # layer_outputs is a tuple with: - # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) - if use_cache is False: - layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:] - - hidden_states, present_key_value_state = layer_outputs[:2] - - # We share the position biases between the layers - the first layer store them - # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), - # (cross-attention position bias), (cross-attention weights) - position_bias = layer_outputs[2] - if self.is_decoder and encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[4 if output_attentions else 3] - # append next layer key value states - if use_cache: - present_key_value_states = present_key_value_states + (present_key_value_state,) - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[3],) - if self.is_decoder: - all_cross_attentions = all_cross_attentions + (layer_outputs[5],) - - # Model Parallel: If it's the last layer for that device, put things on the next device - if self.model_parallel: - for k, v in self.device_map.items(): - if i == v[-1] and "cuda:" + str(k) != self.last_device: - hidden_states = hidden_states.to("cuda:" + str(k + 1)) - - hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.dropout(hidden_states) - - # Add last layer - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - present_key_value_states, - all_hidden_states, - all_attentions, - all_cross_attentions, - ] - if v is not None - ) - # print("all_hidden_states", all_hidden_states.shape) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=present_key_value_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - cross_attentions=all_cross_attentions, - ) - - -MT5_START_DOCSTRING = r""" - - The MT5 model was proposed in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text - Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan - Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. It's an encoder decoder transformer pre-trained in a - text-to-text denoising generative setting. - - This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the - library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads - etc.) - - This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. - Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage - and behavior. - - Parameters: - config ([`MT5Config`]): Model configuration class with all the parameters of the model. - Initializing with a config file does not load the weights associated with the model, only the - configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. -""" - -MT5_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you - should be able to pad the inputs on both the right and the left. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for detail. - - [What are input IDs?](../glossary#input-ids) - - To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). - attention_mask (`torch.FloatTensor` 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) - decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): - Indices of decoder input sequence tokens in the vocabulary. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for details. - - [What are decoder input IDs?](../glossary#decoder-input-ids) - - MT5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` - is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). - - To know more on how to prepare `decoder_input_ids` for pretraining take a look at [MT5 - Training](./mt5#training). - decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): - Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also - be used by default. - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules in the encoder. Mask values selected in `[0, - 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, - 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in - `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): - Tuple consists of (`last_hidden_state`, `optional`: *hidden_states*, `optional`: *attentions*) - `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden states at - the output of the last layer of the encoder. Used in the cross-attention of the decoder. - past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - 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. - decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): - Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded - representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be - input (see `past_key_values`). This is useful if you want more control over how to convert - `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. - - If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value - of `inputs_embeds`. - - 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. -""" - -MT5_ENCODER_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you - should be able to pad the inputs on both the right and the left. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for detail. - - To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). - attention_mask (`torch.FloatTensor` 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) - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - 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. - 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. -""" - -# Warning message for FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask -__HEAD_MASK_WARNING_MSG = """ -The input argument `head_mask` was split into two arguments `head_mask` and `decoder_head_mask`. Currently, -`decoder_head_mask` is set to copy `head_mask`, but this feature is deprecated and will be removed in future versions. -If you do not want to use any `decoder_head_mask` now, please set `decoder_head_mask = torch.ones(num_layers, -num_heads)`. -""" - - -@add_start_docstrings( - "The bare MT5 Model transformer outputting raw hidden-states without any specific head on top.", - MT5_START_DOCSTRING, -) -class MT5Model(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5Model, AutoTokenizer - - >>> model = MT5Model.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> summary = "Weiter Verhandlung in Syrien." - >>> inputs = tokenizer(article, return_tensors="pt") - >>> labels = tokenizer(text_target=summary, return_tensors="pt") - - >>> outputs = model(input_ids=inputs["input_ids"], decoder_input_ids=labels["input_ids"]) - >>> hidden_states = outputs.last_hidden_state - ```""" - model_type = "mt5" - config_class = MT5Config - _keys_to_ignore_on_load_missing = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5Model.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - self.decoder = MT5Stack(decoder_config, self.shared) - - # Initialize weights and apply final processing - self.post_init() - - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5Model.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" - " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'encoder.block.0':" - " 0, 'encoder.block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.decoder.parallelize(self.device_map) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5Model.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.decoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.decoder = self.decoder.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5Model.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_decoder - def get_decoder(self): - return self.decoder - - # Copied from transformers.models.t5.modeling_t5.T5Model._prune_heads - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.layer[layer].attention.prune_heads(heads) - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5Model.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - inputs_embeds: Optional[torch.Tensor] = None, - decoder_inputs_embeds: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], Seq2SeqModelOutput]: - r""" - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MT5Model - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5Model.from_pretrained("mt5-small") - - >>> input_ids = tokenizer( - ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 - - >>> # preprocess: Prepend decoder_input_ids with start token which is pad token for MT5Model. - >>> # This is not needed for torch's MT5ForConditionalGeneration as it does this internally using labels arg. - >>> decoder_input_ids = model._shift_right(decoder_input_ids) - - >>> # forward pass - >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) - >>> last_hidden_states = outputs.last_hidden_state - ```""" - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - hidden_states = hidden_states.to(self.decoder.first_device) - if decoder_input_ids is not None: - decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) - if attention_mask is not None: - attention_mask = attention_mask.to(self.decoder.first_device) - if decoder_attention_mask is not None: - decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - if not return_dict: - return decoder_outputs + encoder_outputs - - return Seq2SeqModelOutput( - last_hidden_state=decoder_outputs.last_hidden_state, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) - - -import torch.nn.functional as F -class PointAttention(nn.Module): - def __init__(self): - super().__init__() - self.hidden_size = 768 - # self.linear_layer = nn.Linear(250778, 1) - - - def forward(self,hidden1_states,hidden2_states,lm_logits,tags = [[0,0,1,2,-1,3,4,5,-1]]): - taglist = [] - finalBiasList = [] - sentence_len = hidden1_states.shape[1] - output_sentence_len = hidden2_states.shape[1] - for tag in tags: - vect,bias = merge_weights(tag,sentence_len) - # print('sentence_len', sentence_len) - taglist.append(vect) - biasList = [] - for i in range(output_sentence_len): - biasList.append(bias) - finalBiasList.append(biasList) - merge = torch.tensor(taglist, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) # 8 512[实则300] 512[指明需要的行] - finalBias = torch.tensor(finalBiasList, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) - # 8 512[实则300] 728 - - # print(hidden1_states.shape) #[8, 512, 768] - # 尝试给merge扩维 - TopK_Shape = hidden1_states.shape[0] - Shape = merge.shape[0] - k = TopK_Shape/Shape - if(k>1): - merge1 = merge.repeat_interleave(int(k),dim=0) - finalBias1 = finalBias.repeat_interleave(int(k),dim=0) - else: - merge1 = merge - finalBias1 = finalBias - # merge = merge.repeat_interleave(,dim=0) - attn_weights = torch.matmul(merge1,hidden1_states) # hidden1_states1 8 512【300】 728 - - attn_weights1 = torch.matmul(hidden2_states, attn_weights.transpose(1, 2))+finalBias1 # hidden2_states 8 512 768 attn_weights 8 512 768 attn_weights1 8 512 512 - score = torch.cat([lm_logits,attn_weights1], dim=2) - # probabilities = F.softmax(score, dim=-1) - # lm_logits[:, :, -30:] = attn_weights1 - return score - - -def get_vect(i,j,sentence_len=512): - vect = [] - for t in range(sentence_len): - vect.append(0) - t = i - count = j-i - while(t', '', '',''} - -def isdigit(s): - if s.replace('.', '', 1).isdigit() or (s.find(".")!=-1 and s.replace('.', '', 1) == ""): - return True - else: - return False -def label_words(original_sentence,lst): - i = 0 - indices = [] - num = 0 - # print(lst) - while i < len(lst): - if isdigit(lst[i]): - # 当前元素是数字字符串，检查下一个元素 - number_str = lst[i] - indices.append(num) - i += 1 - while i < len(lst) and isdigit(lst[i]): - # 下一个元素也是数字字符串，进行拼接 - number_str += lst[i] - indices.append(num) - i += 1 - num+=1 - elif lst[i] in special_tokens: # 特殊字符检查 - indices.append(-1) - num+=1 - i += 1 - else: - indices.append(num) - num+=1 - i += 1 - # print(indices) - return indices - -# 为列表中的每个元素分配序号 -# 中文没考虑数字 -# def label_words(original_sentence,word_list): -# indices = [] -# num = 0 -# for index, word in enumerate(word_list): -# if word in special_tokens: # 特殊字符检查 -# indices.append(-1) -# num += 1 -# else: -# indices.append(index-num) - -# return indices - -# 英文拼接 -# def label_words(original_sentence,tokenized_words): -# labels = [] -# label = 0 -# # word_index = 0 -# current_word = "" -# lastword = "" -# # 判断是不是最后一个句号 -# period = original_sentence.count(".") -# period_num = 1 -# # print(tokenized_words) -# for word in tokenized_words: -# # 如果当前单词是非字母（如空格或特殊符号），单独分配一个标签 -# pmFlag = False -# if not word.isalpha() and not word.isdigit() and word.find("'") == -1 and not is_valid_time_format(word): -# # .不能丢。如果是句子中的最后一个，则... 如果不是句子中的最后一个。则... -# if word == "." and period_num < period: -# period_num += 1 -# elif word == "." and period_num == period: -# labels.append(label+1) -# continue -# else: -# labels.append(-1) -# lastword = word -# continue - -# # 累加拆分的单词片段 -# current_word += word - -# if lastword.isdigit() and original_sentence.find(lastword + word) != -1 and (word.lower() == "pm" or word.lower() == "am"): -# pmFlag = True -# # 检查当前累加的单词片段是否与原句中的单词相匹配 -# # print(original_sentence.split()) - -# if current_word in original_sentence.split() or pmFlag == True: -# if pmFlag == True: -# label+=1 - -# if current_word in original_sentence.split(): -# current_word = "" - -# labels.append(label) -# # print(current_word) -# else: -# labels.append(label) - -# # 如果当前累加的片段已经形成了完整的单词，则增加标签值 -# if current_word == "": -# label += 1 -# lastword = word -# # return changeModel(labels) -# return labels - -def label_tokenized_words_corrected(tokenizer,token_list): - # tokenizer = AutoTokenizer.from_pretrained("/home/lzx/T5-base/tokenizer1/") - original_sentences = [tokenizer.decode(t[:-1]) for t in token_list] - lsts = [] - for tokens in token_list: - decoded_words = [tokenizer.decode([token_id]) for token_id in tokens] - lsts.append(decoded_words) - # print(lsts) - # 遍历分词后的单词 - llables = [] - for i, tokenized_words in enumerate(lsts): - original_sentence = original_sentences[i] - labels = label_words(original_sentence,tokenized_words) - llables.append(labels) - return llables - -@add_start_docstrings("""MT5 Model with a `language modeling` head on top.""", MT5_START_DOCSTRING) -class MT5ForConditionalGeneration(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5ForConditionalGeneration, AutoTokenizer - - >>> model = MT5ForConditionalGeneration.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> summary = "Weiter Verhandlung in Syrien." - >>> inputs = tokenizer(article, text_target=summary, return_tensors="pt") - - >>> outputs = model(**inputs) - >>> loss = outputs.loss - ```""" - - model_type = "mt5" - config_class = MT5Config - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.model_dim = config.d_model - - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - # print("decoder_config", decoder_config) - self.decoder = MT5Stack(decoder_config, self.shared) - - # -512 - self.lm_head = nn.Linear(config.d_model, config.vocab_size-512, bias=False) - # Initialize weights and apply final processing - self.post_init() - self.pointNet = PointAttention() - # /home/lzx/T5-base-lora/tokenizer2/ - # /home/lzx/T5-base/model_cl_multi/mt5-base-trained-final-save - self.tokenizer = AutoTokenizer.from_pretrained("/home/lzx/T5-base-lora/model/mt5-base-trained-9") - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5ForConditionalGeneration.parallelize` is deprecated and will be removed in v5 of Transformers, you" - " should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also" - " provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance" - " {'encoder.block.0': 0, 'encoder.block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.decoder.parallelize(self.device_map) - self.lm_head = self.lm_head.to(self.decoder.first_device) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.decoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.decoder = self.decoder.to("cpu") - self.lm_head = self.lm_head.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.set_output_embeddings - def set_output_embeddings(self, new_embeddings): - self.lm_head = new_embeddings - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_output_embeddings - def get_output_embeddings(self): - return self.lm_head - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_decoder - def get_decoder(self): - return self.decoder - - def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding: - """ - Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`. - - Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method. - - Arguments: - new_num_tokens (`int`, *optional*): - The number of new tokens in the embedding matrix. Increasing the size will add newly initialized - vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just - returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything. - - Return: - `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model. - """ - model_embeds = self._resize_token_embeddings(new_num_tokens) - if new_num_tokens is None: - return model_embeds - - # Update base model and current model config - self.config.vocab_size = new_num_tokens - self.vocab_size = new_num_tokens - - # Tie weights again if needed - self.tie_weights() - - return model_embeds - - def _resize_token_embeddings(self, new_num_tokens): - old_embeddings = self.get_input_embeddings() - new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens) - self.set_input_embeddings(new_embeddings) - - # if word embeddings are not tied, make sure that lm head is resized as well - if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings: - old_lm_head = self.get_output_embeddings() - #512 - new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens-512) - self.set_output_embeddings(new_lm_head) - - return self.get_input_embeddings() - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - labels1: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., - config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for - labels in `[0, ..., config.vocab_size]` - - Returns: - - Examples: - - ```python - >>> from transformers import AutoTokenizer, MT5ForConditionalGeneration - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5ForConditionalGeneration.from_pretrained("mt5-small") - - >>> # training - >>> input_ids = tokenizer("The walks in park", return_tensors="pt").input_ids - >>> labels = tokenizer(" cute dog the ", return_tensors="pt").input_ids - >>> outputs = model(input_ids=input_ids, labels=labels) - >>> loss = outputs.loss - >>> logits = outputs.logits - - >>> # inference - >>> input_ids = tokenizer( - ... "summarize: studies have shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> outputs = model.generate(input_ids) - >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) - >>> # studies have shown that owning a dog is good for you. - ```""" - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - self.tags = label_tokenized_words_corrected(self.tokenizer,input_ids) # [batchsize,sequence] - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - # Convert encoder inputs in embeddings if needed - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - - if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: - # get decoder inputs from shifting lm labels to the right - decoder_input_ids = self._shift_right(labels) - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - hidden_states = hidden_states.to(self.decoder.first_device) - if decoder_input_ids is not None: - decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) - if attention_mask is not None: - attention_mask = attention_mask.to(self.decoder.first_device) - if decoder_attention_mask is not None: - decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions="True", - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = decoder_outputs[0] - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.encoder.first_device) - self.lm_head = self.lm_head.to(self.encoder.first_device) - sequence_output = sequence_output.to(self.lm_head.weight.device) - - if self.config.tie_word_embeddings: - # Rescale output before projecting on vocab - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 - sequence_output = sequence_output * (self.model_dim**-0.5) - - lm_logits1 = self.lm_head(sequence_output) - # print(hidden_states.shape) - # print(sequence_output.shape) - # print(lm_logits.shape) - # cross_attentions = decoder_outputs["cross_attentions"][-1] - # cross_attentions = torch.mean(cross_attentions, dim=1) - # print(cross_attentions.shape) - # print(cross_attentions.shape) - # lm_logits = torch.cat([lm_logits, cross_attentions], dim=2) - # print(lm_logits.shape) - # eps = 1e-7 - # lm_logits = torch.log(lm_logits + eps) - lm_logits = self.pointNet(hidden_states,sequence_output,lm_logits1,self.tags) - # print(labels) - # print(lm_logits) - - loss = None - if labels is not None: - # # 创建NLLLoss对象 - # loss_fct = nn.NLLLoss() - loss_fct = CrossEntropyLoss(ignore_index=-100) - # move labels to correct device to enable PP - labels = labels.to(lm_logits.device) - # print(labels.shape) - # print(lm_logits.view(-1, lm_logits.size(-1)).shape) - loss = loss_fct(lm_logits, labels1) - # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666 - - if not return_dict: - output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs - return ((loss,) + output) if loss is not None else output - - return Seq2SeqLMOutput( - loss=loss, - logits=lm_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_inputs_for_generation - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - attention_mask=None, - head_mask=None, - decoder_head_mask=None, - decoder_attention_mask=None, - cross_attn_head_mask=None, - use_cache=None, - encoder_outputs=None, - **kwargs, - ): - # cut decoder_input_ids if past_key_values is used - if past_key_values is not None: - past_length = past_key_values[0][0].shape[2] - - # Some generation methods already pass only the last input ID - if input_ids.shape[1] > past_length: - remove_prefix_length = past_length - else: - # Default to old behavior: keep only final ID - remove_prefix_length = input_ids.shape[1] - 1 - - input_ids = input_ids[:, remove_prefix_length:] - - return { - "decoder_input_ids": input_ids, - "past_key_values": past_key_values, - "encoder_outputs": encoder_outputs, - "attention_mask": attention_mask, - "head_mask": head_mask, - "decoder_head_mask": decoder_head_mask, - "decoder_attention_mask": decoder_attention_mask, - "cross_attn_head_mask": cross_attn_head_mask, - "use_cache": use_cache, - } - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_decoder_input_ids_from_labels - def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): - return self._shift_right(labels) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration._reorder_cache - def _reorder_cache(self, past_key_values, beam_idx): - # if decoder past is not included in output - # speedy decoding is disabled and no need to reorder - if past_key_values is None: - logger.warning("You might want to consider setting `use_cache=True` to speed up decoding") - return past_key_values - - reordered_decoder_past = () - for layer_past_states in past_key_values: - # get the correct batch idx from layer past batch dim - # batch dim of `past` is at 2nd position - reordered_layer_past_states = () - for layer_past_state in layer_past_states: - # need to set correct `past` for each of the four key / value states - reordered_layer_past_states = reordered_layer_past_states + ( - layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)), - ) - - if reordered_layer_past_states[0].shape != layer_past_states[0].shape: - raise ValueError( - f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched" - ) - if len(reordered_layer_past_states) != len(layer_past_states): - raise ValueError( - f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched" - ) - - reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,) - return reordered_decoder_past - - -@add_start_docstrings( - "The bare MT5 Model transformer outputting encoder's raw hidden-states without any specific head on top.", - MT5_START_DOCSTRING, -) -class MT5EncoderModel(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5EncoderModel, AutoTokenizer - - >>> model = MT5EncoderModel.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> input_ids = tokenizer(article, return_tensors="pt").input_ids - >>> outputs = model(input_ids) - >>> hidden_state = outputs.last_hidden_state - ```""" - - model_type = "mt5" - config_class = MT5Config - _tied_weights_keys = ["encoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - # Initialize weights and apply final processing - self.post_init() - - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5EncoderModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load" - " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," - " 'block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel._prune_heads - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.block[layer].layer[0].SelfAttention.prune_heads(heads) - - @add_start_docstrings_to_model_forward(MT5_ENCODER_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=BaseModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]: - r""" - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MT5EncoderModel - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5EncoderModel.from_pretrained("mt5-small") - >>> input_ids = tokenizer( - ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> outputs = model(input_ids=input_ids) - >>> last_hidden_states = outputs.last_hidden_state - ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - return encoder_outputs - - -@add_start_docstrings( - """ - MT5 model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE - tasks. - """, - MT5_START_DOCSTRING, -) -class MT5ForSequenceClassification(MT5PreTrainedModel): - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.transformer = MT5Model(config) - self.classification_head = MT5ClassificationHead(config) - - # Initialize weights and apply final processing - self.post_init() - - self.model_parallel = False - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.forward - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - decoder_head_mask: Optional[torch.Tensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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, Seq2SeqSequenceClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - Returns: - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if labels is not None: - use_cache = False - - if input_ids is None and inputs_embeds is not None: - raise NotImplementedError( - f"Passing input embeddings is currently not supported for {self.__class__.__name__}" - ) - - # Copied from models.bart.modeling_bart.BartModel.forward different to other models, T5 automatically creates - # decoder_input_ids from input_ids if no decoder_input_ids are provided - if decoder_input_ids is None and decoder_inputs_embeds is None: - if input_ids is None: - raise ValueError( - "If no `decoder_input_ids` or `decoder_inputs_embeds` are " - "passed, `input_ids` cannot be `None`. Please pass either " - "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." - ) - decoder_input_ids = self._shift_right(input_ids) - - outputs = self.transformer( - input_ids, - attention_mask=attention_mask, - decoder_input_ids=decoder_input_ids, - decoder_attention_mask=decoder_attention_mask, - head_mask=head_mask, - decoder_head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - encoder_outputs=encoder_outputs, - inputs_embeds=inputs_embeds, - decoder_inputs_embeds=decoder_inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - sequence_output = outputs[0] - - eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) - - if len(torch.unique_consecutive(eos_mask.sum(1))) > 1: - raise ValueError("All examples must have the same number of tokens.") - batch_size, _, hidden_size = sequence_output.shape - sentence_representation = sequence_output[eos_mask, :].view(batch_size, -1, hidden_size)[:, -1, :] - logits = self.classification_head(sentence_representation) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.config.num_labels == 1: - self.config.problem_type = "regression" - elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.config.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output - - return Seq2SeqSequenceClassifierOutput( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - decoder_hidden_states=outputs.decoder_hidden_states, - decoder_attentions=outputs.decoder_attentions, - cross_attentions=outputs.cross_attentions, - encoder_last_hidden_state=outputs.encoder_last_hidden_state, - encoder_hidden_states=outputs.encoder_hidden_states, - encoder_attentions=outputs.encoder_attentions, - ) - - -@add_start_docstrings( - """ - MT5 Model with a span classification head on top for extractive question-answering tasks like SQuAD (linear layers - on top of the hidden-states output to compute `span start logits` and `span end logits`). - """, - MT5_START_DOCSTRING, -) -class MT5ForQuestionAnswering(MT5PreTrainedModel): - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.model_dim = config.d_model - - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - self.decoder = MT5Stack(decoder_config, self.shared) - - self.num_labels = config.num_labels - self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) - - # Initialize weights and apply final processing - self.post_init() - - self.model_parallel = False - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_decoder - def get_decoder(self): - return self.decoder - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqQuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.forward - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - start_positions: Optional[torch.LongTensor] = None, - end_positions: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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[torch.FloatTensor], Seq2SeqQuestionAnsweringModelOutput]: - r""" - start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the start of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence - are not taken into account for computing the loss. - end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the end of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence - are not taken into account for computing the loss. - Returns: - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - use_cache = use_cache if use_cache is not None else self.config.use_cache - if start_positions is not None and end_positions is not None: - use_cache = False - - # Copied from models.bart.modeling_bart.BartModel.forward - # different to other models, T5 automatically creates decoder_input_ids from - # input_ids if no decoder_input_ids are provided - if decoder_input_ids is None and decoder_inputs_embeds is None: - if input_ids is None: - raise ValueError( - "If no `decoder_input_ids` or `decoder_inputs_embeds` are " - "passed, `input_ids` cannot be `None`. Please pass either " - "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." - ) - decoder_input_ids = self._shift_right(input_ids) - - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=None, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = decoder_outputs[0] - - logits = self.qa_outputs(sequence_output) - start_logits, end_logits = logits.split(1, dim=-1) - start_logits = start_logits.squeeze(-1).contiguous() - end_logits = end_logits.squeeze(-1).contiguous() - - total_loss = None - if start_positions is not None and end_positions is not None: - # If we are on multi-GPU, split add a dimension - if len(start_positions.size()) > 1: - start_positions = start_positions.squeeze(-1).to(start_logits.device) - if len(end_positions.size()) > 1: - end_positions = end_positions.squeeze(-1).to(end_logits.device) - # sometimes the start/end positions are outside our model inputs, we ignore these terms - ignored_index = start_logits.size(1) - start_positions = start_positions.clamp(0, ignored_index) - end_positions = end_positions.clamp(0, ignored_index) - - loss_fct = CrossEntropyLoss(ignore_index=ignored_index) - start_loss = loss_fct(start_logits, start_positions) - end_loss = loss_fct(end_logits, end_positions) - total_loss = (start_loss + end_loss) / 2 - - if not return_dict: - output = (start_logits, end_logits) + decoder_outputs[1:] + encoder_outputs - return ((total_loss,) + output) if total_loss is not None else output - - return Seq2SeqQuestionAnsweringModelOutput( - loss=total_loss, - start_logits=start_logits, - end_logits=end_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) diff --git a/module/MT5_Infer.py b/module/MT5_Infer.py deleted file mode 100644 index e028de3..0000000 --- a/module/MT5_Infer.py +++ /dev/null @@ -1,2611 +0,0 @@ -# coding=utf-8 -# Copyright 2020 Mesh TensorFlow authors, T5 Authors and HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" PyTorch mT5 model.""" - -import copy -import math -import os -import warnings -from typing import List, Optional, Tuple, Union - -import torch -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss - -from transformers.activations import ACT2FN -from transformers.modeling_outputs import ( - BaseModelOutput, - BaseModelOutputWithPastAndCrossAttentions, - Seq2SeqLMOutput, - Seq2SeqModelOutput, - Seq2SeqQuestionAnsweringModelOutput, - Seq2SeqSequenceClassifierOutput, -) -from transformers.modeling_utils import PreTrainedModel -from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer -from transformers.utils import ( - DUMMY_INPUTS, - DUMMY_MASK, - add_start_docstrings, - add_start_docstrings_to_model_forward, - is_torch_fx_proxy, - logging, - replace_return_docstrings, -) -from transformers.utils.model_parallel_utils import assert_device_map, get_device_map -from .configuration_mt5 import MT5Config -from transformers import AutoTokenizer - - -logger = logging.get_logger(__name__) - -_CONFIG_FOR_DOC = "MT5Config" -_CHECKPOINT_FOR_DOC = "mt5-small" - - -PARALLELIZE_DOCSTRING = r""" - This is an experimental feature and is a subject to change at a moment's notice. - - Uses a device map to distribute attention modules of the model across several devices. If no device map is given, - it will evenly distribute blocks across all devices. - - Args: - device_map (`Dict[int, list]`, optional, defaults to None): - A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always - automatically mapped to the first device (for esoteric reasons). That means that the first device should - have fewer attention modules mapped to it than other devices. For reference, the mt5 models have the - following number of attention modules: - - - mt5-small: 6 - - mt5-base: 12 - - mt5-large: 24 - - mt5-xl: 24 - - mt5-xxl: 24 - - Example: - - ```python - # Here is an example of a device map on a machine with 4 GPUs using mt5-xl, which has a total of 24 attention modules: - model = MT5ForConditionalGeneration.from_pretrained("mt5-xl") - device_map = { - 0: [0, 1, 2], - 1: [3, 4, 5, 6, 7, 8, 9], - 2: [10, 11, 12, 13, 14, 15, 16], - 3: [17, 18, 19, 20, 21, 22, 23], - } - model.parallelize(device_map) - ``` -""" -DEPARALLELIZE_DOCSTRING = r""" - Moves the model to cpu from a model parallel state. - - Example: - - ```python - # On a 4 GPU machine with mt5-xl: - model = MT5ForConditionalGeneration.from_pretrained("Mt5-xl") - device_map = { - 0: [0, 1, 2], - 1: [3, 4, 5, 6, 7, 8, 9], - 2: [10, 11, 12, 13, 14, 15, 16], - 3: [17, 18, 19, 20, 21, 22, 23], - } - model.parallelize(device_map) # Splits the model across several devices - model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache() - ``` -""" - - -# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->MT5 -class MT5LayerNorm(nn.Module): - def __init__(self, hidden_size, eps=1e-6): - """ - Construct a layernorm module in the MT5 style. No bias and no subtraction of mean. - """ - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - # MT5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean - # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus varience is calculated - # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for - # half-precision inputs is done in fp32 - - 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 - - -# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->MT5 -class MT5DenseActDense(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - self.wi = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.dropout = nn.Dropout(config.dropout_rate) - self.act = ACT2FN[config.dense_act_fn] - - def forward(self, hidden_states): - hidden_states = self.wi(hidden_states) - hidden_states = self.act(hidden_states) - hidden_states = self.dropout(hidden_states) - if ( - isinstance(self.wo.weight, torch.Tensor) - and hidden_states.dtype != self.wo.weight.dtype - and self.wo.weight.dtype != torch.int8 - ): - hidden_states = hidden_states.to(self.wo.weight.dtype) - hidden_states = self.wo(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->MT5 -class MT5DenseGatedActDense(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False) - self.wo = nn.Linear(config.d_ff, config.d_model, bias=False) - self.dropout = nn.Dropout(config.dropout_rate) - self.act = ACT2FN[config.dense_act_fn] - - def forward(self, hidden_states): - hidden_gelu = self.act(self.wi_0(hidden_states)) - hidden_linear = self.wi_1(hidden_states) - hidden_states = hidden_gelu * hidden_linear - hidden_states = self.dropout(hidden_states) - - # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32. - # See https://github.com/huggingface/transformers/issues/20287 - # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None`` - if ( - isinstance(self.wo.weight, torch.Tensor) - and hidden_states.dtype != self.wo.weight.dtype - and self.wo.weight.dtype != torch.int8 - ): - hidden_states = hidden_states.to(self.wo.weight.dtype) - - hidden_states = self.wo(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->MT5 -class MT5LayerFF(nn.Module): - def __init__(self, config: MT5Config): - super().__init__() - if config.is_gated_act: - self.DenseReluDense = MT5DenseGatedActDense(config) - else: - self.DenseReluDense = MT5DenseActDense(config) - - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward(self, hidden_states): - forwarded_states = self.layer_norm(hidden_states) - forwarded_states = self.DenseReluDense(forwarded_states) - hidden_states = hidden_states + self.dropout(forwarded_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->MT5 -class MT5Attention(nn.Module): - def __init__(self, config: MT5Config, has_relative_attention_bias=False): - super().__init__() - self.is_decoder = config.is_decoder - self.has_relative_attention_bias = has_relative_attention_bias - self.relative_attention_num_buckets = config.relative_attention_num_buckets - self.relative_attention_max_distance = config.relative_attention_max_distance - self.d_model = config.d_model - self.key_value_proj_dim = config.d_kv - self.n_heads = config.num_heads - self.dropout = config.dropout_rate - self.inner_dim = self.n_heads * self.key_value_proj_dim - - # Mesh TensorFlow initialization to avoid scaling before softmax - self.q = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.k = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.v = nn.Linear(self.d_model, self.inner_dim, bias=False) - self.o = nn.Linear(self.inner_dim, self.d_model, bias=False) - - if self.has_relative_attention_bias: - self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) - self.pruned_heads = set() - self.gradient_checkpointing = False - - def prune_heads(self, heads): - if len(heads) == 0: - return - heads, index = find_pruneable_heads_and_indices( - heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads - ) - # Prune linear layers - self.q = prune_linear_layer(self.q, index) - self.k = prune_linear_layer(self.k, index) - self.v = prune_linear_layer(self.v, index) - self.o = prune_linear_layer(self.o, index, dim=1) - # Update hyper params - self.n_heads = self.n_heads - len(heads) - self.inner_dim = self.key_value_proj_dim * self.n_heads - self.pruned_heads = self.pruned_heads.union(heads) - - @staticmethod - def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): - """ - Adapted from Mesh Tensorflow: - https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 - - Translate relative position to a bucket number for relative attention. The relative position is defined as - memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to - position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for - small absolute relative_position and larger buckets for larger absolute relative_positions. All relative - positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. - This should allow for more graceful generalization to longer sequences than the model has been trained on - - Args: - relative_position: an int32 Tensor - bidirectional: a boolean - whether the attention is bidirectional - num_buckets: an integer - max_distance: an integer - - Returns: - a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) - """ - relative_buckets = 0 - if bidirectional: - num_buckets //= 2 - relative_buckets += (relative_position > 0).to(torch.long) * num_buckets - relative_position = torch.abs(relative_position) - else: - relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) - # now relative_position is in the range [0, inf) - - # half of the buckets are for exact increments in positions - max_exact = num_buckets // 2 - is_small = relative_position < max_exact - - # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance - relative_position_if_large = max_exact + ( - torch.log(relative_position.float() / max_exact) - / math.log(max_distance / max_exact) - * (num_buckets - max_exact) - ).to(torch.long) - relative_position_if_large = torch.min( - relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) - ) - - relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) - return relative_buckets - - def compute_bias(self, query_length, key_length, device=None): - """Compute binned relative position bias""" - if device is None: - device = self.relative_attention_bias.weight.device - context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] - memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] - relative_position = memory_position - context_position # shape (query_length, key_length) - relative_position_bucket = self._relative_position_bucket( - relative_position, # shape (query_length, key_length) - bidirectional=(not self.is_decoder), - num_buckets=self.relative_attention_num_buckets, - max_distance=self.relative_attention_max_distance, - ) - values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) - values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length) - return values - - def forward( - self, - hidden_states, - mask=None, - key_value_states=None, - position_bias=None, - past_key_value=None, - layer_head_mask=None, - query_length=None, - use_cache=False, - output_attentions=False, - ): - """ - Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). - """ - # Input is (batch_size, seq_length, dim) - # Mask is (batch_size, key_length) (non-causal) or (batch_size, key_length, key_length) - # past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head) - batch_size, seq_length = hidden_states.shape[:2] - - real_seq_length = seq_length - - if past_key_value is not None: - if len(past_key_value) != 2: - raise ValueError( - f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states" - ) - real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length - - key_length = real_seq_length if key_value_states is None else key_value_states.shape[1] - - def shape(states): - """projection""" - return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) - - def unshape(states): - """reshape""" - return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim) - - def project(hidden_states, proj_layer, key_value_states, past_key_value): - """projects hidden states correctly to key/query states""" - if key_value_states is None: - # self-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(hidden_states)) - elif past_key_value is None: - # cross-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(key_value_states)) - - if past_key_value is not None: - if key_value_states is None: - # self-attn - # (batch_size, n_heads, key_length, dim_per_head) - hidden_states = torch.cat([past_key_value, hidden_states], dim=2) - elif past_key_value.shape[2] != key_value_states.shape[1]: - # checking that the `sequence_length` of the `past_key_value` is the same as - # the provided `key_value_states` to support prefix tuning - # cross-attn - # (batch_size, n_heads, seq_length, dim_per_head) - hidden_states = shape(proj_layer(key_value_states)) - else: - # cross-attn - hidden_states = past_key_value - return hidden_states - - # get query states - query_states = shape(self.q(hidden_states)) # (batch_size, n_heads, seq_length, dim_per_head) - - # get key/value states - key_states = project( - hidden_states, self.k, key_value_states, past_key_value[0] if past_key_value is not None else None - ) - value_states = project( - hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None - ) - - # compute scores - scores = torch.matmul( - query_states, key_states.transpose(3, 2) - ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 - - if position_bias is None: - if not self.has_relative_attention_bias: - position_bias = torch.zeros( - (1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype - ) - if self.gradient_checkpointing and self.training: - position_bias.requires_grad = True - else: - position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device) - - # if key and values are already calculated - # we want only the last query position bias - if past_key_value is not None: - position_bias = position_bias[:, :, -hidden_states.size(1) :, :] - - if mask is not None: - position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length) - - if self.pruned_heads: - mask = torch.ones(position_bias.shape[1]) - mask[list(self.pruned_heads)] = 0 - position_bias_masked = position_bias[:, mask.bool()] - else: - position_bias_masked = position_bias - - scores += position_bias_masked - attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as( - scores - ) # (batch_size, n_heads, seq_length, key_length) - attn_weights = nn.functional.dropout( - attn_weights, p=self.dropout, training=self.training - ) # (batch_size, n_heads, seq_length, key_length) - - # Mask heads if we want to - if layer_head_mask is not None: - attn_weights = attn_weights * layer_head_mask - - attn_output = unshape(torch.matmul(attn_weights, value_states)) # (batch_size, seq_length, dim) - attn_output = self.o(attn_output) - - present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None - outputs = (attn_output,) + (present_key_value_state,) + (position_bias,) - - if output_attentions: - outputs = outputs + (attn_weights,) - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->MT5 -class MT5LayerSelfAttention(nn.Module): - def __init__(self, config, has_relative_attention_bias=False): - super().__init__() - self.SelfAttention = MT5Attention(config, has_relative_attention_bias=has_relative_attention_bias) - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward( - self, - hidden_states, - attention_mask=None, - position_bias=None, - layer_head_mask=None, - past_key_value=None, - use_cache=False, - output_attentions=False, - ): - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.SelfAttention( - normed_hidden_states, - mask=attention_mask, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states = hidden_states + self.dropout(attention_output[0]) - outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->MT5 -class MT5LayerCrossAttention(nn.Module): - def __init__(self, config): - super().__init__() - self.EncDecAttention = MT5Attention(config, has_relative_attention_bias=False) - self.layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - def forward( - self, - hidden_states, - key_value_states, - attention_mask=None, - position_bias=None, - layer_head_mask=None, - past_key_value=None, - use_cache=False, - query_length=None, - output_attentions=False, - ): - normed_hidden_states = self.layer_norm(hidden_states) - attention_output = self.EncDecAttention( - normed_hidden_states, - mask=attention_mask, - key_value_states=key_value_states, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - query_length=query_length, - output_attentions=output_attentions, - ) - layer_output = hidden_states + self.dropout(attention_output[0]) - outputs = (layer_output,) + attention_output[1:] # add attentions if we output them - return outputs - - -# Copied from transformers.models.t5.modeling_t5.T5Block with T5->MT5 -class MT5Block(nn.Module): - def __init__(self, config, has_relative_attention_bias=False): - super().__init__() - self.is_decoder = config.is_decoder - self.layer = nn.ModuleList() - self.layer.append(MT5LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias)) - if self.is_decoder: - self.layer.append(MT5LayerCrossAttention(config)) - - self.layer.append(MT5LayerFF(config)) - - def forward( - self, - hidden_states, - attention_mask=None, - position_bias=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - encoder_decoder_position_bias=None, - layer_head_mask=None, - cross_attn_layer_head_mask=None, - past_key_value=None, - use_cache=False, - output_attentions=False, - return_dict=True, - ): - if past_key_value is not None: - if not self.is_decoder: - logger.warning("`past_key_values` is passed to the encoder. Please make sure this is intended.") - expected_num_past_key_values = 2 if encoder_hidden_states is None else 4 - - if len(past_key_value) != expected_num_past_key_values: - raise ValueError( - f"There should be {expected_num_past_key_values} past states. " - f"{'2 (past / key) for cross attention. ' if expected_num_past_key_values == 4 else ''}" - f"Got {len(past_key_value)} past key / value states" - ) - - self_attn_past_key_value = past_key_value[:2] - cross_attn_past_key_value = past_key_value[2:] - else: - self_attn_past_key_value, cross_attn_past_key_value = None, None - - self_attention_outputs = self.layer[0]( - hidden_states, - attention_mask=attention_mask, - position_bias=position_bias, - layer_head_mask=layer_head_mask, - past_key_value=self_attn_past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states, present_key_value_state = self_attention_outputs[:2] - attention_outputs = self_attention_outputs[2:] # Keep self-attention outputs and relative position weights - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - do_cross_attention = self.is_decoder and encoder_hidden_states is not None - if do_cross_attention: - # the actual query length is unknown for cross attention - # if using past key value states. Need to inject it here - if present_key_value_state is not None: - query_length = present_key_value_state[0].shape[2] - else: - query_length = None - - cross_attention_outputs = self.layer[1]( - hidden_states, - key_value_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - position_bias=encoder_decoder_position_bias, - layer_head_mask=cross_attn_layer_head_mask, - past_key_value=cross_attn_past_key_value, - query_length=query_length, - use_cache=use_cache, - output_attentions=output_attentions, - ) - hidden_states = cross_attention_outputs[0] - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - # Combine self attn and cross attn key value states - if present_key_value_state is not None: - present_key_value_state = present_key_value_state + cross_attention_outputs[1] - - # Keep cross-attention outputs and relative position weights - attention_outputs = attention_outputs + cross_attention_outputs[2:] - - # Apply Feed Forward layer - hidden_states = self.layer[-1](hidden_states) - - # clamp inf values to enable fp16 training - if hidden_states.dtype == torch.float16: - clamp_value = torch.where( - torch.isinf(hidden_states).any(), - torch.finfo(hidden_states.dtype).max - 1000, - torch.finfo(hidden_states.dtype).max, - ) - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - outputs = (hidden_states,) - - if use_cache: - outputs = outputs + (present_key_value_state,) + attention_outputs - else: - outputs = outputs + attention_outputs - - return outputs # hidden-states, present_key_value_states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) - - -def load_tf_weights_in_mt5(model, config, tf_checkpoint_path): - """Load tf checkpoints in a pytorch model.""" - try: - import re - - import numpy as np - import tensorflow as tf - except ImportError: - logger.error( - "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see " - "https://www.tensorflow.org/install/ for installation instructions." - ) - raise - tf_path = os.path.abspath(tf_checkpoint_path) - logger.info(f"Converting TensorFlow checkpoint from {tf_path}") - # Load weights from TF model - init_vars = tf.train.list_variables(tf_path) - names = [] - tf_weights = {} - for name, shape in init_vars: - logger.info(f"Loading TF weight {name} with shape {shape}") - array = tf.train.load_variable(tf_path, name) - names.append(name) - tf_weights[name] = array - - for txt_name in names: - name = txt_name.split("/") - # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v - # which are not required for using pretrained model - if any( - n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"] - for n in name - ): - logger.info(f"Skipping {'/'.join(name)}") - tf_weights.pop(txt_name, None) - continue - if "_slot_" in name[-1]: - logger.info(f"Skipping {'/'.join(name)}") - tf_weights.pop(txt_name, None) - continue - pointer = model - array = tf_weights[txt_name] - - for m_name in name: - if re.fullmatch(r"[A-Za-z]+_\d+", m_name): - scope_names = re.split(r"_(\d+)", m_name) - else: - scope_names = [m_name] - if scope_names[0] in ["kernel", "scale", "embedding"]: - pointer = getattr(pointer, "weight") - elif scope_names[0] == "self_attention": - pointer = getattr(pointer, "layer") - pointer = pointer[0] - elif scope_names[0] == "enc_dec_attention": - pointer = getattr(pointer, "layer") - pointer = pointer[1] - elif scope_names[0] == "dense_relu_dense": - pointer = getattr(pointer, "layer") - pointer = pointer[2] - elif scope_names[0] == "rms_norm": - if hasattr(pointer, "layer_norm"): - pointer = getattr(pointer, "layer_norm") - elif hasattr(pointer, "final_layer_norm"): - pointer = getattr(pointer, "final_layer_norm") - elif scope_names[0] == "scale": - pointer = getattr(pointer, "weight") - elif scope_names[0] == "output_bias" or scope_names[0] == "beta": - pointer = getattr(pointer, "bias") - elif scope_names[0] == "squad": - pointer = getattr(pointer, "classifier") - elif scope_names[0] == "decoder" and name[1] == "logits": - continue - elif scope_names[0] == "logits": - pointer = getattr(pointer, "lm_head") - elif scope_names[0] == "wi" and len(scope_names) > 1 and scope_names[1].isdigit(): - pointer = getattr(pointer, f"wi_{scope_names[1]}") - continue - else: - try: - pointer = getattr(pointer, scope_names[0]) - except AttributeError: - logger.info(f"Skipping {'/'.join(name)}") - continue - if len(scope_names) >= 2: - num = int(scope_names[1]) - pointer = pointer[num] - if scope_names[0] not in ["kernel", "scale", "embedding"]: - pointer = getattr(pointer, "weight") - if scope_names[0] != "embedding": - logger.info(f"Transposing numpy weight of shape {array.shape} for {name}") - array = np.transpose(array) - try: - assert ( - pointer.shape == array.shape - ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched" - except AssertionError as e: - e.args += (pointer.shape, array.shape) - raise - logger.info(f"Initialize PyTorch weight {name}") - pointer.data = torch.from_numpy(array.astype(np.float32)) - tf_weights.pop(txt_name, None) - - logger.info(f"Weights not copied to PyTorch model: {', '.join(tf_weights.keys())}.") - return model - - -# Copied from transformers.models.t5.modeling_t5.T5ClassificationHead with T5->MT5 -class MT5ClassificationHead(nn.Module): - """Head for sentence-level classification tasks.""" - - def __init__(self, config: MT5Config): - super().__init__() - self.dense = nn.Linear(config.d_model, config.d_model) - self.dropout = nn.Dropout(p=config.classifier_dropout) - self.out_proj = nn.Linear(config.d_model, config.num_labels) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.dropout(hidden_states) - hidden_states = self.dense(hidden_states) - hidden_states = torch.tanh(hidden_states) - hidden_states = self.dropout(hidden_states) - hidden_states = self.out_proj(hidden_states) - return hidden_states - - -# Copied from transformers.models.t5.modeling_t5.T5PreTrainedModel with T5->MT5, t5->mt5 -class MT5PreTrainedModel(PreTrainedModel): - """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. - """ - - config_class = MT5Config - load_tf_weights = load_tf_weights_in_mt5 - base_model_prefix = "transformer" - is_parallelizable = True - supports_gradient_checkpointing = True - _no_split_modules = ["MT5Block"] - _keep_in_fp32_modules = ["wo"] - - @property - def dummy_inputs(self): - input_ids = torch.tensor(DUMMY_INPUTS) - input_mask = torch.tensor(DUMMY_MASK) - dummy_inputs = { - "decoder_input_ids": input_ids, - "input_ids": input_ids, - "decoder_attention_mask": input_mask, - } - return dummy_inputs - - def _init_weights(self, module): - """Initialize the weights""" - factor = self.config.initializer_factor # Used for testing weights initialization - if isinstance(module, MT5LayerNorm): - module.weight.data.fill_(factor * 1.0) - elif isinstance( - module, - (MT5Model, MT5ForConditionalGeneration, MT5EncoderModel, MT5ForQuestionAnswering), - ): - # Mesh TensorFlow embeddings initialization - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624 - module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) - if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: - module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0) - if hasattr(module, "qa_outputs"): - module.qa_outputs.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - module.qa_outputs.bias.data.zero_() - elif isinstance(module, MT5ClassificationHead): - module.dense.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.dense, "bias") and module.dense.bias is not None: - module.dense.bias.data.zero_() - module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: - module.out_proj.bias.data.zero_() - elif isinstance(module, MT5DenseActDense): - # Mesh TensorFlow FF initialization - # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56 - # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89 - module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi, "bias") and module.wi.bias is not None: - module.wi.bias.data.zero_() - module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) - if hasattr(module.wo, "bias") and module.wo.bias is not None: - module.wo.bias.data.zero_() - elif isinstance(module, MT5DenseGatedActDense): - module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: - module.wi_0.bias.data.zero_() - module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) - if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: - module.wi_1.bias.data.zero_() - module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) - if hasattr(module.wo, "bias") and module.wo.bias is not None: - module.wo.bias.data.zero_() - elif isinstance(module, MT5Attention): - # Mesh TensorFlow attention initialization to avoid scaling before softmax - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136 - d_model = self.config.d_model - key_value_proj_dim = self.config.d_kv - n_heads = self.config.num_heads - module.q.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5)) - module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) - module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) - module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) - if module.has_relative_attention_bias: - module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5)) - - def _shift_right(self, input_ids): - decoder_start_token_id = self.config.decoder_start_token_id - pad_token_id = self.config.pad_token_id - - if decoder_start_token_id is None: - raise ValueError( - "self.model.config.decoder_start_token_id has to be defined. In MT5 it is usually set to the pad_token_id. " - "See MT5 docs for more information." - ) - - # shift inputs to the right - if is_torch_fx_proxy(input_ids): - # Item assignment is not supported natively for proxies. - shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id) - shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) - else: - shifted_input_ids = input_ids.new_zeros(input_ids.shape) - shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() - shifted_input_ids[..., 0] = decoder_start_token_id - - if pad_token_id is None: - raise ValueError("self.model.config.pad_token_id has to be defined.") - # replace possible -100 values in labels by `pad_token_id` - shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) - - return shifted_input_ids - - -# Copied from transformers.models.t5.modeling_t5.T5Stack with T5->MT5 -class MT5Stack(MT5PreTrainedModel): - def __init__(self, config, embed_tokens=None): - super().__init__(config) - - self.embed_tokens = embed_tokens - self.is_decoder = config.is_decoder - - self.block = nn.ModuleList( - [MT5Block(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)] - ) - self.final_layer_norm = MT5LayerNorm(config.d_model, eps=config.layer_norm_epsilon) - self.dropout = nn.Dropout(config.dropout_rate) - - # Initialize weights and apply final processing - self.post_init() - # Model parallel - self.model_parallel = False - self.device_map = None - self.gradient_checkpointing = False - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - def parallelize(self, device_map=None): - warnings.warn( - "`MT5Stack.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" - " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," - " 'block.1': 1, ...}", - FutureWarning, - ) - # Check validity of device_map - self.device_map = ( - get_device_map(len(self.block), range(torch.cuda.device_count())) if device_map is None else device_map - ) - assert_device_map(self.device_map, len(self.block)) - self.model_parallel = True - self.first_device = "cpu" if "cpu" in self.device_map.keys() else "cuda:" + str(min(self.device_map.keys())) - self.last_device = "cuda:" + str(max(self.device_map.keys())) - # Load onto devices - for k, v in self.device_map.items(): - for layer in v: - cuda_device = "cuda:" + str(k) - self.block[layer] = self.block[layer].to(cuda_device) - - # Set embed_tokens to first layer - self.embed_tokens = self.embed_tokens.to(self.first_device) - # Set final layer norm to last device - self.final_layer_norm = self.final_layer_norm.to(self.last_device) - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.model_parallel = False - self.device_map = None - self.first_device = "cpu" - self.last_device = "cpu" - for i in range(len(self.block)): - self.block[i] = self.block[i].to("cpu") - self.embed_tokens = self.embed_tokens.to("cpu") - self.final_layer_norm = self.final_layer_norm.to("cpu") - torch.cuda.empty_cache() - - def get_input_embeddings(self): - return self.embed_tokens - - def set_input_embeddings(self, new_embeddings): - self.embed_tokens = new_embeddings - - def forward( - self, - input_ids=None, - attention_mask=None, - encoder_hidden_states=None, - encoder_attention_mask=None, - inputs_embeds=None, - head_mask=None, - cross_attn_head_mask=None, - past_key_values=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - ): - # Model parallel - if self.model_parallel: - torch.cuda.set_device(self.first_device) - self.embed_tokens = self.embed_tokens.to(self.first_device) - use_cache = use_cache if use_cache is not None else self.config.use_cache - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if input_ids is not None and inputs_embeds is not None: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError( - f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" - ) - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - err_msg_prefix = "decoder_" if self.is_decoder else "" - raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") - - if inputs_embeds is None: - if self.embed_tokens is None: - raise ValueError("You have to initialize the model with valid token embeddings") - inputs_embeds = self.embed_tokens(input_ids) - # print(input_ids) - batch_size, seq_length = input_shape - - # required mask seq length can be calculated via length of past - mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length - - if use_cache is True: - if not self.is_decoder: - raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") - - if attention_mask is None: - attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) - if self.is_decoder and encoder_attention_mask is None and encoder_hidden_states is not None: - encoder_seq_length = encoder_hidden_states.shape[1] - encoder_attention_mask = torch.ones( - batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long - ) - - # initialize past_key_values with `None` if past does not exist - if past_key_values is None: - past_key_values = [None] * len(self.block) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape) - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.is_decoder and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = None - - 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 - - # Prepare head mask if needed - head_mask = self.get_head_mask(head_mask, self.config.num_layers) - cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers) - present_key_value_states = () if use_cache else None - all_hidden_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - all_cross_attentions = () if (output_attentions and self.is_decoder) else None - position_bias = None - encoder_decoder_position_bias = None - - hidden_states = self.dropout(inputs_embeds) - - for i, (layer_module, past_key_value) in enumerate(zip(self.block, past_key_values)): - layer_head_mask = head_mask[i] - cross_attn_layer_head_mask = cross_attn_head_mask[i] - # Model parallel - if self.model_parallel: - torch.cuda.set_device(hidden_states.device) - # Ensure that attention_mask is always on the same device as hidden_states - if attention_mask is not None: - attention_mask = attention_mask.to(hidden_states.device) - if position_bias is not None: - position_bias = position_bias.to(hidden_states.device) - if encoder_hidden_states is not None: - encoder_hidden_states = encoder_hidden_states.to(hidden_states.device) - if encoder_extended_attention_mask is not None: - encoder_extended_attention_mask = encoder_extended_attention_mask.to(hidden_states.device) - if encoder_decoder_position_bias is not None: - encoder_decoder_position_bias = encoder_decoder_position_bias.to(hidden_states.device) - if layer_head_mask is not None: - layer_head_mask = layer_head_mask.to(hidden_states.device) - if cross_attn_layer_head_mask is not None: - cross_attn_layer_head_mask = cross_attn_layer_head_mask.to(hidden_states.device) - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - layer_module.forward, - hidden_states, - extended_attention_mask, - position_bias, - encoder_hidden_states, - encoder_extended_attention_mask, - encoder_decoder_position_bias, - layer_head_mask, - cross_attn_layer_head_mask, - None, # past_key_value is always None with gradient checkpointing - use_cache, - output_attentions, - ) - - else: - layer_outputs = layer_module( - hidden_states, - attention_mask=extended_attention_mask, - position_bias=position_bias, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - encoder_decoder_position_bias=encoder_decoder_position_bias, - layer_head_mask=layer_head_mask, - cross_attn_layer_head_mask=cross_attn_layer_head_mask, - past_key_value=past_key_value, - use_cache=use_cache, - output_attentions=output_attentions, - ) - - # layer_outputs is a tuple with: - # hidden-states, key-value-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) - if use_cache is False: - layer_outputs = layer_outputs[:1] + (None,) + layer_outputs[1:] - - hidden_states, present_key_value_state = layer_outputs[:2] - - # We share the position biases between the layers - the first layer store them - # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights), - # (cross-attention position bias), (cross-attention weights) - position_bias = layer_outputs[2] - if self.is_decoder and encoder_hidden_states is not None: - encoder_decoder_position_bias = layer_outputs[4 if output_attentions else 3] - # append next layer key value states - if use_cache: - present_key_value_states = present_key_value_states + (present_key_value_state,) - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[3],) - if self.is_decoder: - all_cross_attentions = all_cross_attentions + (layer_outputs[5],) - - # Model Parallel: If it's the last layer for that device, put things on the next device - if self.model_parallel: - for k, v in self.device_map.items(): - if i == v[-1] and "cuda:" + str(k) != self.last_device: - hidden_states = hidden_states.to("cuda:" + str(k + 1)) - - hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.dropout(hidden_states) - - # Add last layer - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - present_key_value_states, - all_hidden_states, - all_attentions, - all_cross_attentions, - ] - if v is not None - ) - # print("all_hidden_states", all_hidden_states.shape) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=present_key_value_states, - hidden_states=all_hidden_states, - attentions=all_attentions, - cross_attentions=all_cross_attentions, - ) - - -MT5_START_DOCSTRING = r""" - - The MT5 model was proposed in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text - Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan - Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu. It's an encoder decoder transformer pre-trained in a - text-to-text denoising generative setting. - - This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the - library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads - etc.) - - This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. - Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage - and behavior. - - Parameters: - config ([`MT5Config`]): Model configuration class with all the parameters of the model. - Initializing with a config file does not load the weights associated with the model, only the - configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. -""" - -MT5_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you - should be able to pad the inputs on both the right and the left. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for detail. - - [What are input IDs?](../glossary#input-ids) - - To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). - attention_mask (`torch.FloatTensor` 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) - decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): - Indices of decoder input sequence tokens in the vocabulary. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for details. - - [What are decoder input IDs?](../glossary#decoder-input-ids) - - MT5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` - is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). - - To know more on how to prepare `decoder_input_ids` for pretraining take a look at [MT5 - Training](./mt5#training). - decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*): - Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also - be used by default. - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules in the encoder. Mask values selected in `[0, - 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0, - 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in - `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): - Tuple consists of (`last_hidden_state`, `optional`: *hidden_states*, `optional`: *attentions*) - `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` is a sequence of hidden states at - the output of the last layer of the encoder. Used in the cross-attention of the decoder. - past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - 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. - decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): - Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded - representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be - input (see `past_key_values`). This is useful if you want more control over how to convert - `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. - - If `decoder_input_ids` and `decoder_inputs_embeds` are both unset, `decoder_inputs_embeds` takes the value - of `inputs_embeds`. - - 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. -""" - -MT5_ENCODER_INPUTS_DOCSTRING = r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. MT5 is a model with relative position embeddings so you - should be able to pad the inputs on both the right and the left. - - Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for detail. - - To know more on how to prepare `input_ids` for pretraining take a look a [MT5 Training](./mt5#training). - attention_mask (`torch.FloatTensor` 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) - head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): - Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - 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. - 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. -""" - -# Warning message for FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask -__HEAD_MASK_WARNING_MSG = """ -The input argument `head_mask` was split into two arguments `head_mask` and `decoder_head_mask`. Currently, -`decoder_head_mask` is set to copy `head_mask`, but this feature is deprecated and will be removed in future versions. -If you do not want to use any `decoder_head_mask` now, please set `decoder_head_mask = torch.ones(num_layers, -num_heads)`. -""" - - -@add_start_docstrings( - "The bare MT5 Model transformer outputting raw hidden-states without any specific head on top.", - MT5_START_DOCSTRING, -) -class MT5Model(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5Model, AutoTokenizer - - >>> model = MT5Model.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> summary = "Weiter Verhandlung in Syrien." - >>> inputs = tokenizer(article, return_tensors="pt") - >>> labels = tokenizer(text_target=summary, return_tensors="pt") - - >>> outputs = model(input_ids=inputs["input_ids"], decoder_input_ids=labels["input_ids"]) - >>> hidden_states = outputs.last_hidden_state - ```""" - model_type = "mt5" - config_class = MT5Config - _keys_to_ignore_on_load_missing = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5Model.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - self.decoder = MT5Stack(decoder_config, self.shared) - - # Initialize weights and apply final processing - self.post_init() - - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5Model.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5Model.parallelize` is deprecated and will be removed in v5 of Transformers, you should load your model" - " with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'encoder.block.0':" - " 0, 'encoder.block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.decoder.parallelize(self.device_map) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5Model.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.decoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.decoder = self.decoder.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5Model.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5Model.get_decoder - def get_decoder(self): - return self.decoder - - # Copied from transformers.models.t5.modeling_t5.T5Model._prune_heads - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.layer[layer].attention.prune_heads(heads) - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5Model.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - inputs_embeds: Optional[torch.Tensor] = None, - decoder_inputs_embeds: Optional[torch.Tensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], Seq2SeqModelOutput]: - r""" - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MT5Model - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5Model.from_pretrained("mt5-small") - - >>> input_ids = tokenizer( - ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1 - - >>> # preprocess: Prepend decoder_input_ids with start token which is pad token for MT5Model. - >>> # This is not needed for torch's MT5ForConditionalGeneration as it does this internally using labels arg. - >>> decoder_input_ids = model._shift_right(decoder_input_ids) - - >>> # forward pass - >>> outputs = model(input_ids=input_ids, decoder_input_ids=decoder_input_ids) - >>> last_hidden_states = outputs.last_hidden_state - ```""" - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - hidden_states = hidden_states.to(self.decoder.first_device) - if decoder_input_ids is not None: - decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) - if attention_mask is not None: - attention_mask = attention_mask.to(self.decoder.first_device) - if decoder_attention_mask is not None: - decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - if not return_dict: - return decoder_outputs + encoder_outputs - - return Seq2SeqModelOutput( - last_hidden_state=decoder_outputs.last_hidden_state, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) - - -import torch.nn.functional as F -class PointAttention(nn.Module): - def __init__(self): - super().__init__() - self.hidden_size = 768 - # self.linear_layer = nn.Linear(250778, 1) - - - def forward(self,hidden1_states,hidden2_states,lm_logits,tags = [[0,0,1,2,-1,3,4,5,-1]]): - taglist = [] - finalBiasList = [] - sentence_len = hidden1_states.shape[1] - output_sentence_len = hidden2_states.shape[1] - for tag in tags: - vect,bias = merge_weights(tag,sentence_len) - # print('sentence_len', sentence_len) - taglist.append(vect) - biasList = [] - for i in range(output_sentence_len): - biasList.append(bias) - finalBiasList.append(biasList) - merge = torch.tensor(taglist, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) # 8 512[实则300] 512[指明需要的行] - finalBias = torch.tensor(finalBiasList, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) - # 8 512[实则300] 728 - - # print(hidden1_states.shape) #[8, 512, 768] - attn_weights = torch.matmul(merge,hidden1_states) # hidden1_states1 8 512【300】 728 - - attn_weights1 = torch.matmul(hidden2_states, attn_weights.transpose(1, 2))+finalBias # hidden2_states 8 512 768 attn_weights 8 512 768 attn_weights1 8 512 512 - score = torch.cat([lm_logits,attn_weights1], dim=2) - # probabilities = F.softmax(score, dim=-1) - # lm_logits[:, :, -30:] = attn_weights1 - return score - - -def get_vect(i,j,sentence_len=512): - vect = [] - for t in range(sentence_len): - vect.append(0) - t = i - count = j-i - while(t', '', '',''} - -def isdigit(s): - if s.replace('.', '', 1).isdigit() or (s.find(".")!=-1 and s.replace('.', '', 1) == ""): - return True - else: - return False -def label_words(original_sentence,lst): - i = 0 - indices = [] - num = 0 - # print(lst) - while i < len(lst): - if isdigit(lst[i]): - # 当前元素是数字字符串，检查下一个元素 - number_str = lst[i] - indices.append(num) - i += 1 - while i < len(lst) and isdigit(lst[i]): - # 下一个元素也是数字字符串，进行拼接 - number_str += lst[i] - indices.append(num) - i += 1 - num+=1 - elif lst[i] in special_tokens: # 特殊字符检查 - indices.append(-1) - num+=1 - i += 1 - else: - indices.append(num) - num+=1 - i += 1 - # print(indices) - return indices - -# 为列表中的每个元素分配序号 -# 中文没考虑数字 -# def label_words(original_sentence,word_list): -# indices = [] -# num = 0 -# for index, word in enumerate(word_list): -# if word in special_tokens: # 特殊字符检查 -# indices.append(-1) -# num += 1 -# else: -# indices.append(index-num) - -# return indices - -# 英文拼接 -# def label_words(original_sentence,tokenized_words): -# labels = [] -# label = 0 -# # word_index = 0 -# current_word = "" -# lastword = "" -# # 判断是不是最后一个句号 -# period = original_sentence.count(".") -# period_num = 1 -# # print(tokenized_words) -# for word in tokenized_words: -# # 如果当前单词是非字母（如空格或特殊符号），单独分配一个标签 -# pmFlag = False -# if not word.isalpha() and not word.isdigit() and word.find("'") == -1 and not is_valid_time_format(word): -# # .不能丢。如果是句子中的最后一个，则... 如果不是句子中的最后一个。则... -# if word == "." and period_num < period: -# period_num += 1 -# elif word == "." and period_num == period: -# labels.append(label+1) -# continue -# else: -# labels.append(-1) -# lastword = word -# continue - -# # 累加拆分的单词片段 -# current_word += word - -# if lastword.isdigit() and original_sentence.find(lastword + word) != -1 and (word.lower() == "pm" or word.lower() == "am"): -# pmFlag = True -# # 检查当前累加的单词片段是否与原句中的单词相匹配 -# # print(original_sentence.split()) - -# if current_word in original_sentence.split() or pmFlag == True: -# if pmFlag == True: -# label+=1 - -# if current_word in original_sentence.split(): -# current_word = "" - -# labels.append(label) -# # print(current_word) -# else: -# labels.append(label) - -# # 如果当前累加的片段已经形成了完整的单词，则增加标签值 -# if current_word == "": -# label += 1 -# lastword = word -# # return changeModel(labels) -# return labels - -def label_tokenized_words_corrected(tokenizer,token_list): - # tokenizer = AutoTokenizer.from_pretrained("/home/lzx/T5-base/tokenizer1/") - original_sentences = [tokenizer.decode(t[:-1]) for t in token_list] - lsts = [] - for tokens in token_list: - decoded_words = [tokenizer.decode([token_id]) for token_id in tokens] - lsts.append(decoded_words) - # print(lsts) - # 遍历分词后的单词 - llables = [] - for i, tokenized_words in enumerate(lsts): - original_sentence = original_sentences[i] - labels = label_words(original_sentence,tokenized_words) - llables.append(labels) - return llables - -@add_start_docstrings("""MT5 Model with a `language modeling` head on top.""", MT5_START_DOCSTRING) -class MT5ForConditionalGeneration(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5ForConditionalGeneration, AutoTokenizer - - >>> model = MT5ForConditionalGeneration.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> summary = "Weiter Verhandlung in Syrien." - >>> inputs = tokenizer(article, text_target=summary, return_tensors="pt") - - >>> outputs = model(**inputs) - >>> loss = outputs.loss - ```""" - - model_type = "mt5" - config_class = MT5Config - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.model_dim = config.d_model - - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - # print("decoder_config", decoder_config) - self.decoder = MT5Stack(decoder_config, self.shared) - - # -512 - self.lm_head = nn.Linear(config.d_model, config.vocab_size-512, bias=False) - # Initialize weights and apply final processing - self.post_init() - self.pointNet = PointAttention() - # /home/lzx/T5-base-lora/tokenizer2/ - # /home/lzx/T5-base/model_cl_multi/mt5-base-trained-final-save - self.tokenizer = AutoTokenizer.from_pretrained("/home/lzx/T5-base-lora/tokenizer2/") - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5ForConditionalGeneration.parallelize` is deprecated and will be removed in v5 of Transformers, you" - " should load your model with `device_map='balanced'` in the call to `from_pretrained`. You can also" - " provide your own `device_map` but it needs to be a dictionary module_name to device, so for instance" - " {'encoder.block.0': 0, 'encoder.block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.decoder.parallelize(self.device_map) - self.lm_head = self.lm_head.to(self.decoder.first_device) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.decoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.decoder = self.decoder.to("cpu") - self.lm_head = self.lm_head.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.set_output_embeddings - def set_output_embeddings(self, new_embeddings): - self.lm_head = new_embeddings - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_output_embeddings - def get_output_embeddings(self): - return self.lm_head - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.get_decoder - def get_decoder(self): - return self.decoder - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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[torch.FloatTensor], Seq2SeqLMOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ..., - config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for - labels in `[0, ..., config.vocab_size]` - - Returns: - - Examples: - - ```python - >>> from transformers import AutoTokenizer, MT5ForConditionalGeneration - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5ForConditionalGeneration.from_pretrained("mt5-small") - - >>> # training - >>> input_ids = tokenizer("The walks in park", return_tensors="pt").input_ids - >>> labels = tokenizer(" cute dog the ", return_tensors="pt").input_ids - >>> outputs = model(input_ids=input_ids, labels=labels) - >>> loss = outputs.loss - >>> logits = outputs.logits - - >>> # inference - >>> input_ids = tokenizer( - ... "summarize: studies have shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> outputs = model.generate(input_ids) - >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True)) - >>> # studies have shown that owning a dog is good for you. - ```""" - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - self.tags = label_tokenized_words_corrected(self.tokenizer,input_ids) # [batchsize,sequence] - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - # Convert encoder inputs in embeddings if needed - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - - if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None: - # get decoder inputs from shifting lm labels to the right - decoder_input_ids = self._shift_right(labels) - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.decoder.first_device) - hidden_states = hidden_states.to(self.decoder.first_device) - if decoder_input_ids is not None: - decoder_input_ids = decoder_input_ids.to(self.decoder.first_device) - if attention_mask is not None: - attention_mask = attention_mask.to(self.decoder.first_device) - if decoder_attention_mask is not None: - decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device) - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=past_key_values, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions="True", - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = decoder_outputs[0] - - # Set device for model parallelism - if self.model_parallel: - torch.cuda.set_device(self.encoder.first_device) - self.lm_head = self.lm_head.to(self.encoder.first_device) - sequence_output = sequence_output.to(self.lm_head.weight.device) - - if self.config.tie_word_embeddings: - # Rescale output before projecting on vocab - # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586 - sequence_output = sequence_output * (self.model_dim**-0.5) - - lm_logits1 = self.lm_head(sequence_output) - # print(hidden_states.shape) - # print(sequence_output.shape) - # print(lm_logits.shape) - # cross_attentions = decoder_outputs["cross_attentions"][-1] - # cross_attentions = torch.mean(cross_attentions, dim=1) - # print(cross_attentions.shape) - # print(cross_attentions.shape) - # lm_logits = torch.cat([lm_logits, cross_attentions], dim=2) - # print(lm_logits.shape) - # eps = 1e-7 - # lm_logits = torch.log(lm_logits + eps) - - lm_logits = self.pointNet(hidden_states,sequence_output,lm_logits1,self.tags) - # print(labels) - # print(lm_logits) - - loss = None - if labels is not None: - # # 创建NLLLoss对象 - # loss_fct = nn.NLLLoss() - loss_fct = CrossEntropyLoss(ignore_index=-100) - # move labels to correct device to enable PP - labels = labels.to(lm_logits.device) - # print(labels.shape) - # print(lm_logits.view(-1, lm_logits.size(-1)).shape) - - loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666 - - if not return_dict: - output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs - return ((loss,) + output) if loss is not None else output - - return Seq2SeqLMOutput( - loss=loss, - logits=lm_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_inputs_for_generation - def prepare_inputs_for_generation( - self, - input_ids, - past_key_values=None, - attention_mask=None, - head_mask=None, - decoder_head_mask=None, - decoder_attention_mask=None, - cross_attn_head_mask=None, - use_cache=None, - encoder_outputs=None, - **kwargs, - ): - # cut decoder_input_ids if past_key_values is used - if past_key_values is not None: - past_length = past_key_values[0][0].shape[2] - - # Some generation methods already pass only the last input ID - if input_ids.shape[1] > past_length: - remove_prefix_length = past_length - else: - # Default to old behavior: keep only final ID - remove_prefix_length = input_ids.shape[1] - 1 - - input_ids = input_ids[:, remove_prefix_length:] - - return { - "decoder_input_ids": input_ids, - "past_key_values": past_key_values, - "encoder_outputs": encoder_outputs, - "attention_mask": attention_mask, - "head_mask": head_mask, - "decoder_head_mask": decoder_head_mask, - "decoder_attention_mask": decoder_attention_mask, - "cross_attn_head_mask": cross_attn_head_mask, - "use_cache": use_cache, - } - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration.prepare_decoder_input_ids_from_labels - def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor): - return self._shift_right(labels) - - # Copied from transformers.models.t5.modeling_t5.T5ForConditionalGeneration._reorder_cache - def _reorder_cache(self, past_key_values, beam_idx): - # if decoder past is not included in output - # speedy decoding is disabled and no need to reorder - if past_key_values is None: - logger.warning("You might want to consider setting `use_cache=True` to speed up decoding") - return past_key_values - - reordered_decoder_past = () - for layer_past_states in past_key_values: - # get the correct batch idx from layer past batch dim - # batch dim of `past` is at 2nd position - reordered_layer_past_states = () - for layer_past_state in layer_past_states: - # need to set correct `past` for each of the four key / value states - reordered_layer_past_states = reordered_layer_past_states + ( - layer_past_state.index_select(0, beam_idx.to(layer_past_state.device)), - ) - - if reordered_layer_past_states[0].shape != layer_past_states[0].shape: - raise ValueError( - f"reordered_layer_past_states[0] shape {reordered_layer_past_states[0].shape} and layer_past_states[0] shape {layer_past_states[0].shape} mismatched" - ) - if len(reordered_layer_past_states) != len(layer_past_states): - raise ValueError( - f"length of reordered_layer_past_states {len(reordered_layer_past_states)} and length of layer_past_states {len(layer_past_states)} mismatched" - ) - - reordered_decoder_past = reordered_decoder_past + (reordered_layer_past_states,) - return reordered_decoder_past - - -@add_start_docstrings( - "The bare MT5 Model transformer outputting encoder's raw hidden-states without any specific head on top.", - MT5_START_DOCSTRING, -) -class MT5EncoderModel(MT5PreTrainedModel): - r""" - Examples: - - ```python - >>> from transformers import MT5EncoderModel, AutoTokenizer - - >>> model = MT5EncoderModel.from_pretrained("google/mt5-small") - >>> tokenizer = AutoTokenizer.from_pretrained("google/mt5-small") - >>> article = "UN Offizier sagt, dass weiter verhandelt werden muss in Syrien." - >>> input_ids = tokenizer(article, return_tensors="pt").input_ids - >>> outputs = model(input_ids) - >>> hidden_state = outputs.last_hidden_state - ```""" - - model_type = "mt5" - config_class = MT5Config - _tied_weights_keys = ["encoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - # Initialize weights and apply final processing - self.post_init() - - # Model parallel - self.model_parallel = False - self.device_map = None - - @add_start_docstrings(PARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.parallelize - def parallelize(self, device_map=None): - warnings.warn( - "`T5EncoderModel.parallelize` is deprecated and will be removed in v5 of Transformers, you should load" - " your model with `device_map='balanced'` in the call to `from_pretrained`. You can also provide your own" - " `device_map` but it needs to be a dictionary module_name to device, so for instance {'block.0': 0," - " 'block.1': 1, ...}", - FutureWarning, - ) - self.device_map = ( - get_device_map(len(self.encoder.block), range(torch.cuda.device_count())) - if device_map is None - else device_map - ) - assert_device_map(self.device_map, len(self.encoder.block)) - self.encoder.parallelize(self.device_map) - self.model_parallel = True - - @add_start_docstrings(DEPARALLELIZE_DOCSTRING) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.deparallelize - def deparallelize(self): - warnings.warn( - "Like `parallelize`, `deparallelize` is deprecated and will be removed in v5 of Transformers.", - FutureWarning, - ) - self.encoder.deparallelize() - self.encoder = self.encoder.to("cpu") - self.model_parallel = False - self.device_map = None - torch.cuda.empty_cache() - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel._prune_heads - def _prune_heads(self, heads_to_prune): - """ - Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base - class PreTrainedModel - """ - for layer, heads in heads_to_prune.items(): - self.encoder.block[layer].layer[0].SelfAttention.prune_heads(heads) - - @add_start_docstrings_to_model_forward(MT5_ENCODER_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=BaseModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5EncoderModel.forward with T5->MT5, t5->mt5 - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]: - r""" - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MT5EncoderModel - - >>> tokenizer = AutoTokenizer.from_pretrained("mt5-small") - >>> model = MT5EncoderModel.from_pretrained("mt5-small") - >>> input_ids = tokenizer( - ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" - ... ).input_ids # Batch size 1 - >>> outputs = model(input_ids=input_ids) - >>> last_hidden_states = outputs.last_hidden_state - ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - return encoder_outputs - - -@add_start_docstrings( - """ - MT5 model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE - tasks. - """, - MT5_START_DOCSTRING, -) -class MT5ForSequenceClassification(MT5PreTrainedModel): - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.transformer = MT5Model(config) - self.classification_head = MT5ClassificationHead(config) - - # Initialize weights and apply final processing - self.post_init() - - self.model_parallel = False - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqSequenceClassifierOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForSequenceClassification.forward - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - decoder_head_mask: Optional[torch.Tensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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, Seq2SeqSequenceClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - Returns: - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - if labels is not None: - use_cache = False - - if input_ids is None and inputs_embeds is not None: - raise NotImplementedError( - f"Passing input embeddings is currently not supported for {self.__class__.__name__}" - ) - - # Copied from models.bart.modeling_bart.BartModel.forward different to other models, T5 automatically creates - # decoder_input_ids from input_ids if no decoder_input_ids are provided - if decoder_input_ids is None and decoder_inputs_embeds is None: - if input_ids is None: - raise ValueError( - "If no `decoder_input_ids` or `decoder_inputs_embeds` are " - "passed, `input_ids` cannot be `None`. Please pass either " - "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." - ) - decoder_input_ids = self._shift_right(input_ids) - - outputs = self.transformer( - input_ids, - attention_mask=attention_mask, - decoder_input_ids=decoder_input_ids, - decoder_attention_mask=decoder_attention_mask, - head_mask=head_mask, - decoder_head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - encoder_outputs=encoder_outputs, - inputs_embeds=inputs_embeds, - decoder_inputs_embeds=decoder_inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - sequence_output = outputs[0] - - eos_mask = input_ids.eq(self.config.eos_token_id).to(sequence_output.device) - - if len(torch.unique_consecutive(eos_mask.sum(1))) > 1: - raise ValueError("All examples must have the same number of tokens.") - batch_size, _, hidden_size = sequence_output.shape - sentence_representation = sequence_output[eos_mask, :].view(batch_size, -1, hidden_size)[:, -1, :] - logits = self.classification_head(sentence_representation) - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.config.num_labels == 1: - self.config.problem_type = "regression" - elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.config.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output - - return Seq2SeqSequenceClassifierOutput( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - decoder_hidden_states=outputs.decoder_hidden_states, - decoder_attentions=outputs.decoder_attentions, - cross_attentions=outputs.cross_attentions, - encoder_last_hidden_state=outputs.encoder_last_hidden_state, - encoder_hidden_states=outputs.encoder_hidden_states, - encoder_attentions=outputs.encoder_attentions, - ) - - -@add_start_docstrings( - """ - MT5 Model with a span classification head on top for extractive question-answering tasks like SQuAD (linear layers - on top of the hidden-states output to compute `span start logits` and `span end logits`). - """, - MT5_START_DOCSTRING, -) -class MT5ForQuestionAnswering(MT5PreTrainedModel): - _keys_to_ignore_on_load_unexpected = ["decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight"] - _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"] - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.__init__ with T5->MT5 - def __init__(self, config: MT5Config): - super().__init__(config) - self.model_dim = config.d_model - - self.shared = nn.Embedding(config.vocab_size, config.d_model) - - encoder_config = copy.deepcopy(config) - encoder_config.is_decoder = False - encoder_config.use_cache = False - encoder_config.is_encoder_decoder = False - self.encoder = MT5Stack(encoder_config, self.shared) - - decoder_config = copy.deepcopy(config) - decoder_config.is_decoder = True - decoder_config.is_encoder_decoder = False - decoder_config.num_layers = config.num_decoder_layers - self.decoder = MT5Stack(decoder_config, self.shared) - - self.num_labels = config.num_labels - self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) - - # Initialize weights and apply final processing - self.post_init() - - self.model_parallel = False - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_input_embeddings - def get_input_embeddings(self): - return self.shared - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.set_input_embeddings - def set_input_embeddings(self, new_embeddings): - self.shared = new_embeddings - self.encoder.set_input_embeddings(new_embeddings) - self.decoder.set_input_embeddings(new_embeddings) - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_encoder - def get_encoder(self): - return self.encoder - - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.get_decoder - def get_decoder(self): - return self.decoder - - @add_start_docstrings_to_model_forward(MT5_INPUTS_DOCSTRING) - @replace_return_docstrings(output_type=Seq2SeqQuestionAnsweringModelOutput, config_class=_CONFIG_FOR_DOC) - # Copied from transformers.models.t5.modeling_t5.T5ForQuestionAnswering.forward - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - decoder_head_mask: Optional[torch.FloatTensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, - start_positions: Optional[torch.LongTensor] = None, - end_positions: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - decoder_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[torch.FloatTensor], Seq2SeqQuestionAnsweringModelOutput]: - r""" - start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the start of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence - are not taken into account for computing the loss. - end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the end of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (*sequence_length*). Position outside of the sequence - are not taken into account for computing the loss. - Returns: - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - use_cache = use_cache if use_cache is not None else self.config.use_cache - if start_positions is not None and end_positions is not None: - use_cache = False - - # Copied from models.bart.modeling_bart.BartModel.forward - # different to other models, T5 automatically creates decoder_input_ids from - # input_ids if no decoder_input_ids are provided - if decoder_input_ids is None and decoder_inputs_embeds is None: - if input_ids is None: - raise ValueError( - "If no `decoder_input_ids` or `decoder_inputs_embeds` are " - "passed, `input_ids` cannot be `None`. Please pass either " - "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`." - ) - decoder_input_ids = self._shift_right(input_ids) - - 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 - - # FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask - if head_mask is not None and decoder_head_mask is None: - if self.config.num_layers == self.config.num_decoder_layers: - warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning) - decoder_head_mask = head_mask - - # Encode if needed (training, first prediction pass) - if encoder_outputs is None: - encoder_outputs = self.encoder( - input_ids=input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None, - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - hidden_states = encoder_outputs[0] - - # Decode - decoder_outputs = self.decoder( - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - inputs_embeds=decoder_inputs_embeds, - past_key_values=None, - encoder_hidden_states=hidden_states, - encoder_attention_mask=attention_mask, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = decoder_outputs[0] - - logits = self.qa_outputs(sequence_output) - start_logits, end_logits = logits.split(1, dim=-1) - start_logits = start_logits.squeeze(-1).contiguous() - end_logits = end_logits.squeeze(-1).contiguous() - - total_loss = None - if start_positions is not None and end_positions is not None: - # If we are on multi-GPU, split add a dimension - if len(start_positions.size()) > 1: - start_positions = start_positions.squeeze(-1).to(start_logits.device) - if len(end_positions.size()) > 1: - end_positions = end_positions.squeeze(-1).to(end_logits.device) - # sometimes the start/end positions are outside our model inputs, we ignore these terms - ignored_index = start_logits.size(1) - start_positions = start_positions.clamp(0, ignored_index) - end_positions = end_positions.clamp(0, ignored_index) - - loss_fct = CrossEntropyLoss(ignore_index=ignored_index) - start_loss = loss_fct(start_logits, start_positions) - end_loss = loss_fct(end_logits, end_positions) - total_loss = (start_loss + end_loss) / 2 - - if not return_dict: - output = (start_logits, end_logits) + decoder_outputs[1:] + encoder_outputs - return ((total_loss,) + output) if total_loss is not None else output - - return Seq2SeqQuestionAnsweringModelOutput( - loss=total_loss, - start_logits=start_logits, - end_logits=end_logits, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_outputs.last_hidden_state, - encoder_hidden_states=encoder_outputs.hidden_states, - encoder_attentions=encoder_outputs.attentions, - ) diff --git a/module/configuration_mt5.py b/module/configuration_mt5.py deleted file mode 100644 index 7b72aa4..0000000 --- a/module/configuration_mt5.py +++ /dev/null @@ -1,178 +0,0 @@ -# coding=utf-8 -# Copyright 2020, The T5 Authors and HuggingFace Inc. -# -# 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. -""" mT5 model configuration""" -from typing import Mapping - -from transformers.configuration_utils import PretrainedConfig -from transformers.onnx import OnnxSeq2SeqConfigWithPast -from transformers.utils import logging - - -logger = logging.get_logger(__name__) - - -class MT5Config(PretrainedConfig): - r""" - This is the configuration class to store the configuration of a [`MT5Model`] or a [`TFMT5Model`]. It is used to - instantiate a mT5 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 mT5 - [google/mt5-small](https://huggingface.co/google/mt5-small) architecture. - - Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PretrainedConfig`] for more information. - - Arguments: - vocab_size (`int`, *optional*, defaults to 250112): - Vocabulary size of the T5 model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`T5Model`] or [`TFT5Model`]. - d_model (`int`, *optional*, defaults to 512): - Size of the encoder layers and the pooler layer. - d_kv (`int`, *optional*, defaults to 64): - Size of the key, query, value projections per attention head. `d_kv` has to be equal to `d_model // - num_heads`. - d_ff (`int`, *optional*, defaults to 1024): - Size of the intermediate feed forward layer in each `T5Block`. - num_layers (`int`, *optional*, defaults to 8): - Number of hidden layers in the Transformer encoder. - num_decoder_layers (`int`, *optional*): - Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set. - num_heads (`int`, *optional*, defaults to 6): - Number of attention heads for each attention layer in the Transformer encoder. - relative_attention_num_buckets (`int`, *optional*, defaults to 32): - The number of buckets to use for each attention layer. - relative_attention_max_distance (`int`, *optional*, defaults to 128): - The maximum distance of the longer sequences for the bucket separation. - dropout_rate (`float`, *optional*, defaults to 0.1): - The ratio for all dropout layers. - layer_norm_eps (`float`, *optional*, defaults to 1e-6): - The epsilon used by the layer normalization layers. - initializer_factor (`float`, *optional*, defaults to 1): - A factor for initializing all weight matrices (should be kept to 1, used internally for initialization - testing). - feed_forward_proj (`string`, *optional*, defaults to `"gated-gelu"`): - Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. - use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). - """ - model_type = "mt5" - keys_to_ignore_at_inference = ["past_key_values"] - - def __init__( - self, - vocab_size=250112, - d_model=512, - d_kv=64, - d_ff=1024, - num_layers=8, - num_decoder_layers=None, - num_heads=6, - relative_attention_num_buckets=32, - relative_attention_max_distance=128, - dropout_rate=0.1, - layer_norm_epsilon=1e-6, - initializer_factor=1.0, - feed_forward_proj="gated-gelu", - is_encoder_decoder=True, - use_cache=True, - tokenizer_class="T5Tokenizer", - tie_word_embeddings=False, - pad_token_id=0, - eos_token_id=1, - decoder_start_token_id=0, - **kwargs, - ): - super().__init__( - is_encoder_decoder=is_encoder_decoder, - tokenizer_class=tokenizer_class, - tie_word_embeddings=tie_word_embeddings, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - decoder_start_token_id=decoder_start_token_id, - **kwargs, - ) - self.vocab_size = vocab_size - self.d_model = d_model - self.d_kv = d_kv - self.d_ff = d_ff - self.num_layers = num_layers - self.num_decoder_layers = ( - num_decoder_layers if num_decoder_layers is not None else self.num_layers - ) # default = symmetry - self.num_heads = num_heads - self.relative_attention_num_buckets = relative_attention_num_buckets - self.relative_attention_max_distance = relative_attention_max_distance - self.dropout_rate = dropout_rate - self.layer_norm_epsilon = layer_norm_epsilon - self.initializer_factor = initializer_factor - self.feed_forward_proj = feed_forward_proj - self.use_cache = use_cache - - act_info = self.feed_forward_proj.split("-") - self.dense_act_fn = act_info[-1] - self.is_gated_act = act_info[0] == "gated" - - if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2: - raise ValueError( - f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer." - "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. " - "'gated-gelu' or 'relu'" - ) - - # for backwards compatibility - if feed_forward_proj == "gated-gelu": - self.dense_act_fn = "gelu_new" - - @property - def hidden_size(self): - return self.d_model - - @property - def num_attention_heads(self): - return self.num_heads - - @property - def num_hidden_layers(self): - return self.num_layers - - -class MT5OnnxConfig(OnnxSeq2SeqConfigWithPast): - @property - # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs - def inputs(self) -> Mapping[str, Mapping[int, str]]: - common_inputs = { - "input_ids": {0: "batch", 1: "encoder_sequence"}, - "attention_mask": {0: "batch", 1: "encoder_sequence"}, - } - if self.use_past: - common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence" - common_inputs["decoder_input_ids"] = {0: "batch"} - common_inputs["decoder_attention_mask"] = {0: "batch", 1: "past_decoder_sequence + sequence"} - else: - common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"} - common_inputs["decoder_attention_mask"] = {0: "batch", 1: "decoder_sequence"} - - if self.use_past: - self.fill_with_past_key_values_(common_inputs, direction="inputs") - - return common_inputs - - @property - # Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset - def default_onnx_opset(self) -> int: - return 13 - - @property - def atol_for_validation(self) -> float: - return 5e-4 diff --git a/module/util.py b/module/util.py deleted file mode 100644 index fd3768e..0000000 --- a/module/util.py +++ /dev/null @@ -1,163 +0,0 @@ -import torch.nn.functional as F -import torch - -class PointAttention(nn.Module): - def __init__(self): - super().__init__() - self.hidden_size = 768 - # self.linear_layer = nn.Linear(250778, 1) - - - def forward(self,hidden1_states,hidden2_states,lm_logits,tags = [[0,0,1,2,-1,3,4,5,-1]]): - taglist = [] - finalBiasList = [] - sentence_len = hidden1_states.shape[1] - output_sentence_len = hidden2_states.shape[1] - for tag in tags: - vect,bias = merge_weights(tag,sentence_len) - # print('sentence_len', sentence_len) - taglist.append(vect) - biasList = [] - for i in range(output_sentence_len): - biasList.append(bias) - finalBiasList.append(biasList) - merge = torch.tensor(taglist, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) # 8 512[实则300] 512[指明需要的行] - finalBias = torch.tensor(finalBiasList, dtype=torch.float32, requires_grad=True).to(hidden1_states.device) - # 8 512[实则300] 728 - - # print(hidden1_states.shape) #[8, 512, 768] - attn_weights = torch.matmul(merge,hidden1_states) # hidden1_states1 8 512【300】 728 - - attn_weights1 = torch.matmul(hidden2_states, attn_weights.transpose(1, 2))+finalBias # hidden2_states 8 512 768 attn_weights 8 512 768 attn_weights1 8 512 512 - score = torch.cat([lm_logits,attn_weights1], dim=2) - # probabilities = F.softmax(score, dim=-1) - # lm_logits[:, :, -30:] = attn_weights1 - return score - - -def get_vect(i,j,sentence_len=512): - vect = [] - for t in range(sentence_len): - vect.append(0) - t = i - count = j-i - while(t', '', '',''} - -def isdigit(s): - if s.replace('.', '', 1).isdigit() or (s.find(".")!=-1 and s.replace('.', '', 1) == ""): - return True - else: - return False -def label_words(original_sentence,lst): - i = 0 - indices = [] - num = 0 - # print(lst) - while i < len(lst): - if isdigit(lst[i]): - # 当前元素是数字字符串，检查下一个元素 - number_str = lst[i] - indices.append(num) - i += 1 - while i < len(lst) and isdigit(lst[i]): - # 下一个元素也是数字字符串，进行拼接 - number_str += lst[i] - indices.append(num) - i += 1 - num+=1 - elif lst[i] in special_tokens: # 特殊字符检查 - indices.append(-1) - num+=1 - i += 1 - else: - indices.append(num) - num+=1 - i += 1 - # print(indices) - return indices - -def label_tokenized_words_corrected(tokenizer,token_list): - # tokenizer = AutoTokenizer.from_pretrained("/home/lzx/T5-base/tokenizer1/") - original_sentences = [tokenizer.decode(t[:-1]) for t in token_list] - lsts = [] - for tokens in token_list: - decoded_words = [tokenizer.decode([token_id]) for token_id in tokens] - lsts.append(decoded_words) - # print(lsts) - # 遍历分词后的单词 - llables = [] - for i, tokenized_words in enumerate(lsts): - original_sentence = original_sentences[i] - labels = label_words(original_sentence,tokenized_words) - llables.append(labels) - return llables diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..641e169 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,175 @@ +absl-py==2.0.0 +accelerate==0.24.1 +aiohttp==3.8.6 +aiosignal==1.3.1 +annotated-types==0.6.0 +anyio==3.7.1 +astunparse==1.6.3 +async-timeout==4.0.3 +attrs==23.1.0 +backoff==2.2.1 +beartype==0.16.4 +bitsandbytes==0.41.2.post2 +blinker==1.7.0 +cachetools==5.3.2 +certifi==2023.7.22 +charset-normalizer==3.3.2 +click==8.1.7 +cmake==3.25.0 +cohere==4.32 +dataclasses-json==0.6.1 +datasets==2.14.6 +deepspeed==0.11.1 +dill==0.3.7 +distro==1.9.0 +einops==0.7.0 +evaluate==0.4.1 +exceptiongroup==1.1.3 +experta==1.9.4 +fastavro==1.8.2 +fastcore==1.5.29 +filelock==3.13.1 +Flask==3.0.1 +flatbuffers==23.5.26 +frozendict==1.2 +frozenlist==1.4.0 +fsspec==2023.10.0 +ftfy==6.1.1 +fuzzywuzzy==0.18.0 +gast==0.5.4 +gevent==23.9.1 +google-api-core==2.12.0 +google-api-python-client==2.106.0 +google-auth==2.23.4 +google-auth-httplib2==0.1.1 +google-auth-oauthlib==1.2.0 +google-pasta==0.2.0 +googleapis-common-protos==1.61.0 +graphviz==0.20.1 +greenlet==3.0.1 +grpcio==1.60.0 +h11==0.14.0 +h5py==3.10.0 +hanlp==2.1.0b63 +hanlp_common==0.0.21 +hanlp_downloader==0.0.25 +hanlp_trie==0.0.5 +hjson==3.1.0 +httpcore==1.0.7 +httplib2==0.22.0 +httpx==0.28.1 +huggingface-hub==0.17.3 +idna==3.4 +importlib-metadata==6.8.0 +itsdangerous==2.1.2 +jaraco.context==4.3.0 +jieba==0.42.1 +Jinja2==3.1.2 +jiter==0.8.2 +joblib==1.3.2 +json5==0.10.0 +jsonlines==4.0.0 +jsonpatch==1.33 +jsonpointer==2.4 +keras==2.15.0 +langchain==0.0.329 +langsmith==0.0.56 +Levenshtein==0.26.1 +libclang==16.0.6 +lit==15.0.7 +Markdown==3.5.1 +MarkupSafe==2.1.3 +marshmallow==3.20.1 +ml-dtypes==0.2.0 +more-itertools==10.1.0 +mpmath==1.3.0 +multidict==6.0.4 +multiprocess==0.70.15 +mypy-extensions==1.0.0 +networkx==3.2.1 +ninja==1.11.1.1 +numpy==1.26.1 +nvidia-cublas-cu12==12.1.3.1 +nvidia-cuda-cupti-cu12==12.1.105 +nvidia-cuda-nvrtc-cu12==12.1.105 +nvidia-cuda-runtime-cu12==12.1.105 +nvidia-cudnn-cu12==8.9.2.26 +nvidia-cufft-cu12==11.0.2.54 +nvidia-curand-cu12==10.3.2.106 +nvidia-cusolver-cu12==11.4.5.107 +nvidia-cusparse-cu12==12.1.0.106 +nvidia-ml-py==12.560.30 +nvidia-nccl-cu12==2.18.1 +nvidia-nvjitlink-cu12==12.3.52 +nvidia-nvtx-cu12==12.1.105 +oauthlib==3.2.2 +openai==1.59.6 +opt-einsum==3.3.0 +packaging==23.2 +pandas==2.1.2 +phrasetree==0.0.9 +Pillow==10.1.0 +protobuf==4.23.4 +psutil==5.9.6 +py-cpuinfo==9.0.0 +pyarrow==14.0.0 +pyasn1==0.5.0 +pyasn1-modules==0.3.0 +pydantic==1.10.13 +pydantic_core==2.10.1 +pynvml==12.0.0 +pyparsing==3.1.1 +python-dateutil==2.8.2 +pytz==2023.3.post1 +PyYAML==6.0.1 +RapidFuzz==3.10.1 +regex==2023.10.3 +requests==2.31.0 +requests-oauthlib==1.3.1 +responses==0.18.0 +rouge==1.0.1 +rsa==4.9 +safetensors==0.4.0 +schema==0.6.7 +scikit-learn==1.3.2 +scipy==1.11.3 +sentencepiece==0.1.99 +six==1.16.0 +sniffio==1.3.0 +SQLAlchemy==2.0.23 +sympy==1.12 +tenacity==8.2.3 +tensorboard==2.15.1 +tensorboard-data-server==0.7.2 +tensorflow==2.15.0.post1 +tensorflow-estimator==2.15.0 +tensorflow-io-gcs-filesystem==0.34.0 +termcolor==2.4.0 +threadpoolctl==3.2.0 +tokenizers==0.13.3 +toolformer-pytorch==0.0.29 +toposort==1.5 +torch==2.1.0 +torchvision==0.16.0 +torchviz==0.0.2 +tqdm==4.66.1 +transformers==4.31.0 +triton==2.1.0 +typing-inspect==0.9.0 +typing_extensions==4.12.2 +tzdata==2023.3 +uritemplate==4.1.1 +urllib3==2.0.7 +wcwidth==0.2.9 +Werkzeug==3.0.1 +wolframalpha==5.0.0 +wrapt==1.14.1 +x-clip==0.14.4 +xformers==0.0.22.post7 +xmltodict==0.13.0 +xxhash==3.4.1 +yarl==1.9.2 +zipp==3.17.0 +zope.event==5.0 +zope.interface==6.1 +ray==2.41.0 \ No newline at end of file diff --git a/test_model.py b/test_model.py new file mode 100644 index 0000000..d9c8eaa --- /dev/null +++ b/test_model.py @@ -0,0 +1,395 @@ +import os +import warnings + +import sys +from typing import Literal + +sys.path.append('/usr/src/app/semantic_parsing/code') # 按理说不需要，但服务器有短暂的时间出现索引错误，稳妥起见补一个ad hoc的处理 + +import yaml + +# os.environ["CUDA_VISIBLE_DEVICES"] = "6" + +import json +import torch +import torch.optim as optim +from torch.utils.data import DataLoader + +from datasets import load_dataset +from transformers import get_scheduler + +import re +import string +# import jsonlines +from tqdm import tqdm + +def format_time_string(input_string): + english_punctuation = string.punctuation.replace(':', '').replace("'", '') + time_pattern = r'(\d{1,2})(:)?(\d{2})?(am|pm|AM|PM)?' + def replacer(match): + hour = match.group(1) + colon = " : " if match.group(2) else "" + minute = match.group(3) if match.group(3) else "" + period = f" {match.group(4)}" if match.group(4) else "" + return f"{hour}{colon}{minute}{period}" + formatted_string = re.sub(time_pattern, replacer, input_string) + + def insert_spaces_around_punctuation(formatted_string): + i = 0 + while i < len(formatted_string): + if formatted_string[i] in english_punctuation: + if i > 0 and formatted_string[i - 1] != " ": + formatted_string = formatted_string[:i] + " " + formatted_string[i:] + i += 1 # 跳过新插入的空格 + if (i + 1) < len(formatted_string) and formatted_string[i + 1] != " ": + formatted_string = formatted_string[:i + 1] + " " + formatted_string[i + 1:] + i += 1 # 跳过新插入的空格 + i += 1 + return formatted_string + + formatted_string = insert_spaces_around_punctuation(formatted_string) + return re.sub(r"([a-zA-Z0-9])'s", r"\1 's", formatted_string) + + +import re +def in_input(content, input_str): + if f" {content} " not in input_str: + if f" {content}?" not in input_str and f" {content}." not in input_str and f" {content}," not in input_str: + if f" {content}" not in input_str: + return False + elif input_str.endswith(f" {content}") == False: + return False + + return True + +def edit_label(examples): + input = examples["utterance"] + output = examples["semantic_parse"] + + pattern = r"(?<=\[SL:WEATHER_TEMPERATURE_UNIT\s)(.*?)(?=\s\])" + if "[SL:WEATHER_TEMPERATURE_UNIT" in output: + match = re.search(pattern, output) + content = match.group(0) + if in_input(content, input) == False: + in_input1 = content.capitalize() + if in_input(in_input1, input) == True: + output = re.sub(pattern, in_input1, output) + elif in_input(content.lower(), input) == True: + output = re.sub(pattern, content.lower(), output) + else: + return examples + + examples["semantic_parse"] = output + + return examples + +from utils.remove_non_slot_leaf import remove_non_slot_leaf_nodes +from utils.sort_label import sort_string + +def ptr_change(examples): + examples["semantic_parse"] = edit_label(examples)["semantic_parse"] + + examples["utterance"] = format_time_string(examples["utterance"]) + # 删标签 + examples["semantic_parse"] = remove_non_slot_leaf_nodes(examples["semantic_parse"]) + # 排序 + examples["semantic_parse"] = sort_string(examples["semantic_parse"]) + # print(examples) + return examples + + +def preprocess_dataset(dataset): + dataset = dataset.map(ptr_change, load_from_cache_file=False) + return dataset + + +def tokenize_function(examples, tokenizer): + tokenized_inputs = tokenizer(examples['utterance'], padding='max_length', truncation=True, max_length=64, + return_tensors="pt") + tokenized_labels = tokenizer(examples['semantic_parse'], padding='max_length', truncation=True, max_length=64, + return_tensors="pt") + + tokenized_inputs['labels'] = tokenized_labels['input_ids'] + return tokenized_inputs + + +def tokenizer_dataset(tokenizer, dataset): + tokenized_datasets = dataset.map(tokenize_function, fn_kwargs={"tokenizer": tokenizer}, load_from_cache_file=False) + tokenized_datasets = tokenized_datasets.remove_columns(["utterance"]) + tokenized_datasets = tokenized_datasets.remove_columns(["semantic_parse"]) + if "domain" in tokenized_datasets.column_names: + tokenized_datasets = tokenized_datasets.remove_columns(["domain"]) + + return tokenized_datasets + + +def mycollate_trainer(examples): + """ + 这里面不应当包含多余的key + """ + remove_keys = [] + for example in examples: + for key in example: + try: + example[key] = torch.tensor(example[key]) + except Exception as e: + remove_keys.append(key) + + for key in set(remove_keys): + for example in examples: + del example[key] + + batch = {} + for key in examples[0]: + batch[key] = torch.stack([example[key][0] for example in examples]) + + return batch + +def extract_slots(label): + # 按空格切分字符串 + parts = label.split()[1:-1] + result = [] # 用来存储最外层的SL标签 + stack = [] # 用来处理嵌套情况 + current_slot = "" # 当前的SL标签 + for part in parts: + if part.startswith('[sl:'): # 如果是SL标签 + if current_slot == "": # 如果sl为空，表示是最外层的SL标签 + current_slot = part # 开始新的最外层SL标签 + else: # 如果栈不为空，表示在处理嵌套的SL标签 + current_slot += ' ' + part # 拼接当前的SL标签 + stack.append('[') # 入栈表示有一个嵌套 + elif part.startswith('[in:'): + current_slot += ' ' + part + stack.append('[') + elif part == ']' or part.find("]") != -1: # 处理结束的右括号 + try: + stack.pop() # 出栈 + current_slot += ' ' + part # 拼接右括号 + if not stack: # 栈为空，表示这个 ] 就是当前最外层的slot + result.append(current_slot) # 将最外层的SL标签存入结果 + current_slot = "" # 重置当前SL标签 + except: + continue + else: + current_slot += ' ' + part # 拼接标签内容 + # 返回最终结果 + return result + +def get_f1(label, predict): + predict_slot_lst = extract_slots(predict.lower()) + label_lst = extract_slots(label.lower()) + + true_index = [] + # 算出来正确的有几个 + true_predict = 0 + for predict_slot in predict_slot_lst: + predict_fix = predict_slot.split()[0] + predict_content = predict_slot.split()[1:-1] + predict_content_str = " ".join(predict_content) + for l, label_slot in enumerate(label_lst): + if predict_fix not in label_slot or l in true_index: + continue + label_content = label_slot.split()[1:-1] + label_content_str = " ".join(label_content) + if predict_content_str.replace(" ","") in label_content_str.replace(" ","") or label_content_str.replace(" ","") in predict_content_str.replace(" ",""): + true_predict += 1 + true_index.append(l) + + result = { + "true_predict": true_predict, + "predict_slot_lst_len": len(predict_slot_lst), + "label_lst_len": len(label_lst), + } + + return result + +def remove_prepositions(text): + prepositions = ["in", "on", "at", "by", "with", "about","the", + "for", "under", "over", "between", "during", + "through", "of"] + words = text.split() + filtered_words = [word for word in words if word.lower() not in prepositions] + return " ".join(filtered_words) + + +def test_model(model, tokenizer, dataset, args=None, device=None, id_sign="", train_machine: Literal['our', 'ray'] = 'our'): + # 在 GPU 上测试（如果可用） + if args != None: + device = args.device + correct = 0 + data_length = 0 + + true_predict = 0 + predict_slot_lst_len = 0 + label_lst_len = 0 + # DataLoader 用于批量测试 + # template = "<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n" + # for data in dataset["validation"]: + # print(data) + # data["semantic_parse"] = template.format(content=data["semantic_parse"]) + + if train_machine == 'our': + for key in dataset: + dataset[key] = tokenizer_dataset(tokenizer, preprocess_dataset(dataset[key])) + + test_loader = DataLoader(dataset["validation"], batch_size=128, collate_fn=mycollate_trainer) # 你可以调整 batch_size + + for i, batch in enumerate(test_loader): + batch = {k: v.to(device) for k, v in batch.items()} + with torch.no_grad(): + generate_ids = model.generate(batch['input_ids'], max_length=128) + + input_ids = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True) + decoded_predicted = tokenizer.batch_decode(generate_ids, skip_special_tokens=True) + decoded_labels = tokenizer.batch_decode(batch["labels"], skip_special_tokens=True) + k = 0 + + num_correct_sentences = 0 + for input_id, pred, label in zip(input_ids, decoded_predicted, decoded_labels): + k += 1 + result = get_f1(label, pred) + + pred = remove_prepositions(pred).lower() + label = remove_prepositions(label.lower()) + + true_predict += result["true_predict"] + predict_slot_lst_len += result["predict_slot_lst_len"] + label_lst_len += result["label_lst_len"] + + pred = pred[pred.index('['):] if '[' in pred else pred + print("结果：", pred, '分隔符', label) + + pred = pred.replace(" ", "").strip() + pred = _extract_pred(pred) + label = label.replace(" ", "").strip() + + if pred == label: + num_correct_sentences += 1 + + data_length += k + correct += num_correct_sentences + + accuracy = correct / data_length + acc = true_predict / predict_slot_lst_len if predict_slot_lst_len > 0 else 0 + recall = true_predict / label_lst_len if label_lst_len > 0 else 0 + f1_score = 2 * acc * recall / (acc + recall) if (acc + recall) > 0 else 0 + print(f"本轮结果({id_sign}): {accuracy, f1_score}") + return accuracy, f1_score + + +# def _load_dataset(path): +# dataset = load_dataset(path) +# for key in dataset: +# dataset[key] = tokenizer_dataset(tokenizer, preprocess_dataset(dataset[key])) +# +# return dataset + + +def _extract_pred(s: str) -> str: + s = s.strip() + if not s: + return s + + if s[0] != '[': + warnings.warn("s[0] must be a [") + return s + + result = '' + depth = 0 + + for i, c in enumerate(s): + result += c + if c == '[': + depth += 1 + elif c == ']': + depth -= 1 + if depth == 0: + return result + + return result + + +def test_all_models(model_output_dir, model_name, save_path, **kwargs) -> list[str]: + tasks = ['event', 'reminder', 'weather'] + dataset_types = ['top', 'our_data', 'our_add_top'] + exp_settings = ['SPIS25', 'SPIS50', 'full'] + train_strategy = 'full' + + preprocessed_data_save_dir = '../LLaMA-Factory/data' + + results = [] + + for task in tasks: + for dataset_type in dataset_types: + for exp_setting in exp_settings: + if dataset_type == 'our_data' and exp_setting == 'SPIS50': + file_name = f"{task}_{dataset_type}_{'SPIS25'}" + else: + file_name = f"{task}_{dataset_type}_{exp_setting}" + model_dir = f'{model_output_dir}/{model_name}/{train_strategy}/sft_{file_name}' + + # try: + tokenizer = AutoTokenizer.from_pretrained(model_dir, padding_side='left') + model = AutoModelForCausalLM.from_pretrained(model_dir).cuda() + + data_path = os.path.join(preprocessed_data_save_dir, task, dataset_type, exp_setting) + dataset = load_dataset(data_path) # fixme: 这里缺少处理 + + acc, f1 = test_model(model, tokenizer, dataset, device='cuda') + + del model + # except Exception as e: + # print("测试错误", model_dir, e) + # acc, f1 = -1, -1 + + result = f"{task} {dataset_type} {exp_setting} {model_name}: acc is {acc} and f1 is {f1}" + + with open(result_save_path, 'a', encoding='utf-8') as f: + f.write(result + '\n') + + results.append(result) + + + return results + + +if __name__ == '__main__': + # Load model directly + from transformers import AutoTokenizer, AutoModelForCausalLM + + # model_name = 'qwen3_8b' + # model_path = f'/usr/src/app/saves/{model_name}/full/sft_event_top_SPIS25' + # + # tokenizer = AutoTokenizer.from_pretrained(model_path) + # model = AutoModelForCausalLM.from_pretrained(model_path).cuda() + # + # data_path = '/home/cyz/data/semantic_parsing/LLaMA-Factory/data/event/top/SPIS25' + # dataset = load_dataset(data_path) + # acc, f1 = test_model(model, tokenizer, dataset, device='cuda') + # print(acc, f1) + + machine_path = '/home/cyz/data/semantic_parsing/semantic_parsing' # '/usr/src/app' + saved_models_dir = '/var/lib/docker/data/cyz/semantic_parsing' # '/usr/src/app' + + model_info = { + 'qwen3_8b': { + 'model_output_dir': f'{saved_models_dir}/saves', + 'model_name': 'qwen3_8b', + }, + 'llama3_8b': { + 'model_output_dir': f'{saved_models_dir}/saves', + 'model_name': 'llama3_8b', + }, + 'Phi_4': { + 'model_output_dir': f'{saved_models_dir}/saves', + 'model_name': 'Phi_4', + } + } + + result_save_path = f'{machine_path}/saves/results.txt' + with open(result_save_path, "w", encoding="utf-8") as f: + pass + + for key in model_info: + test_all_models(save_path=result_save_path, **model_info[key]) + diff --git a/train/ray_train.py b/train/ray_train.py new file mode 100644 index 0000000..f721d60 --- /dev/null +++ b/train/ray_train.py @@ -0,0 +1,90 @@ +import ray +from datasets import load_dataset +from ray import tune +from ray.tune import CLIReporter +from torch import optim + +from transformers import AutoModelForSeq2SeqLM + +from test_model import test_model, preprocess_dataset +from .self_train import train_model_self_train + + +def get_optimizer(optimizer, model, args): + arguments = ['lr'] + input_args = {arg: getattr(args, arg) for arg in arguments} + + if optimizer == 'Adam': + return optim.Adam(model.parameters(), **input_args) + + raise ValueError(f"Unknown optimizer: {optimizer}") + + +def train_tune(config, args, tokenizer, dataset, optimizer, model): + """ + Ray Tune 的试验函数：根据 config 中的超参数训练模型、评估，并报告指标给 Tune。 + """ + # 更新超参数 + args.batch_size = config["batch_size"] + args.learn_rate = config["learn_rate"] + args.max_length = config["max_length"] + args.epoch = config["epoch"] + + # 构造优化器（注意每个试验内部新建 optimizer） + optimizer = get_optimizer(args.optimizer, model, args) + + # 训练模型 + model = train_model_self_train(model, tokenizer, optimizer, dataset, args) + + # 加载测试模型 + dataset = preprocess_dataset(load_dataset(args.train_dataset_dir)) + + # 在训练后评估模型 + accuracy, f1 = test_model(model, tokenizer, dataset, args, id_sign='train_tune', train_machine='ray') + + from main import _save_metrics_to_file + _save_metrics_to_file(accuracy, f1, args, args.train_dataset_dir) + + # 报告指标给 Ray Tune（Ray Tune 会根据这些指标进行调度和选择最佳试验） + tune.report({"accuracy": accuracy, "f1": f1}) + + + +def tune_hyperparameters_ray(tokenizer, dataset, args, model_save_path, optimizer, model): + """ + 利用 Ray Tune 进行超参数网格搜索，每个试验分配一个 GPU。 + """ + # 初始化 Ray（若 Ray 已经初始化，可忽略 ignore_reinit_error 参数） + ray.init(ignore_reinit_error=True) + + # 定义超参数搜索空间 + config = { + "batch_size": tune.grid_search([args.batch_size]), + "learn_rate": tune.grid_search([args.lr]), + "max_length": tune.grid_search([args.max_length]), + "epoch": tune.grid_search([args.epoch]) + } + + # 设置一个 CLI 报告器，可以在命令行中看到进度 + reporter = CLIReporter( + metric_columns=["accuracy", "avg_loss", "training_iteration"] + ) + + # 调用 tune.run 开始超参数搜索 + analysis = tune.run( + tune.with_parameters(train_tune, args=args, tokenizer=tokenizer, dataset=dataset, optimizer=optimizer, model=model), + resources_per_trial={"gpu": 1}, # 每个试验分配 1 个 GPU；如果你的机器有多 GPU，就能实现不同试验分别在不同卡上运行 + config=config, + metric="accuracy", + mode="max", + progress_reporter=reporter, + storage_path=model_save_path, # 日志和检查点保存目录 + name="tune_experiment" + ) + + # 输出最佳超参数组合 + best_config = analysis.get_best_config(metric="accuracy", mode="max") + + with open("best_config", "w", encoding="utf-8") as f: + f.write(str(best_config)) + print("Best config: ", best_config) diff --git a/train/self_train.py b/train/self_train.py index 8a52290..1467272 100644 --- a/train/self_train.py +++ b/train/self_train.py @@ -1,3 +1,4 @@ +import os from typing import Dict import numpy as np @@ -6,8 +7,12 @@ import heapq from torch import device +from torch.utils.data import DataLoader from transformers import TrainingArguments, Trainer +import json +import jsonlines + from utils.dataset import mycollate_trainer, self_train_collate, AssertionExample, SelfTrainDataset import torch.nn.functional as F @@ -146,19 +151,19 @@ def get_soft_label_dataloader(self): for example in new_examples: example.weight /= sum_scores - def calc_weighted_loss(loss_fct, outputs, batch, scores): - inputs, labels = batch["input_ids"], batch["labels"] - loss_per_token = loss_fct(outputs.logits.view(-1, model.config.vocab_size), labels.view(-1)) - - # 由于每个样例的长度不同，需要计算每个样例的平均loss - loss_per_token = loss_per_token.view(inputs.size(0), -1) - loss_per_example = loss_per_token.sum(dim=1) / (labels != tokenizer.pad_token_id).sum(dim=1) - - whole_loss = sum([l * s for l, s in zip(loss_per_example, scores)]) / len(scores) - # 打印每个样例的loss - # for i, loss in enumerate(loss_per_example): - # ImportError}: {loss.item()}") - return whole_loss + # def calc_weighted_loss(loss_fct, outputs, batch, scores): + # inputs, labels = batch["input_ids"], batch["labels"] + # loss_per_token = loss_fct(outputs.logits.view(-1, model.config.vocab_size), labels.view(-1)) + # + # # 由于每个样例的长度不同，需要计算每个样例的平均loss + # loss_per_token = loss_per_token.view(inputs.size(0), -1) + # loss_per_example = loss_per_token.sum(dim=1) / (labels != tokenizer.pad_token_id).sum(dim=1) + # + # whole_loss = sum([l * s for l, s in zip(loss_per_example, scores)]) / len(scores) + # # 打印每个样例的loss + # # for i, loss in enumerate(loss_per_example): + # # ImportError}: {loss.item()}") + # return whole_loss def compute_loss(self, model, inputs, return_outputs=False): ## compute loss这个步骤实际上定义了 forward和loss的计算过程 score = torch.tensor(inputs.pop("weight")) @@ -174,6 +179,276 @@ def compute_loss(self, model, inputs, return_outputs=False): ## compute loss这 # return ( # loss, {'outputs': outputs}) if return_outputs else loss # 一定要记得 compute loss的时候！！！outputs要用字典返回 + +from tqdm import tqdm +import torch.nn.functional as F + + +def unlabel_filter(data, domain): + weather = ["[IN:UNSUPPORTED_WEATHER", "[IN:GET_WEATHER", "[IN:GET_SUNSET", "[IN:GET_SUNRISE", "[IN:GET_LOCATION", + "[SL:DATE_TIME", "[SL:WEATHER_TEMPERATURE_UNIT", "[SL:LOCATION", "[SL:WEATHER_ATTRIBUTE", + "[SL:LOCATION_USER", "[SL:SEARCH_RADIUS", "[SL:LOCATION_MODIFIER"] + reminder = ["IN:GET_MESSAGE]", "[SL:CONTACT", "[SL:MUTUAL_EMPLOYER", "[SL:METHOD_RETRIEVAL_REMINDER", "[SL:TODO", + "[SL:RECIPIENT", "[IN:DELETE_REMINDER", "[SL:CATEGORY_EVENT", "[IN:GET_REMINDER", "[IN:GET_TODO", + "SL:ATTENDEE_EVENT]", "[IN:GET_RECURRING_DATE_TIME", "SL:CONTENT_EXACT]", "IN:GET_RECURRING_DATE_TIME]", + "[SL:PERSON_REMINDED", "[IN:SEND_MESSAGE", "IN:REPLY_MESSAGE]", "[IN:GET_CONTACT", "[SL:AMOUNT", + "SL:CONTACT_RELATED]", "SL:FREQUENCY]", "SL:DATE_TIME]", "[SL:RECURRING_DATE_TIME", "[SL:CONTENT_EXACT", + "[IN:GET_MESSAGE", "[SL:DATE_TIME", "SL:ATTENDEE]", "IN:GET_CONTACT]", "SL:RECURRING_DATE_TIME]", + "[SL:FREQUENCY", "[IN:CREATE_REMINDER", "[SL:CONTACT_RELATED", "[IN:REPLY_MESSAGE", + "[SL:ATTENDEE_EVENT", "[SL:ORDINAL", "[SL:ATTENDEE", "SL:PERSON_REMINDED]", "SL:AMOUNT]", "SL:CONTACT]", + "SL:TYPE_RELATION]", "IN:GET_EVENT]", "[SL:TYPE_RELATION", "[IN:GET_EVENT"] + event = ["[SL:CATEGORY_LOCATION", "SL:CATEGORY_LOCATION]", "[SL:ORGANIZER_EVENT", "SL:ORGANIZER_EVENT]", + "[SL:CATEGORY_EVENT", "SL:CATEGORY_EVENT]", "[IN:GET_LOCATION", "IN:GET_LOCATION]", "[SL:ORDINAL", + "SL:ORDINAL]", "[SL:NAME_EVENT", "SL:NAME_EVENT]", "[SL:LOCATION", "SL:LOCATION]", "[SL:LOCATION_MODIFIER", + "SL:LOCATION_MODIFIER]", "[SL:DATE_TIME", "SL:DATE_TIME]", "[SL:ATTRIBUTE_EVENT", "SL:ATTRIBUTE_EVENT]", + "[SL:POINT_ON_MAP", "SL:POINT_ON_MAP]", "[IN:GET_EVENT", "IN:GET_EVENT]", ] + domain_label = {"weather": weather, "event": event, "reminder": reminder} + + label = data['pred'] + for word in label.split(): + if word != ']' and word not in data['question'] and word not in domain_label[domain]: + return False + + return True + +def get_unlabel_data(unlabel_train_loader, model,tokenizer, save_path: str, domain:str, p_rate: float): + + # 用于暂存所有样本及其置信度 + all_data = [] + + # 遍历无标签数据的 DataLoader + sum = 0 + true = 0 + for batch in tqdm(unlabel_train_loader): + batch = {k: v.to(device) for k, v in batch.items()} + with torch.no_grad(): + # 关键：需要打开 return_dict_in_generate 和 output_scores + outputs = model.generate( + batch['input_ids'], + max_length=128, + return_dict_in_generate=True, + output_scores=True, + ) + # 使用模型的 compute_transition_scores 方法计算转移分数 + transition_scores = model.copmute_transition_scores( + outputs.sequences, outputs.scores, beam_indices=None, normalize_logits=False + ) + # 在第 0 个维度上归一化 + # transition_scores = F.softmax(transition_scores, dim=0) + confidence_scores = transition_scores.sum(dim=1) + # 计算每条生成序列的置信度 + # outputs.scores 是一个列表，长度等于生成的 token 步数 + # 每个元素 shape = [batch_size, vocab_size] + batch_size = batch['input_ids'].shape[0] + for b_idx in range(batch_size): + # 解码输入序列和生成的序列（可根据需要 skip_special_tokens 等参数） + input_text = tokenizer.decode(batch["input_ids"][b_idx], skip_special_tokens=True) + pred_text = tokenizer.decode(outputs.sequences[b_idx], skip_special_tokens=True) + conf = confidence_scores[b_idx].item() + + data_dict = { + "question": input_text, + "pred": pred_text, + "confidence": conf + } + if unlabel_filter(data_dict,domain): + all_data.append(data_dict) + + # 排序并取置信度处于前 30% 的数据 + all_data = sorted(all_data, key=lambda x: x["confidence"], reverse=True) + top_n = int(len(all_data) * p_rate) + top_data = all_data[:top_n] + + true0 = 0 + for item in top_data: + if item['pred'] == item['label']: + true0 += 1 + if true0 / len(top_data) > 0.85: + p_rate += 0.1 + if p_rate >= 1: + p_rate = 1 + + p_lst = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] + for p1 in p_lst: + top_n1 = int(len(all_data) * p1) + test_data = all_data[:top_n1] + # 看每個區間下面的分數 + true1 = 0 + for item in test_data: + if item['pred'] == item['label']: + true1 += 1 + + + # 写入文件 + with open(save_path, 'w', encoding='utf-8') as f: + for item in top_data: + json_line = json.dumps(item, ensure_ascii=False) + f.write(json_line + '\n') + + return p_rate + + +def get_selftrain_model(tokenizer, unlabel_path: str): + train_data = [] + with jsonlines.open(unlabel_path) as reader: + for line in reader: + source_text = line['question'] + target_text = line['pred'] + # 对输入文本进行编码 + source_encoding = tokenizer(source_text, truncation=True, padding='max_length', max_length=128, return_tensors="pt") + input_ids = source_encoding["input_ids"] + attention_mask = source_encoding["attention_mask"] + + # 对目标文本进行编码 + target_encoding = tokenizer(target_text, truncation=True, padding='max_length', max_length=128, return_tensors="pt") + labels = target_encoding["input_ids"] + + train_data.append({ + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels + }) + + return train_data + + +# def train_model_self_train(model, tokenizer, optimizer, dataset, args): +# """ +# 1. 先训练基础模型 +# 2. 用基础模型预测unlabeled数据集 +# 3. 选择topk数据集 +# 4. 用topk数据集fine-tune基础模型 +# 5. 重复2-4 +# """ +# # print("###", len(dataset["train"])) +# # for d in dataset["train"]: +# # print(d) +# +# +# +# +# +# # selftrain_args = SelfTrainingArguments(output_dir=args.save_dir, +# # num_train_epochs=args.selftrain_iteration, # 这个指每个self_train里面的epoch +# # per_device_train_batch_size=args.batch_size, +# # save_steps=1000, +# # learning_rate=args.lr, +# # evaluation_strategy="epoch", +# # selftrain_topk=2, +# # no_cuda=False if args.device == "cuda" else True) +# +# # model.resize_token_embeddings(len(tokenizer)) +# +# for epoch in range(25): +# if args.given_model: # 如果基础模型已经训过了，就先训self +# args.given_model = False +# else: +# # 1. 先训练基础模型 +# if epoch != 0: +# # train_args = TrainingArguments(output_dir=args.save_dir, +# # num_train_epochs=25, # 300个epoch的ft + 25*（10个selftrain + 25个ft） +# # per_device_train_batch_size=args.batch_size, +# # save_steps=1000, +# # learning_rate=args.lr, +# # evaluation_strategy="epoch", +# # do_eval=True if "eval" in dataset else False, +# # no_cuda=False if args.device == "cuda" else True, +# # fp16=True, +# # ) +# train_args = TrainingArguments( +# output_dir=args.save_dir, +# num_train_epochs=args.epoch, +# per_device_train_batch_size=args.batch_size, +# save_steps=1000, +# save_total_limit=1, +# learning_rate=args.lr, +# evaluation_strategy="epoch", +# do_eval=True if "eval" in dataset else False, +# no_cuda=False if args.device == "cuda" else True, +# dataloader_num_workers=1, # 根据gpu数量调整 +# fp16=True, # 如果你有显存足够，可以开启fp16加速训练 +# local_rank=int(os.environ.get("LOCAL_RANK", -1)), # 设置用于分布式训练 +# ) +# else: +# # train_args = TrainingArguments(output_dir=args.save_dir, +# # num_train_epochs=args.epoch, +# # per_device_train_batch_size=args.batch_size, +# # save_total_limit=1, +# # save_steps=1000, +# # learning_rate=args.lr, +# # evaluation_strategy="epoch", +# # do_eval=True if "eval" in dataset else False, +# # no_cuda=False if args.device == "cuda" else True, +# # fp16=True,) +# train_args = TrainingArguments( +# output_dir=args.save_dir, +# num_train_epochs=args.epoch, +# per_device_train_batch_size=args.batch_size, +# save_steps=1000, +# save_total_limit=1, +# learning_rate=args.lr, +# evaluation_strategy="epoch", +# do_eval=True if "eval" in dataset else False, +# no_cuda=False if args.device == "cuda" else True, +# dataloader_num_workers=1, # 根据gpu数量调整 +# fp16=True, # 如果你有显存足够，可以开启fp16加速训练 +# local_rank=int(os.environ.get("LOCAL_RANK", -1)), # 设置用于分布式训练 +# ) +# +# trainer = Trainer(model=model, +# args=train_args, +# data_collator=mycollate_trainer, # 要么给自己的，要么在定义trainer后面单独写一个data_collator=None，不然代码里有默认collate +# train_dataset=dataset["train"], +# eval_dataset=dataset["validation"] if "validation" in dataset else None, +# tokenizer=tokenizer, +# optimizers=(optimizer, None)) # 缺了学习率调度器 +# +# trainer.train() +# model = trainer.model +# model.to(args.device) +# +# print("初步训练完成") +# +# if args.close_selftrain: +# return model +# +# unlabeled_dataset = dataset["unlabeled"] # 我们假设这里是个question_list +# # unlabeled_dataset = SelfTrainDataset(question_list=unlabeled_dataset) +# # 自训练 +# unlabel_train_loader = DataLoader(unlabeled_dataset, batch_size=128, collate_fn=mycollate_trainer) +# p_rate = 0.2 +# unlabel_path = "" +# # 得到数据 +# p_rate = get_unlabel_data(unlabel_train_loader, model,tokenizer, unlabel_path,"weather",p_rate) +# # 自训练 +# unlabel_train_dataset = get_selftrain_model(tokenizer, unlabel_path) +# selftrain_args = TrainingArguments( +# output_dir=f"{args.save_dir}/normal", +# num_train_epochs=args.selftrain_iteration, +# per_device_train_batch_size=128, +# learning_rate=args.sf_lr, +# do_eval=False, +# no_cuda=False +# ) +# +# # 定义Trainer实例 +# selftrainer = Trainer( +# model=model, +# args=selftrain_args, +# data_collator=mycollate_trainer, +# train_dataset=unlabel_train_dataset +# ) +# # 开始自训练 +# selftrainer.train() +# +# model = selftrainer.model +# model.save_pretrained(f"/data/lbq/models/mt5_1000_{epoch}") +# +# return model + + +tokenizer1 = None + def train_model_self_train(model, tokenizer, optimizer, dataset, args): """ 1. 先训练基础模型 @@ -182,62 +457,35 @@ def train_model_self_train(model, tokenizer, optimizer, dataset, args): 4. 用topk数据集fine-tune基础模型 5. 重复2-4 """ - train_args = TrainingArguments(output_dir=args.save_dir, - num_train_epochs=1,#args.epoch, # 这个指每个self_train里面的epoch - per_device_train_batch_size=args.batch_size, - save_steps=1000, - learning_rate=args.lr, - evaluation_strategy="epoch", - do_eval=True if "eval" in dataset else False, - no_cuda=False if args.device == "cuda" else True) - - unlabeled_dataset = dataset["unlabeled"] # 我们假设这里是个question_list - # unlabeled_dataset = SelfTrainDataset(question_list=unlabeled_dataset) - - selftrain_args = SelfTrainingArguments(output_dir=args.save_dir, - num_train_epochs=args.selftrain_iteration, # 这个指每个self_train里面的epoch - per_device_train_batch_size=args.batch_size, - save_steps=1000, - learning_rate=args.lr, - evaluation_strategy="epoch", - selftrain_topk=2, - no_cuda=False if args.device == "cuda" else True) - - # model.resize_token_embeddings(len(tokenizer)) - - for epoch in range(3): - if args.given_model: # 如果基础模型已经训过了，就先训self - args.given_model = False - else: - # 1. 先训练基础模型 - trainer = Trainer(model=model, - args=train_args, - data_collator=mycollate_trainer, # 要么给自己的，要么在定义trainer后面单独写一个data_collator=None，不然代码里有默认collate - train_dataset=dataset["train"][:10], - eval_dataset=dataset["train"][:10],#dataset["eval"] if "eval" in dataset else None, - tokenizer=tokenizer, - optimizers=(optimizer, None)) # 缺了学习率调度器 - - trainer.train() - model = trainer.model - model.to(args.device) - - print("初步训练完成") - - # 2. 用基础模型预测unlabeled数据集 - self_trainer = SelfTrainTrainer(model=model, - train_args=selftrain_args, - data_collator=self_train_collate, - train_dataset=unlabeled_dataset, - tokenizer=tokenizer, - optimizers=(optimizer, None)) - - unlabeled_dataset.eval() - self_trainer.get_soft_label_dataloader() - - unlabeled_dataset.train(tokenizer, args.max_length) - self_trainer.train() - - model = self_trainer.model - - model.save_pretrained(f"/data/lbq/models/mt5_1000_{epoch}") \ No newline at end of file + from transformers import AutoTokenizer, AdamW, get_scheduler + global tokenizer1 + tokenizer1 = tokenizer + + train_args = TrainingArguments( + output_dir=args.save_dir, + num_train_epochs=args.epoch, + per_device_train_batch_size=args.batch_size, + save_steps=1000, + save_total_limit=1, + learning_rate=args.lr, + evaluation_strategy="epoch", + do_eval=True if "eval" in dataset else False, + no_cuda=False if args.device == "cuda" else True, + dataloader_num_workers=1, # 根据gpu数量调整 + fp16=True, # 如果你有显存足够，可以开启fp16加速训练 + local_rank=int(os.environ.get("LOCAL_RANK", -1)), # 设置用于分布式训练 + ) + + trainer = Trainer(model=model, + args=train_args, + data_collator=mycollate_trainer, # 要么给自己的，要么在定义trainer后面单独写一个data_collator=None，不然代码里有默认collate + train_dataset=dataset["train"], + eval_dataset=dataset["validation"], # dataset["eval"] if "eval" in dataset else None, + tokenizer=tokenizer, + optimizers=(optimizer, None)) # 缺了学习率调度器 + + trainer.train() + model = trainer.model + model.to(args.device) + + return model diff --git a/utils/data_preprocess.py b/utils/data_preprocess.py index b8a2890..e3768b1 100644 --- a/utils/data_preprocess.py +++ b/utils/data_preprocess.py @@ -11,8 +11,11 @@ from utils.dataset import AssertionExample, PreliminaryDataset, SelfTrainDataset from utils.operators_concepts import operator_dict from utils.text_utils import add_space_after_chinese, find_long_string_in_list +from utils.remove_non_slot_leaf import remove_non_slot_leaf_nodes +from utils.sort_label import sort_string import pandas as pd +import json @DatasetsProcessorNameSpace.register("Default") @@ -42,6 +45,7 @@ def split_dataset(dataset: Union[list, Dataset, DatasetDict], split_ratio: Union def read_ours_from_dir(directory_path: str) -> List[AssertionExample]: dataset = [] + s = 0 # 遍历目录下的所有文件 for filename in os.listdir(directory_path): file_path = os.path.join(directory_path, filename) @@ -50,12 +54,13 @@ def read_ours_from_dir(directory_path: str) -> List[AssertionExample]: line = eval(f.readlines()[0]) for e in line: - expression, natural_sentence = e['表达式'], random.choice(e['自然语句']) - if expression == "None": + # 存在e['自然语句'] = []的情况导致random.choice(e['自然语句'])报错 + if(e['自然语句'] == []): continue - + expression, natural_sentence = e['表达式'], random.choice(e['自然语句']) + # if expression == "None": + # continue dataset.append(AssertionExample(expression, natural_sentence)) - return dataset @DatasetsReaderNameSpace.register("preliminary_ours") @@ -69,26 +74,36 @@ def read_dataset(directory_path: str) -> Union[Dataset, DatasetDict]: return PreliminaryDataset(dataset) +from torch.utils.data import random_split @DatasetsReaderNameSpace.register("self-train_ours") def read_dataset(directory_path: str) -> Union[Dict, DatasetDict]: """ 一般情况下，read的时候就改成train eval test """ - dataset = read_ours_from_dir(directory_path) + train_dataset = read_ours_from_dir(directory_path) + + # 计算训练集和开发集的大小 + # train_size = int(0.95 * len(train_dataset)) + # dev_size = len(train_dataset) - train_size - return {'train': PreliminaryDataset(dataset)} + # # 使用random_split划分数据集 + # train_subset, dev_subset = random_split(train_dataset, [train_size, dev_size]) + return {'train': PreliminaryDataset(train_dataset)} @DatasetsReaderNameSpace.register("self-train_topv2") def read_dataset(directory_path: str) -> DatasetDict: """ 一般情况下，read的时候就改成train eval test """ + # data_files = {'train': '/home/cyz/data/semantic_parsing/code/TOPv2/low_resource_splits/our/data/our_weather_train.tsv', + # 'test': '/home/cyz/data/semantic_parsing/code/TOPv2/low_resource_splits/our/data/weather_valid_50spis.tsv'} dataset = load_dataset(directory_path) + train_dataset = dataset["train"] return dataset @DatasetsReaderNameSpace.register("self-train_zcl") def read_dataset(directory_path: str) -> DatasetDict: - dataset1000 = load_dataset("/data/lbq/datasets/1000spis", split='train') + dataset1000 = load_dataset(directory_path, split='train') # dataset1000 = load_dataset("/home/lzx/old_selftrain/data/weather/1000spis", split='train') dataset = dataset1000.train_test_split(test_size=0.5) return {"train": dataset["train"]} @@ -96,15 +111,15 @@ def read_dataset(directory_path: str) -> DatasetDict: @DatasetsReaderNameSpace.register("self-train_zcl_mixed") def read_dataset(directory_path: str) -> DatasetDict: def filter_function(example): - return example["seqlogical"].find("[IN:UNSUPPORTED_WEATHER") != -1 + return example["semantic_parse"].find("[IN:UNSUPPORTED_WEATHER") != -1 - dataset = load_dataset("/data/lbq/datasets/1000spis", split="train") + dataset = load_dataset(directory_path, split="train") dataset1 = dataset.filter(filter_function) dataset1 = dataset1.train_test_split(test_size=0.5) # 再度入一个TopV2的，取0.2 def filter_function(example): - return example["seqlogical"].find("[IN:UNSUPPORTED_WEATHER") == -1 + return example["semantic_parse"].find("[IN:UNSUPPORTED_WEATHER") == -1 dataset2 = dataset.filter(filter_function) dataset2 = dataset2.train_test_split(test_size=0.8) @@ -113,7 +128,7 @@ def filter_function(example): with open("/data/lbq/datasets/train_weather.json", 'r') as file: dataset3 = json.load(file) data_dict = { - "seqlogical": [item["expression"] for item in dataset3], + "semantic_parse": [item["expression"] for item in dataset3], "utterance": [item["sentence"] for item in dataset3], } @@ -126,13 +141,28 @@ def filter_function(example): @DatasetsReaderNameSpace.register("ours") def read_unlabeled_dataset(directory_path: str): - dataset = load_dataset(directory_path) - return SelfTrainDataset(init_question_list=[l[0] for l in dataset["train"]["自然语句"]]) + # 修改为 ↓ + for filename in os.listdir(directory_path): + # 读取json文件，里面有"split_and_filter"和"origin" + # 把split_and_filter里的全部取出来就行 + with open(directory_path + "/" + filename) as f: # replace 'yourfilename' with your actual file name + data = json.load(f) + flattened_list = [] + for i in range(len(data)): + flattened_list += data[i]["split_and_filter"] + return SelfTrainDataset(init_question_list=flattened_list) @DatasetsReaderNameSpace.register("topv2") def read_unlabeled_dataset(directory_path: str): - dataset = load_dataset(directory_path) - return dataset["eval"]["utterance"] + # dataset = load_dataset(directory_path) + # 用于存储所有字典数据的列表 + query_list = [] + with open(directory_path, "r", encoding="utf-8") as f: + for s in f: + dict1 = json.loads(json.loads(s)) + query_list.append(dict1["question"]) + return SelfTrainDataset(init_question_list=query_list) + # return dataset["eval"]["utterance"] # zcl def read_unlabeled_dataset_zcl(directory_path: str): @@ -196,6 +226,7 @@ def unify_format(example: AssertionExample): @DatasetsProcessorNameSpace.register("ours") def ptr_change(example: AssertionExample): e = unify_format(example) + e.natural_sentence = add_space_after_chinese(e.natural_sentence.replace("得到", "")) # encode = tokenizer.encode(e.natural_sentence) word_list = e.natural_sentence.split() @@ -228,6 +259,10 @@ def ptr_change(example: AssertionExample): rhs_indexes = find_long_string_in_list(word_list, rhs) + # 改成留着第一个和最后一个，看看能不能好点 + if len(rhs_indexes) > 0: + rhs_indexes = [rhs_indexes[0],rhs_indexes[-1]] + rhs_ptr = f"[{operator_dict[predicate][-1]}：" + "".join( ["@ptr_" + str(item + 1) for item in rhs_indexes]) + "]" variable_list = [] @@ -236,6 +271,8 @@ def ptr_change(example: AssertionExample): new_structural_tokens.append(f"[{operator_dict[predicate][-1]}：") for concept, variable in zip(operator_dict[predicate], variables): variable_indexes = find_long_string_in_list(word_list, variable) + if len(variable_indexes) > 0: + variable_indexes = [variable_indexes[0],variable_indexes[-1]] result = "".join(["@ptr_" + str(item + 1) for item in variable_indexes]) variable_list.append(f"[{concept}：{result}]") new_structural_tokens.append(f"[{concept}：") @@ -253,30 +290,115 @@ def ptr_change(example: AssertionExample): if i < len(result_list) - 1: Result += " , " e.expression = Result + # print(e) return e +import re +import string +# 英文标点符号 +def format_time_string(input_string): + english_punctuation = string.punctuation.replace(':', '') + # 正则表达式匹配时间格式 + time_pattern = r'(\d{1,2})(:)?(\d{2})?(am|pm|AM|PM)?' + + # 替换逻辑 + def replacer(match): + hour = match.group(1) + colon = " : " if match.group(2) else "" + minute = match.group(3) if match.group(3) else "" + period = f" {match.group(4)}" if match.group(4) else "" + return f"{hour}{colon}{minute}{period}" + + # 对字符串进行替换 + formatted_string = re.sub(time_pattern, replacer, input_string) + + # 判断末尾字符是否是english_punctuation内的标点符号且前面无空格 + for i in range(len(formatted_string)): + if formatted_string[i] in english_punctuation and formatted_string[i-1]!= " ": + formatted_string = formatted_string[:i] + " " + formatted_string[i:] + + return re.sub(r"([a-zA-Z])'s", r"\1 's", formatted_string) + +def in_input(content, input_str): + if f" {content} " not in input_str: + if f" {content}?" not in input_str and f" {content}." not in input_str and f" {content}," not in input_str: + if f" {content}" not in input_str: + return False + elif input_str.endswith(f" {content}") == False: + return False + + return True + +def edit_label(examples): + input = examples["utterance"] + output = examples["semantic_parse"] + + pattern = r"(?<=\[SL:WEATHER_TEMPERATURE_UNIT\s)(.*?)(?=\s\])" + if "[SL:WEATHER_TEMPERATURE_UNIT" in output: + match = re.search(pattern, output) + content = match.group(0) + if in_input(content, input) == False: + in_input1 = content.capitalize() + if in_input(in_input1, input) == True: + output = re.sub(pattern, in_input1, output) + elif in_input(content.lower(), input) == True: + output = re.sub(pattern, content.lower(), output) + else: + return examples + + examples["semantic_parse"] = output + + return examples + +def filter(examples): + input = examples["utterance"] + output = examples["semantic_parse"] + + # 过滤掉today误生成的 + pattern_today = r"\[SL:DATE_TIME[^\]]*today[^\]]*\]" + pattern_Today = r"\[SL:DATE_TIME[^\]]*Today[^\]]*\]" + # 过滤掉today误生成的 + if input.find("today") != -1 and re.search(pattern_today, output) == None: + return False + if input.find("Today") != -1 and re.search(pattern_Today, output) == None: + return False + + # 判断[SL:WEATHER_TEMPERATURE_UNIT的内容在label里有没有 + pattern = r"(?<=\[SL:WEATHER_TEMPERATURE_UNIT\s)(.*?)(?=\s\])" + if "[SL:WEATHER_TEMPERATURE_UNIT" in output: + match = re.search(pattern, output) + content = match.group(0) + return in_input(content, input) + + return True + @DatasetsProcessorNameSpace.register("topv2") def ptr_change(examples): """ 将semantic_parse里面的的词，换成utterance里对应的ptr_x """ - for i, st in enumerate(examples["semantic_parse"]): - changed_item = [] - # ut_list = ut.split(' ') - cnt = 1 - # print(st) - for s in st.split(' '): - if s.startswith('[') or s == ']': - # print(s) - # exit() - changed_item.append(s) - else: - # print(s) - # print(f"@ptr_{cnt}") - changed_item.append(f"@ptr_{cnt}") - cnt += 1 - - examples["semantic_parse"][i] = ' '.join(changed_item) + # 改词 + examples["semantic_parse"] = edit_label(examples)["semantic_parse"] + # print(examples) + # st = examples["semantic_parse"] + # changed_item = [] + # cnt = 1 + + # for s in st.split(' '): + # # 如果是以 [ 开头 或 是 ] 或 是英文标点符号，则保留原样 + # if s.startswith('[') or s == ']': + # changed_item.append(s) + # else: + # changed_item.append(f"@ptr_{cnt}") + # cnt += 1 + + # examples["semantic_parse"] = ' '.join(changed_item) + examples["utterance"] = format_time_string(examples["utterance"]) + # 删标签 + examples["semantic_parse"] = remove_non_slot_leaf_nodes(examples["semantic_parse"]) + # 排序 + examples["semantic_parse"] = sort_string(examples["semantic_parse"]) + # print(examples) return examples #zcl @@ -284,7 +406,7 @@ def ptr_change_zcl(examples): """ 将semantic_parse里面的的词，换成utterance里对应的ptr_x """ - st = examples["seqlogical"] + st = examples["semantic_parse"] changed_item = [] # ut_list = ut.split(' ') cnt = 1 @@ -300,7 +422,7 @@ def ptr_change_zcl(examples): changed_item.append(f"@ptr_{cnt}") cnt += 1 - examples["seqlogical"] = ' '.join(changed_item) + examples["semantic_parse"] = ' '.join(changed_item) return examples @DatasetsProcessorNameSpace.register("zcl") def ptr_change(examples): @@ -322,7 +444,36 @@ def filter_function(example): return example def preprocess_dataset(dataset): - dataset = dataset.map(ptr_change) - dataset = dataset.filter(filter_function) + dataset = dataset.filter(filter) + dataset = dataset.map(ptr_change, load_from_cache_file=False) + + return dataset - return dataset \ No newline at end of file +# if __name__ == '__main__': +# s = "[IN:GET_ALARM Check my alarms . ]" +# import string +# english_punctuation = string.punctuation.replace(':', '') +# +# def ptr_change(examples): +# """ +# 将semantic_parse里面的的词，换成utterance里对应的ptr_x +# """ +# # print(examples) +# st = examples["semantic_parse"] +# changed_item = [] +# cnt = 1 +# +# for s in st.split(' '): +# # 如果是以 [ 开头 或 是 ] 或 是英文标点符号，则保留原样 +# if s.startswith('[') or s == ']' or all(c in english_punctuation for c in s): +# changed_item.append(s) +# else: +# changed_item.append(f"@ptr_{cnt}") +# cnt += 1 +# +# examples["semantic_parse"] = ' '.join(changed_item) +# examples["utterance"] = format_time_string(examples["utterance"]) +# # print(examples) +# return examples +# +# print(ptr_change({"semantic_parse": s, "utterance": s})) diff --git a/utils/dataset.py b/utils/dataset.py index 49d0c00..286cb15 100644 --- a/utils/dataset.py +++ b/utils/dataset.py @@ -7,13 +7,14 @@ class AssertionExample: - def __init__(self, expression, natural_sentence, weight=-1): + def __init__(self, expression,natural_sentence, score=None,weight=-1): self.expression = expression self.natural_sentence = natural_sentence self.weight = weight # 只是self train时候才用 - - self.__dict__.update({"expression": expression, - "natural_sentence": natural_sentence, + self.score = score + self.__dict__.update({"natural_sentence": natural_sentence, + "expression": expression, + "score" : score, "weight": weight}) def __repr__(self): @@ -59,7 +60,7 @@ def __len__(self): return len(self.examples) def append(self, expression, natural_sentence): - self.examples.append(AssertionExample(expression, natural_sentence)) + self.examples.append(AssertionExample(expression=expression, natural_sentence=natural_sentence)) def map(self, func, *args, **kwargs): return PreliminaryDataset([func(e, *args, **kwargs) for e in self]) @@ -85,21 +86,29 @@ def __init__(self, init_question_list: str, topk=5):#unlabeled_dataset: List[Lis 综上所述，暂定还是list """ super().__init__() - self.key_to_index = {q: i for i, q in enumerate(init_question_list)} + + t = 0 + self.key_to_index = {} + for i, q in enumerate(init_question_list): + if q not in self.key_to_index: + self.key_to_index[q] = i - t + else: + t += 1 + # self.key_to_index = {q: i for i, q in enumerate(init_question_list)} + # for i, data_list in enumerate(self.unlabeled_dataset): # 这里list没关系，因为每次都是更新所有的score，所以每次整个把data_list删掉重建 # # 因为即便用dict，修改方便但每次还要排序 # key = data_list[0].expression # self.key_to_index[key] = i - self.unlabeled_dataset = [[] for _ in range(len(init_question_list))] + self.unlabeled_dataset = [[] for _ in range(len(self.key_to_index))] self.sent_to_instance_list = [] # 用于避免重复 for i, data_list in enumerate(self.unlabeled_dataset): # sent_to_instance = {data.natural_sentence: data for data in data_list}# 这个地方就不应该有重复 sent_to_instance = {} self.sent_to_instance_list.append(sent_to_instance) - self.sorted_sign = [False] * len(init_question_list) # 用于减少排序开销 - + self.sorted_sign = [False] * len(self.key_to_index) # 用于减少排序开销 self.tokenized_dataset = None self.tokenized_sign = False @@ -149,21 +158,21 @@ def __contains__(self, item): return False - def append(self, natural_sentence, expression, score=-1): + def append(self, natural_sentence, expression, score,weight=-1): """ 判断新旧再append，这里面不控制 """ key = natural_sentence if key in self.key_to_index: - self.unlabeled_dataset[self.key_to_index[key]].append(AssertionExample(expression, natural_sentence, score)) #? + self.unlabeled_dataset[self.key_to_index[key]].append(AssertionExample(natural_sentence=natural_sentence,expression=expression, score=score,weight=weight)) #? self.sorted_sign[self.key_to_index[key]] = False else: self.key_to_index[key] = len(self.unlabeled_dataset) - self.unlabeled_dataset.append([AssertionExample(expression, natural_sentence, score)]) + self.unlabeled_dataset.append([AssertionExample(natural_sentence=natural_sentence,expression=expression, score=score,weight=weight)]) self.sorted_sign.append(False) - def train(self, tokenizer, max_length=512): + def train(self, tokenizer, max_length=256): """进入train的状态，此时应该改为返回tokenize_dataset""" self.tokenized_dataset = self._return_tokenized_dataset(tokenizer, max_length) # 每次重新算吧，毕竟labels在更新，也浪费不了多少时间 self.tokenized_sign = True # 顺序不能颠倒，不然getitem会错 @@ -173,7 +182,7 @@ def eval(self): del self.tokenized_dataset self.tokenized_sign = False - def _return_tokenized_dataset(self, tokenizer, max_length=512) -> List[List[Dict[str, torch.Tensor]]]: + def _return_tokenized_dataset(self, tokenizer, max_length=256) -> List[List[Dict[str, torch.Tensor]]]: """ 这个地方直接返回可以进dataloader的dataset/list, 不过这里因为有topk，所以要多一层 """ @@ -185,19 +194,29 @@ def tokenize_example(input_text): input_ids = torch.tensor(input_text + [tokenizer.pad_token_id] * (max_length - len(input_text))) input_ids = input_ids.unsqueeze(0) # Add batch dimension elif isinstance(input_text, torch.Tensor): - if input_text.size(0) < max_length: + if input_text.size(0) <= max_length: padding = torch.tensor([tokenizer.pad_token_id] * (max_length - input_text.size(1))).unsqueeze(dim=0) input_ids = torch.cat((input_text.cpu(), padding), dim=1) else: raise ValueError("Invalid input_text type {}.".format(type(input_text))) - - return input_ids.cpu() # 这个地方固定住to.cpu也没关系，因为目测没有to(cuda)的需求 - + return input_ids.to(torch.int64).cpu() # 这个地方固定住to.cpu也没关系，因为目测没有to(cuda)的需求 + + #相似化的处理 + def tokenize_example1(input_text): + from .text_utils import add_space_after_chinese + input_text = add_space_after_chinese(input_text.replace("得到", "")) + from .tokenization import delete_blank + tokenized_inputs = tokenizer(input_text, padding='max_length', truncation=True, max_length=max_length, return_tensors="pt") + tokenized_inputs = delete_blank(tokenized_inputs)["input_ids"] + return tokenized_inputs.to(torch.int64).cpu() # 这个地方固定住to.cpu也没关系，因为目测没有to(cuda)的需求 + + # 这里是不是要和正常的对齐,即加空格，再删空格 tokenized_dataset = [] for topk_examples in self: tokenized_dataset.append( - [{"input_ids": tokenize_example(example.natural_sentence), + [{"input_ids": tokenize_example1(example.natural_sentence), "labels": tokenize_example(example.expression), + "score" : example.score, "weight": example.weight} for example in topk_examples]) # 因为getitem时候已经取了topk @@ -246,12 +265,11 @@ def self_train_collate(examples): """ 和mycollate一样，主要是每四个要单独处理一个y^ """ + for topk_examples in examples: weights_sum = sum([example['weight'] for example in topk_examples]) - print("examples", examples) for topk_examples in examples: - print(topk_examples) for sub_example in topk_examples: # 指topk for key in sub_example: if key == "weight": @@ -263,7 +281,6 @@ def self_train_collate(examples): pass examples = [e for topk_examples in examples for e in topk_examples] - batch = {} for key in examples[0]: try: diff --git a/utils/sort_label.py b/utils/sort_label.py new file mode 100644 index 0000000..80f2e82 --- /dev/null +++ b/utils/sort_label.py @@ -0,0 +1,149 @@ +def sort_string(s): + """ + 输入一个类似下面格式的字符串： + [IN:GET_WEATHER [SL:WEATHER_ATTRIBUTE warm ] [SL:DATE_TIME tomorrow morning ] [SL:WEATHER_TEMPERATURE_UNIT fahrenheit ] ] + 输出排序后的结果，排序规则为：在每个 [IN:…] 块内，将直接子项中所有 [SL:XXX ...] 按照 XXX（即冒号后第一个单词）的字母序排序， + 如果多个 [SL:XXX ...] 的 XXX 部分相同，则保持原有顺序。 + 嵌套的 [IN:…] 块独立处理。 + """ + token, _ = parse_token(s, 0) + sort_token(token) + return token_to_string(token) + + +def parse_token(s, i): + """ + 递归解析，从 s[i] 开始解析一个 token（应以 '[' 开头），返回 (token, new_index)。 + token 以字典表示，结构为： + { + 'type': "IN" 或 "SL", + 'command': 第一个单词（如 "GET_WEATHER" 或 "WEATHER_ATTRIBUTE"）， + 'args': 剩余的纯文本参数（可能为空）， + 'children': 直接嵌套的子 token 列表 + } + """ + assert s[i] == '[', f"预期 '[' 开头，当前字符：{s[i]}" + i += 1 # 跳过 '[' + + # 读取 token 类型，直到遇到冒号 + token_type = "" + while i < len(s) and s[i] != ':': + token_type += s[i] + i += 1 + token_type = token_type.strip() + i += 1 # 跳过冒号 + + # 读取 command（第一个单词），直到遇到空白、'[' 或 ']' + command = "" + while i < len(s) and s[i] not in [' ', '\t', '\n', '[', ']']: + command += s[i] + i += 1 + command = command.strip() + + # 跳过 command 后的空白 + while i < len(s) and s[i].isspace(): + i += 1 + + args = "" + children = [] + + # 读取 token 内部内容，直到遇到对应的 ']' + while i < len(s) and s[i] != ']': + if s[i] == '[': + # 如果遇到子 token，则递归解析 + child, i = parse_token(s, i) + children.append(child) + # 解析完后跳过可能的空白 + while i < len(s) and s[i].isspace(): + i += 1 + else: + # 否则读取纯文本（参数部分），直到遇到 '[' 或 ']' + start = i + while i < len(s) and s[i] not in ['[', ']']: + i += 1 + args += s[start:i] + # 去掉首尾多余空格，但注意保留中间的空格 + # （这里可以根据需要调整，默认用 strip()） + args = args.strip() # 可选，根据需要是否保留原始空白 + # 如果后面还有空格则跳过 + while i < len(s) and s[i].isspace(): + i += 1 + + # 跳过 ']' + i += 1 + + token = { + 'type': token_type, + 'command': command, + 'args': args, + 'children': children + } + return token, i + + +def sort_token(token): + """ + 对 token 进行排序： + 如果 token 是 [IN:...] 类型，则将其直接 children 中所有类型为 SL 的 token 按其 command 字母序排序（保持稳定性）。 + 同时对子 token 递归调用 sort_token。 + """ + # 先对子 token 递归排序 + for child in token['children']: + sort_token(child) + + if token['type'] == 'IN': + # 只对直接 children 中类型为 SL 的 token 排序， + # 若 children 中存在其他类型的 token，则保持它们在原位置不变。 + # 记录 SL token 的原始索引及 token 对象 + sl_indices = [] + sl_tokens = [] + for idx, child in enumerate(token['children']): + if child['type'] == 'SL': + sl_indices.append(idx) + sl_tokens.append(child) + # 对 sl_tokens 按 command 排序（稳定排序） + sl_tokens_sorted = sorted(sl_tokens, key=lambda x: x['command']) + # 将排序后的 SL token 放回原来的索引位置 + for pos, sorted_token_obj in zip(sl_indices, sl_tokens_sorted): + token['children'][pos] = sorted_token_obj + + +def token_to_string(token): + """ + 将 token 按照格式重构为字符串。 + 输出格式示例： + [IN:GET_WEATHER [SL:DATE_TIME tomorrow morning ] [SL:WEATHER_ATTRIBUTE warm ] [SL:WEATHER_TEMPERATURE_UNIT fahrenheit ] ] + """ + parts = [token['command']] + if token['args']: + parts.append(token['args']) + # 如果有 children，则每个子 token前加一个空格分隔 + for child in token['children']: + parts.append(token_to_string(child)) + inner = " ".join(parts) + return f"[{token['type']}:{inner} ]" + + +# --------------------------- +# 以下为简单测试代码 +if __name__ == "__main__": + # 示例1 + # input_str = "[IN:GET_WEATHER [SL:WEATHER_ATTRIBUTE warm ] [SL:DATE_TIME tomorrow morning ] [SL:WEATHER_TEMPERATURE_UNIT fahrenheit ] ]" + # output_str = sort_string(input_str) + # print("示例1输出：") + # print(output_str) + # # 预期输出： + # # [IN:GET_WEATHER [SL:DATE_TIME tomorrow morning ] [SL:WEATHER_ATTRIBUTE warm ] [SL:WEATHER_TEMPERATURE_UNIT fahrenheit ] ] + # + # # 示例2（包含嵌套） + # input_str2 = "[IN:GET_WEATHER [SL:LOCATION [IN:GET_LOCATION [SL:LOCATION_USER here ] ] ] [SL:DATE_TIME for next week ] ]" + # output_str2 = sort_string(input_str2) + # print("\n示例2输出：") + # print(output_str2) + # # 对于嵌套部分：[IN:GET_LOCATION ...] 内部只有一个 token，不受影响；外层仅排序直接的 [SL:DATE_TIME ...] 与 [SL:LOCATION ...] + file1 = open("preprocess2.txt","w",encoding="utf-8") + with open("/home/lzx2000/T5-base-lora/TOPv2/low_resource_splits/preprocess.txt", "r", encoding="utf-8") as f: + for line in f.readlines(): + file1.write(sort_string(line)) + file1.write("\n") + diff --git a/utils/tokenization.py b/utils/tokenization.py index ee154a5..80b06d1 100644 --- a/utils/tokenization.py +++ b/utils/tokenization.py @@ -6,9 +6,9 @@ from .text_utils import add_space_after_chinese from transformers import AutoModelForSeq2SeqLM, AutoTokenizer -tokenizer1 = AutoTokenizer.from_pretrained("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again")#("/home/lzx/T5-base/model3/mt5-base-trained-final-500+500-2-7_again")#("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again")#("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again")# +#("/home/lzx/T5-base/model3/mt5-base-trained-final-500+500-2-7_again")#("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again")#("/data/lbq/models/mt5-base-trained-final-500+500-2-7_again")# -def delete_blank(tokenized_inputs, max_seq=512): +def delete_blank(tokenized_inputs, max_seq=256): new_tokenized_inputs = defaultdict(list) for example_ids, example_mask in zip(tokenized_inputs['input_ids'], tokenized_inputs['attention_mask']): @@ -20,17 +20,20 @@ def delete_blank(tokenized_inputs, max_seq=512): new_tokenized_inputs['input_ids'].extend([0]*(max_seq - len(new_tokenized_inputs['input_ids']))) new_tokenized_inputs['attention_mask'].extend([0]*(max_seq - len(new_tokenized_inputs['attention_mask']))) - tokenized_inputs = {k: torch.tensor(v) for k,v in new_tokenized_inputs.items()} + tokenized_inputs = {k: torch.tensor([v,]) for k,v in new_tokenized_inputs.items()} return tokenized_inputs + @DatasetsProcessorNameSpace.register("ours") def tokenize_function(examples, tokenizer): # examples = ptr_change(examples) - examples.natural_sentence = [add_space_after_chinese(s.replace("得到","") ) for s in examples.natural_sentence] - tokenized_inputs = tokenizer(" ".join(examples.natural_sentence), padding="max_length", truncation=True, max_length=512, return_tensors="pt") - # tokenized_inputs = delete_blank(tokenized_inputs) + # examples.natural_sentence = [add_space_after_chinese(s.replace("得到","") ) for s in examples.natural_sentence] + tokenized_inputs = tokenizer(examples.natural_sentence, padding="max_length", truncation=True, max_length=256, return_tensors="pt") + + tokenized_inputs = delete_blank(tokenized_inputs) + # print([tokenizer.decode(i) for i in tokenized_inputs['input_ids'][0]]) - tokenized_labels = tokenizer(examples.expression, padding="max_length", truncation=True, max_length=512, return_tensors="pt") + tokenized_labels = tokenizer(examples.expression, padding="max_length", truncation=True, max_length=256, return_tensors="pt") tokenized_inputs['labels'] = tokenized_labels['input_ids'] tokenized_inputs['expression'] = examples.expression @@ -40,9 +43,9 @@ def tokenize_function(examples, tokenizer): @DatasetsProcessorNameSpace.register("topv2") def tokenize_function(examples, tokenizer): - examples = ptr_change(examples) - tokenized_inputs = tokenizer(examples['utterance'], padding='max_length', truncation=True, max_length=30, return_tensors="pt") - tokenized_labels = tokenizer(examples['semantic_parse'], padding='max_length', truncation=True, max_length=30, return_tensors="pt") + # examples = ptr_change(examples) + tokenized_inputs = tokenizer(examples['utterance'], padding='max_length', truncation=True, max_length=128, return_tensors="pt") + tokenized_labels = tokenizer(examples['semantic_parse'], padding='max_length', truncation=True, max_length=128, return_tensors="pt") tokenized_inputs['labels'] = tokenized_labels['input_ids'] return tokenized_inputs @@ -50,8 +53,7 @@ def tokenize_function(examples, tokenizer): #zcl def tokenize_function_zcl(examples, tokenizer): examples = ptr_change(examples) - global tokenizer1 - tokenizer = tokenizer1 + tokenizer = AutoTokenizer.from_pretrained("../tokenizer/") tokenized_inputs = tokenizer(examples['utterance'], padding='max_length', truncation=True, max_length=128, return_tensors="pt") tokenized_labels = tokenizer(examples['seqlogical'], padding='max_length', truncation=True, max_length=128, return_tensors="pt") @@ -76,7 +78,7 @@ def tokenizer_dataset(tokenizer, dataset: PreliminaryDataset) -> PreliminaryData @DatasetsProcessorNameSpace.register("topv2") def tokenizer_dataset(tokenizer, dataset): - tokenized_datasets = dataset.map(tokenize_function, tokenizer) + tokenized_datasets = dataset.map(tokenize_function, fn_kwargs={"tokenizer": tokenizer}) tokenized_datasets = tokenized_datasets.remove_columns(["utterance"]) tokenized_datasets = tokenized_datasets.remove_columns(["semantic_parse"]) tokenized_datasets = tokenized_datasets.remove_columns(["domain"])