wehub-resource-sync
384 lines
17 KiB
Python
384 lines
17 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2020-06-08 16:31
|
|
import logging
|
|
from abc import ABC
|
|
from typing import Callable, Union
|
|
from typing import List
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader
|
|
|
|
from hanlp_common.constant import IDX
|
|
from hanlp.common.dataset import TableDataset, SortingSampler, PadSequenceDataLoader, TransformableDataset
|
|
from hanlp.common.torch_component import TorchComponent
|
|
from hanlp.common.vocab import Vocab
|
|
from hanlp.components.distillation.schedulers import LinearTeacherAnnealingScheduler
|
|
from hanlp.layers.scalar_mix import ScalarMixWithDropoutBuilder
|
|
from hanlp.layers.transformers.encoder import TransformerEncoder
|
|
from hanlp.layers.transformers.pt_imports import PreTrainedModel, AutoTokenizer, BertTokenizer, AutoTokenizer_
|
|
from hanlp.layers.transformers.utils import transformer_sliding_window, build_optimizer_scheduler_with_transformer
|
|
from hanlp.metrics.accuracy import CategoricalAccuracy
|
|
from hanlp.transform.transformer_tokenizer import TransformerTextTokenizer
|
|
from hanlp.utils.time_util import CountdownTimer
|
|
from hanlp_common.util import merge_locals_kwargs, merge_dict, isdebugging
|
|
|
|
|
|
class TransformerClassificationModel(nn.Module):
|
|
|
|
def __init__(self,
|
|
transformer: PreTrainedModel,
|
|
num_labels: int,
|
|
max_seq_length=512) -> None:
|
|
super().__init__()
|
|
self.max_seq_length = max_seq_length
|
|
self.transformer = transformer
|
|
self.dropout = nn.Dropout(transformer.config.hidden_dropout_prob)
|
|
self.classifier = nn.Linear(transformer.config.hidden_size, num_labels)
|
|
|
|
def forward(self, input_ids, attention_mask, token_type_ids):
|
|
seq_length = input_ids.size(-1)
|
|
if seq_length > self.max_seq_length:
|
|
sequence_output = transformer_sliding_window(self.transformer, input_ids,
|
|
max_pieces=self.max_seq_length, ret_cls='max')
|
|
else:
|
|
sequence_output = self.transformer(input_ids, attention_mask, token_type_ids)[0][:, 0, :]
|
|
sequence_output = self.dropout(sequence_output)
|
|
logits = self.classifier(sequence_output)
|
|
return logits
|
|
|
|
|
|
class TransformerComponent(TorchComponent, ABC):
|
|
def __init__(self, **kwargs) -> None:
|
|
""" The base class for transorfmer based components. If offers methods to build transformer tokenizers
|
|
, optimizers and models.
|
|
|
|
Args:
|
|
**kwargs: Passed to config.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self.transformer_tokenizer = None
|
|
|
|
def build_optimizer(self,
|
|
trn,
|
|
epochs,
|
|
lr,
|
|
adam_epsilon,
|
|
weight_decay,
|
|
warmup_steps,
|
|
transformer_lr=None,
|
|
teacher=None,
|
|
**kwargs):
|
|
num_training_steps = len(trn) * epochs // self.config.get('gradient_accumulation', 1)
|
|
if transformer_lr is None:
|
|
transformer_lr = lr
|
|
transformer = self.model.encoder.transformer
|
|
optimizer, scheduler = build_optimizer_scheduler_with_transformer(self.model, transformer,
|
|
lr, transformer_lr,
|
|
num_training_steps, warmup_steps,
|
|
weight_decay, adam_epsilon)
|
|
if teacher:
|
|
lambda_scheduler = LinearTeacherAnnealingScheduler(num_training_steps)
|
|
scheduler = (scheduler, lambda_scheduler)
|
|
return optimizer, scheduler
|
|
|
|
def fit(self, trn_data, dev_data, save_dir,
|
|
transformer=None,
|
|
lr=5e-5,
|
|
transformer_lr=None,
|
|
adam_epsilon=1e-8,
|
|
weight_decay=0,
|
|
warmup_steps=0.1,
|
|
batch_size=32,
|
|
gradient_accumulation=1,
|
|
grad_norm=5.0,
|
|
transformer_grad_norm=None,
|
|
average_subwords=False,
|
|
scalar_mix: Union[ScalarMixWithDropoutBuilder, int] = None,
|
|
word_dropout=None,
|
|
hidden_dropout=None,
|
|
max_seq_len=None,
|
|
ret_raw_hidden_states=False,
|
|
batch_max_tokens=None,
|
|
epochs=3,
|
|
logger=None,
|
|
devices: Union[float, int, List[int]] = None,
|
|
**kwargs):
|
|
return super().fit(**merge_locals_kwargs(locals(), kwargs))
|
|
|
|
def on_config_ready(self, **kwargs):
|
|
super().on_config_ready(**kwargs)
|
|
if 'albert_chinese' in self.config.transformer:
|
|
self.transformer_tokenizer = BertTokenizer.from_pretrained(self.config.transformer, use_fast=True)
|
|
else:
|
|
self.transformer_tokenizer = AutoTokenizer_.from_pretrained(self.config.transformer, use_fast=True)
|
|
|
|
def build_transformer(self, training=True):
|
|
transformer = TransformerEncoder(self.config.transformer, self.transformer_tokenizer,
|
|
self.config.average_subwords,
|
|
self.config.scalar_mix, self.config.word_dropout,
|
|
ret_raw_hidden_states=self.config.ret_raw_hidden_states,
|
|
training=training)
|
|
transformer_layers = self.config.get('transformer_layers', None)
|
|
if transformer_layers:
|
|
transformer.transformer.encoder.layer = transformer.transformer.encoder.layer[:transformer_layers]
|
|
return transformer
|
|
|
|
|
|
class TransformerClassifier(TransformerComponent):
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
"""A classifier using transformer as encoder.
|
|
|
|
Args:
|
|
**kwargs: Passed to config.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self.model: TransformerClassificationModel = None
|
|
|
|
def build_criterion(self, **kwargs):
|
|
criterion = nn.CrossEntropyLoss()
|
|
return criterion
|
|
|
|
def build_metric(self, **kwargs):
|
|
return CategoricalAccuracy()
|
|
|
|
def execute_training_loop(self, trn: DataLoader, dev: DataLoader, epochs, criterion, optimizer, metric, save_dir,
|
|
logger: logging.Logger, devices, **kwargs):
|
|
best_epoch, best_metric = 0, -1
|
|
timer = CountdownTimer(epochs)
|
|
ratio_width = len(f'{len(trn)}/{len(trn)}')
|
|
for epoch in range(1, epochs + 1):
|
|
logger.info(f"[yellow]Epoch {epoch} / {epochs}:[/yellow]")
|
|
self.fit_dataloader(trn, criterion, optimizer, metric, logger)
|
|
if dev:
|
|
self.evaluate_dataloader(dev, criterion, metric, logger, ratio_width=ratio_width)
|
|
report = f'{timer.elapsed_human}/{timer.total_time_human}'
|
|
dev_score = metric.get_metric()
|
|
if dev_score > best_metric:
|
|
self.save_weights(save_dir)
|
|
best_metric = dev_score
|
|
report += ' [red]saved[/red]'
|
|
timer.log(report, ratio_percentage=False, newline=True, ratio=False)
|
|
|
|
@property
|
|
def label_vocab(self):
|
|
return self.vocabs[self.config.label_key]
|
|
|
|
def fit_dataloader(self, trn: DataLoader, criterion, optimizer, metric, logger: logging.Logger, **kwargs):
|
|
self.model.train()
|
|
timer = CountdownTimer(len(trn))
|
|
optimizer, scheduler = optimizer
|
|
total_loss = 0
|
|
metric.reset()
|
|
for batch in trn:
|
|
optimizer.zero_grad()
|
|
logits = self.feed_batch(batch)
|
|
target = batch['label_id']
|
|
loss = self.compute_loss(criterion, logits, target, batch)
|
|
loss.backward()
|
|
optimizer.step()
|
|
scheduler.step()
|
|
total_loss += loss.item()
|
|
self.update_metric(metric, logits, target)
|
|
timer.log(f'loss: {total_loss / (timer.current + 1):.4f} acc: {metric.get_metric():.2%}',
|
|
ratio_percentage=None,
|
|
logger=logger)
|
|
del loss
|
|
return total_loss / timer.total
|
|
|
|
def update_metric(self, metric, logits: torch.Tensor, target, output=None):
|
|
metric(logits, target)
|
|
if output:
|
|
label_ids = logits.argmax(-1)
|
|
return label_ids
|
|
|
|
def compute_loss(self, criterion, logits, target, batch):
|
|
loss = criterion(logits, target)
|
|
return loss
|
|
|
|
def feed_batch(self, batch) -> torch.LongTensor:
|
|
logits = self.model(*[batch[key] for key in ['input_ids', 'attention_mask', 'token_type_ids']])
|
|
return logits
|
|
|
|
# noinspection PyMethodOverriding
|
|
def evaluate_dataloader(self,
|
|
data: DataLoader,
|
|
criterion: Callable,
|
|
metric,
|
|
logger,
|
|
ratio_width=None,
|
|
filename=None,
|
|
output=None,
|
|
**kwargs):
|
|
self.model.eval()
|
|
timer = CountdownTimer(len(data))
|
|
total_loss = 0
|
|
metric.reset()
|
|
num_samples = 0
|
|
if output:
|
|
output = open(output, 'w')
|
|
for batch in data:
|
|
logits = self.feed_batch(batch)
|
|
target = batch['label_id']
|
|
loss = self.compute_loss(criterion, logits, target, batch)
|
|
total_loss += loss.item()
|
|
label_ids = self.update_metric(metric, logits, target, output)
|
|
if output:
|
|
labels = [self.vocabs[self.config.label_key].idx_to_token[i] for i in label_ids.tolist()]
|
|
for i, label in enumerate(labels):
|
|
# text_a text_b pred gold
|
|
columns = [batch[self.config.text_a_key][i]]
|
|
if self.config.text_b_key:
|
|
columns.append(batch[self.config.text_b_key][i])
|
|
columns.append(label)
|
|
columns.append(batch[self.config.label_key][i])
|
|
output.write('\t'.join(columns))
|
|
output.write('\n')
|
|
num_samples += len(target)
|
|
report = f'loss: {total_loss / (timer.current + 1):.4f} acc: {metric.get_metric():.2%}'
|
|
if filename:
|
|
report = f'{filename} {report} {num_samples / timer.elapsed:.0f} samples/sec'
|
|
timer.log(report, ratio_percentage=None, logger=logger, ratio_width=ratio_width)
|
|
if output:
|
|
output.close()
|
|
return total_loss / timer.total
|
|
|
|
# noinspection PyMethodOverriding
|
|
def build_model(self, transformer, training=True, **kwargs) -> torch.nn.Module:
|
|
# config: PretrainedConfig = AutoConfig.from_pretrained(transformer)
|
|
# config.num_labels = len(self.vocabs.label)
|
|
# config.hidden_dropout_prob = self.config.hidden_dropout_prob
|
|
transformer = self.build_transformer(training=training).transformer
|
|
model = TransformerClassificationModel(transformer, len(self.vocabs.label))
|
|
# truncated_normal_(model.classifier.weight, mean=0.02, std=0.05)
|
|
return model
|
|
|
|
# noinspection PyMethodOverriding
|
|
def build_dataloader(self, data, batch_size, shuffle, device, text_a_key, text_b_key,
|
|
label_key,
|
|
logger: logging.Logger = None,
|
|
sorting=True,
|
|
**kwargs) -> DataLoader:
|
|
if not batch_size:
|
|
batch_size = self.config.batch_size
|
|
dataset = self.build_dataset(data)
|
|
dataset.append_transform(self.vocabs)
|
|
if self.vocabs.mutable:
|
|
if not any([text_a_key, text_b_key]):
|
|
if len(dataset.headers) == 2:
|
|
self.config.text_a_key = dataset.headers[0]
|
|
self.config.label_key = dataset.headers[1]
|
|
elif len(dataset.headers) >= 3:
|
|
self.config.text_a_key, self.config.text_b_key, self.config.label_key = dataset.headers[0], \
|
|
dataset.headers[1], \
|
|
dataset.headers[-1]
|
|
else:
|
|
raise ValueError('Wrong dataset format')
|
|
report = {'text_a_key', 'text_b_key', 'label_key'}
|
|
report = dict((k, self.config[k]) for k in report)
|
|
report = [f'{k}={v}' for k, v in report.items() if v]
|
|
report = ', '.join(report)
|
|
logger.info(f'Guess [bold][blue]{report}[/blue][/bold] according to the headers of training dataset: '
|
|
f'[blue]{dataset}[/blue]')
|
|
self.build_vocabs(dataset, logger)
|
|
dataset.purge_cache()
|
|
# if self.config.transform:
|
|
# dataset.append_transform(self.config.transform)
|
|
dataset.append_transform(TransformerTextTokenizer(tokenizer=self.transformer_tokenizer,
|
|
text_a_key=self.config.text_a_key,
|
|
text_b_key=self.config.text_b_key,
|
|
max_seq_length=self.config.max_seq_length,
|
|
truncate_long_sequences=self.config.truncate_long_sequences,
|
|
output_key=''))
|
|
batch_sampler = None
|
|
if sorting and not isdebugging():
|
|
if dataset.cache and len(dataset) > 1000:
|
|
timer = CountdownTimer(len(dataset))
|
|
lens = []
|
|
for idx, sample in enumerate(dataset):
|
|
lens.append(len(sample['input_ids']))
|
|
timer.log('Pre-processing and caching dataset [blink][yellow]...[/yellow][/blink]',
|
|
ratio_percentage=None)
|
|
else:
|
|
lens = [len(sample['input_ids']) for sample in dataset]
|
|
batch_sampler = SortingSampler(lens, batch_size=batch_size, shuffle=shuffle,
|
|
batch_max_tokens=self.config.batch_max_tokens)
|
|
return PadSequenceDataLoader(dataset, batch_size, shuffle, batch_sampler=batch_sampler, device=device)
|
|
|
|
def build_dataset(self, data) -> TransformableDataset:
|
|
if isinstance(data, str):
|
|
dataset = TableDataset(data, cache=True)
|
|
elif isinstance(data, TableDataset):
|
|
dataset = data
|
|
elif isinstance(data, list):
|
|
dataset = TableDataset(data)
|
|
else:
|
|
raise ValueError(f'Unsupported data {data}')
|
|
return dataset
|
|
|
|
def predict(self, data: Union[str, List[str]], batch_size: int = None, **kwargs):
|
|
if not data:
|
|
return []
|
|
flat = isinstance(data, str) or isinstance(data, tuple)
|
|
if flat:
|
|
data = [data]
|
|
samples = []
|
|
for idx, d in enumerate(data):
|
|
sample = {IDX: idx}
|
|
if self.config.text_b_key:
|
|
sample[self.config.text_a_key] = d[0]
|
|
sample[self.config.text_b_key] = d[1]
|
|
else:
|
|
sample[self.config.text_a_key] = d
|
|
samples.append(sample)
|
|
dataloader = self.build_dataloader(samples,
|
|
sorting=False,
|
|
**merge_dict(self.config,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
device=self.device,
|
|
overwrite=True)
|
|
)
|
|
labels = [None] * len(data)
|
|
vocab = self.vocabs.label
|
|
for batch in dataloader:
|
|
logits = self.feed_batch(batch)
|
|
pred = logits.argmax(-1)
|
|
pred = pred.tolist()
|
|
for idx, tag in zip(batch[IDX], pred):
|
|
labels[idx] = vocab.idx_to_token[tag]
|
|
if flat:
|
|
return labels[0]
|
|
return labels
|
|
|
|
def fit(self, trn_data, dev_data, save_dir,
|
|
text_a_key=None,
|
|
text_b_key=None,
|
|
label_key=None,
|
|
transformer=None,
|
|
max_seq_len=512,
|
|
truncate_long_sequences=True,
|
|
# hidden_dropout_prob=0.0,
|
|
lr=5e-5,
|
|
transformer_lr=None,
|
|
adam_epsilon=1e-6,
|
|
weight_decay=0,
|
|
warmup_steps=0.1,
|
|
batch_size=32,
|
|
batch_max_tokens=None,
|
|
epochs=3,
|
|
logger=None,
|
|
# transform=None,
|
|
devices: Union[float, int, List[int]] = None,
|
|
**kwargs):
|
|
return super().fit(**merge_locals_kwargs(locals(), kwargs))
|
|
|
|
def build_vocabs(self, trn, logger, **kwargs):
|
|
self.vocabs.label = Vocab(pad_token=None, unk_token=None)
|
|
for each in trn:
|
|
pass
|
|
self.vocabs.lock()
|
|
self.vocabs.summary(logger)
|