wehub-resource-sync
194 lines
9.4 KiB
Python
194 lines
9.4 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2019-11-10 13:19
|
|
|
|
import math
|
|
from typing import Union, Tuple, Any, Iterable
|
|
import tensorflow as tf
|
|
from hanlp.common.keras_component import KerasComponent
|
|
from hanlp_common.structure import SerializableDict
|
|
from hanlp.layers.transformers.loader_tf import build_transformer
|
|
from hanlp.optimizers.adamw import create_optimizer
|
|
from hanlp.transform.table_tf import TableTransform
|
|
from hanlp.utils.log_util import logger
|
|
from hanlp_common.util import merge_locals_kwargs
|
|
from transformers.tokenization_utils import PreTrainedTokenizer
|
|
|
|
class TransformerTextTransform(TableTransform):
|
|
|
|
def __init__(self, config: SerializableDict = None, map_x=False, map_y=True, x_columns=None,
|
|
y_column=-1, skip_header=True, delimiter='auto', multi_label=False, **kwargs) -> None:
|
|
super().__init__(config, map_x, map_y, x_columns, y_column, multi_label, skip_header, delimiter, **kwargs)
|
|
self.tokenizer: PreTrainedTokenizer = None
|
|
|
|
def inputs_to_samples(self, inputs, gold=False):
|
|
tokenizer = self.tokenizer
|
|
max_length = self.config.max_length
|
|
num_features = None
|
|
pad_token = None if self.label_vocab.mutable else tokenizer.convert_tokens_to_ids(['[PAD]'])[0]
|
|
for (X, Y) in super().inputs_to_samples(inputs, gold):
|
|
if self.label_vocab.mutable:
|
|
yield None, Y
|
|
continue
|
|
if isinstance(X, str):
|
|
X = (X,)
|
|
if num_features is None:
|
|
num_features = self.config.num_features
|
|
assert num_features == len(X), f'Numbers of features {num_features} ' \
|
|
f'inconsistent with current {len(X)}={X}'
|
|
text_a = X[0]
|
|
text_b = X[1] if len(X) > 1 else None
|
|
tokens_a = self.tokenizer.tokenize(text_a)
|
|
tokens_b = self.tokenizer.tokenize(text_b) if text_b else None
|
|
tokens = ["[CLS]"] + tokens_a + ["[SEP]"]
|
|
segment_ids = [0] * len(tokens)
|
|
if tokens_b:
|
|
tokens += tokens_b
|
|
segment_ids += [1] * len(tokens_b)
|
|
token_ids = self.tokenizer.convert_tokens_to_ids(tokens)
|
|
attention_mask = [1] * len(token_ids)
|
|
diff = max_length - len(token_ids)
|
|
if diff < 0:
|
|
# logger.warning(
|
|
# f'Input tokens {tokens} exceed the max sequence length of {max_length - 2}. '
|
|
# f'The exceeded part will be truncated and ignored. '
|
|
# f'You are recommended to split your long text into several sentences within '
|
|
# f'{max_length - 2} tokens beforehand.')
|
|
token_ids = token_ids[:max_length]
|
|
attention_mask = attention_mask[:max_length]
|
|
segment_ids = segment_ids[:max_length]
|
|
elif diff > 0:
|
|
token_ids += [pad_token] * diff
|
|
attention_mask += [0] * diff
|
|
segment_ids += [0] * diff
|
|
|
|
assert len(token_ids) == max_length, "Error with input length {} vs {}".format(len(token_ids), max_length)
|
|
assert len(attention_mask) == max_length, "Error with input length {} vs {}".format(len(attention_mask),
|
|
max_length)
|
|
assert len(segment_ids) == max_length, "Error with input length {} vs {}".format(len(segment_ids),
|
|
max_length)
|
|
|
|
label = Y
|
|
yield (token_ids, attention_mask, segment_ids), label
|
|
|
|
def create_types_shapes_values(self) -> Tuple[Tuple, Tuple, Tuple]:
|
|
max_length = self.config.max_length
|
|
types = (tf.int32, tf.int32, tf.int32), tf.string
|
|
shapes = ([max_length], [max_length], [max_length]), [None, ] if self.config.get('multi_label', None) else []
|
|
values = (0, 0, 0), self.label_vocab.safe_pad_token
|
|
return types, shapes, values
|
|
|
|
def x_to_idx(self, x) -> Union[tf.Tensor, Tuple]:
|
|
logger.fatal('map_x should always be set to True')
|
|
exit(1)
|
|
|
|
def y_to_idx(self, y) -> tf.Tensor:
|
|
if self.config.get('multi_label', None):
|
|
# need to change index to binary vector
|
|
mapped = tf.map_fn(fn=lambda x: tf.cast(self.label_vocab.lookup(x), tf.int32), elems=y,
|
|
fn_output_signature=tf.TensorSpec(dtype=tf.dtypes.int32, shape=[None, ]))
|
|
one_hots = tf.one_hot(mapped, len(self.label_vocab))
|
|
idx = tf.reduce_sum(one_hots, -2)
|
|
else:
|
|
idx = self.label_vocab.lookup(y)
|
|
return idx
|
|
|
|
def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None,
|
|
batch=None) -> Iterable:
|
|
# Prediction to be Y > 0:
|
|
if self.config.get('multi_label', None):
|
|
preds = Y
|
|
else:
|
|
preds = tf.argmax(Y, axis=-1)
|
|
for y in preds:
|
|
yield self.label_vocab.idx_to_token[y]
|
|
|
|
def input_is_single_sample(self, input: Any) -> bool:
|
|
return isinstance(input, (str, tuple))
|
|
|
|
|
|
class TransformerClassifierTF(KerasComponent):
|
|
|
|
def __init__(self, bert_text_transform=None) -> None:
|
|
if not bert_text_transform:
|
|
bert_text_transform = TransformerTextTransform()
|
|
super().__init__(bert_text_transform)
|
|
self.model: tf.keras.Model
|
|
self.transform: TransformerTextTransform = bert_text_transform
|
|
|
|
# noinspection PyMethodOverriding
|
|
def fit(self, trn_data: Any, dev_data: Any, save_dir: str, transformer: str, max_length: int = 128,
|
|
optimizer='adamw', warmup_steps_ratio=0.1, use_amp=False, batch_size=32,
|
|
epochs=3, logger=None, verbose=1, **kwargs):
|
|
return super().fit(**merge_locals_kwargs(locals(), kwargs))
|
|
|
|
def evaluate_output(self, tst_data, out, num_batches, metric):
|
|
out.write('sentence\tpred\tgold\n')
|
|
total, correct, score = 0, 0, 0
|
|
for idx, batch in enumerate(tst_data):
|
|
outputs = self.model.predict_on_batch(batch[0])
|
|
outputs = tf.argmax(outputs, axis=1)
|
|
for X, Y_pred, Y_gold, in zip(batch[0][0], outputs, batch[1]):
|
|
feature = ' '.join(self.transform.tokenizer.convert_ids_to_tokens(X.numpy()))
|
|
feature = feature.replace(' ##', '') # fix sub-word generated by BERT tagger
|
|
out.write('{}\t{}\t{}\n'.format(feature,
|
|
self._y_id_to_str(Y_pred),
|
|
self._y_id_to_str(Y_gold)))
|
|
total += 1
|
|
correct += int(tf.equal(Y_pred, Y_gold).numpy())
|
|
score = correct / total
|
|
print('\r{}/{} {}: {:.2f}'.format(idx + 1, num_batches, metric, score * 100), end='')
|
|
print()
|
|
return score
|
|
|
|
def _y_id_to_str(self, Y_pred) -> str:
|
|
return self.transform.label_vocab.idx_to_token[Y_pred.numpy()]
|
|
|
|
def build_loss(self, loss, **kwargs):
|
|
if loss:
|
|
assert isinstance(loss, tf.keras.losses.loss), 'Must specify loss as an instance in tf.keras.losses'
|
|
return loss
|
|
elif self.config.get('multi_label', None):
|
|
# Loss to be BinaryCrossentropy for multi-label:
|
|
loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
|
|
else:
|
|
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
|
|
return loss
|
|
|
|
# noinspection PyMethodOverriding
|
|
def build_optimizer(self, optimizer, use_amp, train_steps, warmup_steps, **kwargs):
|
|
if optimizer == 'adamw':
|
|
opt = create_optimizer(init_lr=5e-5, num_train_steps=train_steps, num_warmup_steps=warmup_steps)
|
|
# opt = tfa.optimizers.AdamW(learning_rate=3e-5, epsilon=1e-08, weight_decay=0.01)
|
|
# opt = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08)
|
|
self.config.optimizer = tf.keras.utils.serialize_keras_object(opt)
|
|
lr_config = self.config.optimizer['config']['learning_rate']['config']
|
|
if hasattr(lr_config['decay_schedule_fn'], 'get_config'):
|
|
lr_config['decay_schedule_fn'] = dict(
|
|
(k, v) for k, v in lr_config['decay_schedule_fn'].config().items() if not k.startswith('_'))
|
|
else:
|
|
opt = super().build_optimizer(optimizer)
|
|
if use_amp:
|
|
# loss scaling is currently required when using mixed precision
|
|
opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(opt, 'dynamic')
|
|
return opt
|
|
|
|
# noinspection PyMethodOverriding
|
|
def build_model(self, transformer, max_length, **kwargs):
|
|
model, self.transform.tokenizer = build_transformer(transformer, max_length, len(self.transform.label_vocab),
|
|
tagging=False)
|
|
return model
|
|
|
|
def build_vocab(self, trn_data, logger):
|
|
train_examples = super().build_vocab(trn_data, logger)
|
|
warmup_steps_per_epoch = math.ceil(train_examples * self.config.warmup_steps_ratio / self.config.batch_size)
|
|
self.config.warmup_steps = warmup_steps_per_epoch * self.config.epochs
|
|
return train_examples
|
|
|
|
def build_metrics(self, metrics, logger, **kwargs):
|
|
if self.config.get('multi_label', None):
|
|
metric = tf.keras.metrics.BinaryAccuracy('binary_accuracy')
|
|
else:
|
|
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
|
|
return [metric]
|