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

<commit_msg>`transformer_layers` means number of bottom layers<commit_after># -*- 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


{'fraction_non_alphanumeric': 0.06968001398845952, 'fraction_numerical': 0.005770239552369295, 'mean_word_length': 2.189015890716476, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 38, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '251755', 'n_tokens_mistral': 3046, 'n_tokens_neox': 2931, 'n_words': 579}