text: # -*- 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,
                                           self.config.max_seq_len, 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)
  