wehub-resource-sync
370 lines
17 KiB
Python
370 lines
17 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2020-06-15 21:22
|
|
from collections import defaultdict
|
|
from typing import Tuple, Union
|
|
|
|
import torch
|
|
from torch.nn import functional as F
|
|
|
|
from hanlp.components.parsers.ud import udify_util as util
|
|
from hanlp.layers.transformers.pt_imports import PreTrainedModel
|
|
|
|
|
|
def transformer_encode(transformer: PreTrainedModel,
|
|
input_ids,
|
|
attention_mask=None,
|
|
token_type_ids=None,
|
|
token_span=None,
|
|
layer_range: Union[int, Tuple[int, int]] = 0,
|
|
max_sequence_length=None,
|
|
average_subwords=False,
|
|
ret_raw_hidden_states=False):
|
|
"""Run transformer and pool its outputs.
|
|
|
|
Args:
|
|
transformer: A transformer model.
|
|
input_ids: Indices of subwords.
|
|
attention_mask: Mask for these subwords.
|
|
token_type_ids: Type ids for each subword.
|
|
token_span: The spans of tokens.
|
|
layer_range: The range of layers to use. Note that the 0-th layer means embedding layer, so the last 3 layers
|
|
of a 12-layer BERT will be (10, 13).
|
|
max_sequence_length: The maximum sequence length. Sequence longer than this will be handled by sliding
|
|
window.
|
|
average_subwords: ``True`` to average subword representations.
|
|
ret_raw_hidden_states: ``True`` to return hidden states of each layer.
|
|
|
|
Returns:
|
|
Pooled outputs.
|
|
|
|
"""
|
|
if max_sequence_length and input_ids.size(-1) > max_sequence_length:
|
|
# TODO: split token type ids in transformer_sliding_window if token type ids are not always 1
|
|
outputs = transformer_sliding_window(transformer, input_ids, max_pieces=max_sequence_length)
|
|
else:
|
|
if attention_mask is None:
|
|
attention_mask = input_ids.ne(0)
|
|
if transformer.config.output_hidden_states:
|
|
outputs = transformer(input_ids, attention_mask, token_type_ids)[-1]
|
|
else:
|
|
outputs = transformer(input_ids, attention_mask, token_type_ids)[0]
|
|
if transformer.config.output_hidden_states:
|
|
if isinstance(layer_range, int):
|
|
outputs = outputs[layer_range:]
|
|
else:
|
|
outputs = outputs[layer_range[0], layer_range[1]]
|
|
# Slow pick
|
|
# hs = []
|
|
# for h in outputs:
|
|
# hs.append(pick_tensor_for_each_token(h, token_span, average_subwords))
|
|
# Fast pick
|
|
if not isinstance(outputs, torch.Tensor):
|
|
x = torch.stack(outputs)
|
|
else:
|
|
x = outputs
|
|
L, B, T, F = x.size()
|
|
x = x.flatten(end_dim=1)
|
|
# tile token_span as x
|
|
if token_span is not None:
|
|
token_span = token_span.repeat(L, 1, 1)
|
|
hs = pick_tensor_for_each_token(x, token_span, average_subwords).view(L, B, -1, F)
|
|
if ret_raw_hidden_states:
|
|
return hs, outputs
|
|
return hs
|
|
else:
|
|
if ret_raw_hidden_states:
|
|
return pick_tensor_for_each_token(outputs, token_span, average_subwords), outputs
|
|
return pick_tensor_for_each_token(outputs, token_span, average_subwords)
|
|
|
|
|
|
def pick_tensor_for_each_token(h, token_span, average_subwords):
|
|
if token_span is None:
|
|
return h
|
|
if average_subwords and token_span.size(-1) > 1:
|
|
batch_size = h.size(0)
|
|
h_span = h.gather(1, token_span.view(batch_size, -1).unsqueeze(-1).expand(-1, -1, h.shape[-1]))
|
|
h_span = h_span.view(batch_size, *token_span.shape[1:], -1)
|
|
n_sub_tokens = token_span.ne(0)
|
|
n_sub_tokens[:, 0, 0] = True
|
|
h_span = (h_span * n_sub_tokens.unsqueeze(-1)).sum(2)
|
|
n_sub_tokens = n_sub_tokens.sum(-1).unsqueeze(-1)
|
|
zero_mask = n_sub_tokens == 0
|
|
if torch.any(zero_mask):
|
|
n_sub_tokens[zero_mask] = 1 # avoid dividing by zero
|
|
embed = h_span / n_sub_tokens
|
|
else:
|
|
embed = h.gather(1, token_span[:, :, 0].unsqueeze(-1).expand(-1, -1, h.size(-1)))
|
|
return embed
|
|
|
|
|
|
def transformer_sliding_window(transformer: PreTrainedModel,
|
|
input_ids: torch.LongTensor,
|
|
input_mask=None,
|
|
offsets: torch.LongTensor = None,
|
|
token_type_ids: torch.LongTensor = None,
|
|
max_pieces=512,
|
|
start_tokens: int = 1,
|
|
end_tokens: int = 1,
|
|
ret_cls=None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
|
|
Args:
|
|
transformer:
|
|
input_ids: torch.LongTensor:
|
|
input_mask: (Default value = None)
|
|
offsets: torch.LongTensor: (Default value = None)
|
|
token_type_ids: torch.LongTensor: (Default value = None)
|
|
max_pieces: (Default value = 512)
|
|
start_tokens: int: (Default value = 1)
|
|
end_tokens: int: (Default value = 1)
|
|
ret_cls: (Default value = None)
|
|
|
|
Returns:
|
|
|
|
|
|
"""
|
|
# pylint: disable=arguments-differ
|
|
batch_size, full_seq_len = input_ids.size(0), input_ids.size(-1)
|
|
initial_dims = list(input_ids.shape[:-1])
|
|
|
|
# The embedder may receive an input tensor that has a sequence length longer than can
|
|
# be fit. In that case, we should expect the wordpiece indexer to create padded windows
|
|
# of length `max_pieces` for us, and have them concatenated into one long sequence.
|
|
# E.g., "[CLS] I went to the [SEP] [CLS] to the store to [SEP] ..."
|
|
# We can then split the sequence into sub-sequences of that length, and concatenate them
|
|
# along the batch dimension so we effectively have one huge batch of partial sentences.
|
|
# This can then be fed into BERT without any sentence length issues. Keep in mind
|
|
# that the memory consumption can dramatically increase for large batches with extremely
|
|
# long sentences.
|
|
needs_split = full_seq_len > max_pieces
|
|
if needs_split:
|
|
input_ids = split_to_sliding_window(input_ids, max_pieces)
|
|
|
|
# if token_type_ids is None:
|
|
# token_type_ids = torch.zeros_like(input_ids)
|
|
if input_mask is None:
|
|
input_mask = (input_ids != 0).long()
|
|
|
|
# input_ids may have extra dimensions, so we reshape down to 2-d
|
|
# before calling the BERT model and then reshape back at the end.
|
|
outputs = transformer(input_ids=util.combine_initial_dims_to_1d_or_2d(input_ids),
|
|
# token_type_ids=util.combine_initial_dims_to_1d_or_2d(token_type_ids),
|
|
attention_mask=util.combine_initial_dims_to_1d_or_2d(input_mask)).to_tuple()
|
|
if len(outputs) == 3:
|
|
all_encoder_layers = outputs.hidden_states
|
|
all_encoder_layers = torch.stack(all_encoder_layers)
|
|
elif len(outputs) == 2:
|
|
all_encoder_layers, _ = outputs[:2]
|
|
else:
|
|
all_encoder_layers = outputs[0]
|
|
|
|
if needs_split:
|
|
if ret_cls is not None:
|
|
cls_mask = input_ids[:, 0] == input_ids[0][0]
|
|
cls_hidden = all_encoder_layers[:, 0, :]
|
|
if ret_cls == 'max':
|
|
cls_hidden[~cls_mask] = -1e20
|
|
else:
|
|
cls_hidden[~cls_mask] = 0
|
|
cls_mask = cls_mask.view(-1, batch_size).transpose(0, 1)
|
|
cls_hidden = cls_hidden.reshape(cls_mask.size(1), batch_size, -1).transpose(0, 1)
|
|
if ret_cls == 'max':
|
|
cls_hidden = cls_hidden.max(1)[0]
|
|
elif ret_cls == 'raw':
|
|
return cls_hidden, cls_mask
|
|
else:
|
|
cls_hidden = torch.sum(cls_hidden, dim=1)
|
|
cls_hidden /= torch.sum(cls_mask, dim=1, keepdim=True)
|
|
return cls_hidden
|
|
else:
|
|
recombined_embeddings, select_indices = restore_from_sliding_window(all_encoder_layers, batch_size,
|
|
max_pieces, full_seq_len, start_tokens,
|
|
end_tokens)
|
|
|
|
initial_dims.append(len(select_indices))
|
|
else:
|
|
recombined_embeddings = all_encoder_layers
|
|
|
|
# Recombine the outputs of all layers
|
|
# (layers, batch_size * d1 * ... * dn, sequence_length, embedding_dim)
|
|
# recombined = torch.cat(combined, dim=2)
|
|
# input_mask = (recombined_embeddings != 0).long()
|
|
|
|
# At this point, mix is (batch_size * d1 * ... * dn, sequence_length, embedding_dim)
|
|
|
|
if offsets is None:
|
|
# Resize to (batch_size, d1, ..., dn, sequence_length, embedding_dim)
|
|
dims = initial_dims if needs_split else input_ids.size()
|
|
layers = util.uncombine_initial_dims(recombined_embeddings, dims)
|
|
else:
|
|
# offsets is (batch_size, d1, ..., dn, orig_sequence_length)
|
|
offsets2d = util.combine_initial_dims_to_1d_or_2d(offsets)
|
|
# now offsets is (batch_size * d1 * ... * dn, orig_sequence_length)
|
|
range_vector = util.get_range_vector(offsets2d.size(0),
|
|
device=util.get_device_of(recombined_embeddings)).unsqueeze(1)
|
|
# selected embeddings is also (batch_size * d1 * ... * dn, orig_sequence_length)
|
|
selected_embeddings = recombined_embeddings[:, range_vector, offsets2d]
|
|
|
|
layers = util.uncombine_initial_dims(selected_embeddings, offsets.size())
|
|
|
|
return layers
|
|
|
|
|
|
def split_to_sliding_window(input_ids, max_pieces):
|
|
# Split the flattened list by the window size, `max_pieces`
|
|
split_input_ids = list(input_ids.split(max_pieces, dim=-1))
|
|
# We want all sequences to be the same length, so pad the last sequence
|
|
last_window_size = split_input_ids[-1].size(-1)
|
|
padding_amount = max_pieces - last_window_size
|
|
split_input_ids[-1] = F.pad(split_input_ids[-1], pad=[0, padding_amount], value=0)
|
|
# Now combine the sequences along the batch dimension
|
|
input_ids = torch.cat(split_input_ids, dim=0)
|
|
return input_ids
|
|
|
|
|
|
def restore_from_sliding_window(all_encoder_layers, batch_size, max_pieces, full_seq_len, start_tokens, end_tokens):
|
|
# First, unpack the output embeddings into one long sequence again
|
|
unpacked_embeddings = torch.split(all_encoder_layers, batch_size, dim=-3)
|
|
unpacked_embeddings = torch.cat(unpacked_embeddings, dim=-2)
|
|
# Next, select indices of the sequence such that it will result in embeddings representing the original
|
|
# sentence. To capture maximal context, the indices will be the middle part of each embedded window
|
|
# sub-sequence (plus any leftover start and final edge windows), e.g.,
|
|
# 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
|
|
# "[CLS] I went to the very fine [SEP] [CLS] the very fine store to eat [SEP]"
|
|
# with max_pieces = 8 should produce max context indices [2, 3, 4, 10, 11, 12] with additional start
|
|
# and final windows with indices [0, 1] and [14, 15] respectively.
|
|
# Find the stride as half the max pieces, ignoring the special start and end tokens
|
|
# Calculate an offset to extract the centermost embeddings of each window
|
|
stride = (max_pieces - start_tokens - end_tokens) // 2
|
|
stride_offset = stride // 2 + start_tokens
|
|
first_window = list(range(stride_offset))
|
|
max_context_windows = [i for i in range(full_seq_len)
|
|
if stride_offset - 1 < i % max_pieces < stride_offset + stride]
|
|
final_window_start = max_context_windows[-1] + 1
|
|
final_window = list(range(final_window_start, full_seq_len))
|
|
select_indices = first_window + max_context_windows + final_window
|
|
select_indices = torch.LongTensor(select_indices).to(unpacked_embeddings.device)
|
|
recombined_embeddings = unpacked_embeddings.index_select(-2, select_indices)
|
|
return recombined_embeddings, select_indices
|
|
|
|
|
|
def build_optimizer_for_pretrained(model: torch.nn.Module,
|
|
pretrained: torch.nn.Module,
|
|
lr=1e-5,
|
|
weight_decay=0.01,
|
|
eps=1e-8,
|
|
transformer_lr=None,
|
|
transformer_weight_decay=None,
|
|
no_decay=('bias', 'LayerNorm.bias', 'LayerNorm.weight'),
|
|
**kwargs):
|
|
if transformer_lr is None:
|
|
transformer_lr = lr
|
|
if transformer_weight_decay is None:
|
|
transformer_weight_decay = weight_decay
|
|
params = defaultdict(lambda: defaultdict(list))
|
|
pretrained = set(pretrained.parameters())
|
|
if isinstance(no_decay, tuple):
|
|
def no_decay_fn(name):
|
|
return any(nd in name for nd in no_decay)
|
|
else:
|
|
assert callable(no_decay), 'no_decay has to be callable or a tuple of str'
|
|
no_decay_fn = no_decay
|
|
for n, p in model.named_parameters():
|
|
is_pretrained = 'pretrained' if p in pretrained else 'non_pretrained'
|
|
is_no_decay = 'no_decay' if no_decay_fn(n) else 'decay'
|
|
params[is_pretrained][is_no_decay].append(p)
|
|
|
|
grouped_parameters = [
|
|
{'params': params['pretrained']['decay'], 'weight_decay': transformer_weight_decay, 'lr': transformer_lr},
|
|
{'params': params['pretrained']['no_decay'], 'weight_decay': 0.0, 'lr': transformer_lr},
|
|
{'params': params['non_pretrained']['decay'], 'weight_decay': weight_decay, 'lr': lr},
|
|
{'params': params['non_pretrained']['no_decay'], 'weight_decay': 0.0, 'lr': lr},
|
|
]
|
|
|
|
from transformers import optimization
|
|
return optimization.AdamW(
|
|
grouped_parameters,
|
|
lr=lr,
|
|
weight_decay=weight_decay,
|
|
eps=eps,
|
|
no_deprecation_warning=True, # For backwards compatability
|
|
**kwargs)
|
|
|
|
|
|
def build_optimizer_scheduler_with_transformer(model: torch.nn.Module,
|
|
transformer: torch.nn.Module,
|
|
lr: float,
|
|
transformer_lr: float,
|
|
num_training_steps: int,
|
|
warmup_steps: Union[float, int],
|
|
weight_decay: float,
|
|
adam_epsilon: float,
|
|
no_decay=('bias', 'LayerNorm.bias', 'LayerNorm.weight')):
|
|
optimizer = build_optimizer_for_pretrained(model,
|
|
transformer,
|
|
lr,
|
|
weight_decay,
|
|
eps=adam_epsilon,
|
|
transformer_lr=transformer_lr,
|
|
no_decay=no_decay)
|
|
if isinstance(warmup_steps, float):
|
|
assert 0 < warmup_steps < 1, 'warmup_steps has to fall in range (0, 1) when it is float.'
|
|
warmup_steps = num_training_steps * warmup_steps
|
|
from transformers import optimization
|
|
scheduler = optimization.get_linear_schedule_with_warmup(optimizer, warmup_steps, num_training_steps)
|
|
return optimizer, scheduler
|
|
|
|
|
|
def get_optimizers(
|
|
model: torch.nn.Module,
|
|
num_training_steps: int,
|
|
learning_rate=5e-5,
|
|
adam_epsilon=1e-8,
|
|
weight_decay=0.0,
|
|
warmup_steps=0.1,
|
|
) -> Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]:
|
|
"""
|
|
Modified from https://github.com/huggingface/transformers/blob/7b75aa9fa55bee577e2c7403301ed31103125a35/src/transformers/trainer.py#L232
|
|
Setup the optimizer and the learning rate scheduler.
|
|
|
|
We provide a reasonable default that works well.
|
|
"""
|
|
if isinstance(warmup_steps, float):
|
|
assert 0 < warmup_steps < 1
|
|
warmup_steps = int(num_training_steps * warmup_steps)
|
|
# Prepare optimizer and schedule (linear warmup and decay)
|
|
no_decay = ["bias", "LayerNorm.weight"]
|
|
optimizer_grouped_parameters = [
|
|
{
|
|
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
|
|
"weight_decay": weight_decay,
|
|
},
|
|
{
|
|
"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
|
|
"weight_decay": 0.0,
|
|
},
|
|
]
|
|
from transformers import AdamW, get_linear_schedule_with_warmup
|
|
optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=adam_epsilon)
|
|
scheduler = get_linear_schedule_with_warmup(
|
|
optimizer, num_warmup_steps=warmup_steps, num_training_steps=num_training_steps
|
|
)
|
|
return optimizer, scheduler
|
|
|
|
|
|
def collect_decay_params(model, weight_decay):
|
|
no_decay = ["bias", "LayerNorm.weight"]
|
|
optimizer_grouped_parameters = [
|
|
{
|
|
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
|
|
"weight_decay": weight_decay,
|
|
},
|
|
{
|
|
"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
|
|
"weight_decay": 0.0,
|
|
},
|
|
]
|
|
return optimizer_grouped_parameters
|