# -*- 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