wehub-resource-sync
188 lines
9.6 KiB
Python
188 lines
9.6 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2020-07-05 13:50
|
|
from typing import Optional, Union, List, Any, Dict, Tuple
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from hanlp_common.configurable import AutoConfigurable
|
|
from hanlp.layers.embeddings.embedding import Embedding
|
|
from hanlp.layers.scalar_mix import ScalarMixWithDropoutBuilder
|
|
from hanlp.layers.transformers.encoder import TransformerEncoder
|
|
from hanlp.layers.transformers.pt_imports import PreTrainedTokenizer, AutoConfig_, AutoTokenizer_
|
|
from hanlp.transform.transformer_tokenizer import TransformerSequenceTokenizer
|
|
|
|
|
|
class ContextualWordEmbeddingModule(TransformerEncoder):
|
|
def __init__(self,
|
|
field: str,
|
|
transformer: str,
|
|
transformer_tokenizer: PreTrainedTokenizer,
|
|
average_subwords=False,
|
|
scalar_mix: Union[ScalarMixWithDropoutBuilder, int] = None,
|
|
word_dropout=None,
|
|
max_sequence_length=None,
|
|
ret_raw_hidden_states=False,
|
|
transformer_args: Dict[str, Any] = None,
|
|
trainable=True,
|
|
training=True) -> None:
|
|
"""A contextualized word embedding module.
|
|
|
|
Args:
|
|
field: The field to work on. Usually some token fields.
|
|
transformer: An identifier of a ``PreTrainedModel``.
|
|
transformer_tokenizer:
|
|
average_subwords: ``True`` to average subword representations.
|
|
scalar_mix: Layer attention.
|
|
word_dropout: Dropout rate of randomly replacing a subword with MASK.
|
|
max_sequence_length: The maximum sequence length. Sequence longer than this will be handled by sliding
|
|
window.
|
|
ret_raw_hidden_states: ``True`` to return hidden states of each layer.
|
|
transformer_args: Extra arguments passed to the transformer.
|
|
trainable: ``False`` to use static embeddings.
|
|
training: ``False`` to skip loading weights from pre-trained transformers.
|
|
"""
|
|
super().__init__(transformer, transformer_tokenizer, average_subwords, scalar_mix, word_dropout,
|
|
max_sequence_length, ret_raw_hidden_states, transformer_args, trainable,
|
|
training)
|
|
self.field = field
|
|
|
|
# noinspection PyMethodOverriding
|
|
# noinspection PyTypeChecker
|
|
def forward(self, batch: dict, mask=None, **kwargs):
|
|
input_ids: torch.LongTensor = batch[f'{self.field}_input_ids']
|
|
token_span: torch.LongTensor = batch.get(f'{self.field}_token_span', None)
|
|
# input_device = input_ids.device
|
|
# this_device = self.get_device()
|
|
# if input_device != this_device:
|
|
# input_ids = input_ids.to(this_device)
|
|
# token_span = token_span.to(this_device)
|
|
# We might want to apply mask here
|
|
output: Union[torch.Tensor, List[torch.Tensor]] = super().forward(input_ids, token_span=token_span, **kwargs)
|
|
# if input_device != this_device:
|
|
# if isinstance(output, torch.Tensor):
|
|
# output = output.to(input_device)
|
|
# else:
|
|
# output = [x.to(input_device) for x in output]
|
|
return output
|
|
|
|
def get_output_dim(self):
|
|
return self.transformer.config.hidden_size
|
|
|
|
def get_device(self):
|
|
device: torch.device = next(self.parameters()).device
|
|
return device
|
|
|
|
|
|
class ContextualWordEmbedding(Embedding, AutoConfigurable):
|
|
def __init__(self, field: str,
|
|
transformer: str,
|
|
average_subwords=False,
|
|
scalar_mix: Union[ScalarMixWithDropoutBuilder, int] = None,
|
|
word_dropout: Optional[Union[float, Tuple[float, str]]] = None,
|
|
max_sequence_length=None,
|
|
truncate_long_sequences=False,
|
|
cls_is_bos=False,
|
|
sep_is_eos=False,
|
|
ret_token_span=True,
|
|
ret_subtokens=False,
|
|
ret_subtokens_group=False,
|
|
ret_prefix_mask=False,
|
|
ret_raw_hidden_states=False,
|
|
transformer_args: Dict[str, Any] = None,
|
|
use_fast=True,
|
|
do_basic_tokenize=True,
|
|
trainable=True) -> None:
|
|
"""A contextual word embedding builder which builds a
|
|
:class:`~hanlp.layers.embeddings.contextual_word_embedding.ContextualWordEmbeddingModule` and a
|
|
:class:`~hanlp.transform.transformer_tokenizer.TransformerSequenceTokenizer`.
|
|
|
|
Args:
|
|
field: The field to work on. Usually some token fields.
|
|
transformer: An identifier of a ``PreTrainedModel``.
|
|
average_subwords: ``True`` to average subword representations.
|
|
scalar_mix: Layer attention.
|
|
word_dropout: Dropout rate of randomly replacing a subword with MASK.
|
|
max_sequence_length: The maximum sequence length. Sequence longer than this will be handled by sliding
|
|
window.
|
|
truncate_long_sequences: ``True`` to return hidden states of each layer.
|
|
cls_is_bos: ``True`` means the first token of input is treated as [CLS] no matter what its surface form is.
|
|
``False`` (default) means the first token is not [CLS], it will have its own embedding other than
|
|
the embedding of [CLS].
|
|
sep_is_eos: ``True`` means the last token of input is [SEP].
|
|
``False`` means it's not but [SEP] will be appended,
|
|
``None`` means it dependents on `input[-1] == [EOS]`.
|
|
ret_token_span: ``True`` to return span of each token measured by subtoken offsets.
|
|
ret_subtokens: ``True`` to return list of subtokens belonging to each token.
|
|
ret_subtokens_group: ``True`` to return list of offsets of subtokens belonging to each token.
|
|
ret_prefix_mask: ``True`` to generate a mask where each non-zero element corresponds to a prefix of a token.
|
|
ret_raw_hidden_states: ``True`` to return hidden states of each layer.
|
|
transformer_args: Extra arguments passed to the transformer.
|
|
use_fast: Whether or not to try to load the fast version of the tokenizer.
|
|
do_basic_tokenize: Whether to do basic tokenization before wordpiece.
|
|
trainable: ``False`` to use static embeddings.
|
|
"""
|
|
super().__init__()
|
|
self.truncate_long_sequences = truncate_long_sequences
|
|
self.transformer_args = transformer_args
|
|
self.trainable = trainable
|
|
self.ret_subtokens_group = ret_subtokens_group
|
|
self.ret_subtokens = ret_subtokens
|
|
self.ret_raw_hidden_states = ret_raw_hidden_states
|
|
self.sep_is_eos = sep_is_eos
|
|
self.cls_is_bos = cls_is_bos
|
|
self.max_sequence_length = max_sequence_length
|
|
self.word_dropout = word_dropout
|
|
self.scalar_mix = scalar_mix
|
|
self.average_subwords = average_subwords
|
|
self.transformer = transformer
|
|
self.field = field
|
|
self._transformer_tokenizer = AutoTokenizer_.from_pretrained(self.transformer,
|
|
use_fast=use_fast,
|
|
do_basic_tokenize=do_basic_tokenize)
|
|
self._tokenizer_transform = TransformerSequenceTokenizer(self._transformer_tokenizer,
|
|
field,
|
|
truncate_long_sequences=truncate_long_sequences,
|
|
ret_prefix_mask=ret_prefix_mask,
|
|
ret_token_span=ret_token_span,
|
|
cls_is_bos=cls_is_bos,
|
|
sep_is_eos=sep_is_eos,
|
|
ret_subtokens=ret_subtokens,
|
|
ret_subtokens_group=ret_subtokens_group,
|
|
max_seq_length=self.max_sequence_length
|
|
)
|
|
|
|
def transform(self, **kwargs) -> TransformerSequenceTokenizer:
|
|
return self._tokenizer_transform
|
|
|
|
def module(self, training=True, **kwargs) -> Optional[nn.Module]:
|
|
return ContextualWordEmbeddingModule(self.field,
|
|
self.transformer,
|
|
self._transformer_tokenizer,
|
|
self.average_subwords,
|
|
self.scalar_mix,
|
|
self.word_dropout,
|
|
self.max_sequence_length,
|
|
self.ret_raw_hidden_states,
|
|
self.transformer_args,
|
|
self.trainable,
|
|
training=training)
|
|
|
|
def get_output_dim(self):
|
|
config = AutoConfig_.from_pretrained(self.transformer)
|
|
return config.hidden_size
|
|
|
|
def get_tokenizer(self):
|
|
return self._transformer_tokenizer
|
|
|
|
|
|
def find_transformer(embed: nn.Module):
|
|
if isinstance(embed, ContextualWordEmbeddingModule):
|
|
return embed
|
|
if isinstance(embed, nn.ModuleList):
|
|
for child in embed:
|
|
found = find_transformer(child)
|
|
if found:
|
|
return found
|