text: # -*- coding:utf-8 -*-
  # Author: hankcs
  # Date: 2020-12-12 16:05
  import logging
  from typing import Dict, Any, List, Union, Iterable, Callable
  
  import torch
  from stog.data.dataset_readers.amr_parsing.amr import AMRGraph
  from stog.data.dataset_readers.amr_parsing.node_utils import NodeUtilities
  from stog.data.dataset_readers.amr_parsing.postprocess.node_restore import NodeRestore
  from torch.utils.data import DataLoader
  
  from hanlp_common.constant import CLS
  from hanlp.common.dataset import PrefetchDataLoader, SamplerBuilder
  from hanlp.common.transform import VocabDict
  from hanlp.components.amr.amr_parser.graph_amr_decoder import GraphAbstractMeaningRepresentationDecoder
  from hanlp.components.amr.amr_parser.graph_parser import GraphAbstractMeaningRepresentationParser
  from hanlp.components.amr.amr_parser.postprocess import PostProcessor
  from hanlp.components.amr.amr_parser.work import parse_batch
  from hanlp.components.mtl.tasks import Task
  from hanlp.datasets.parsing.amr import batchify, get_concepts
  from hanlp.layers.scalar_mix import ScalarMixWithDropoutBuilder
  from hanlp.metrics.amr.smatch_eval import SmatchScores, get_amr_utils
  from hanlp.metrics.f1 import F1_
  from hanlp.metrics.metric import Metric
  from hanlp.metrics.mtl import MetricDict
  from hanlp.utils.io_util import get_resource
  from hanlp_common.util import merge_list_of_dict, merge_locals_kwargs
  
  
  class GraphAbstractMeaningRepresentationParsing(Task, GraphAbstractMeaningRepresentationParser):
  
      def __init__(self,
                   trn: str = None,
                   dev: str = None,
                   tst: str = None,
                   sampler_builder: SamplerBuilder = None,
                   dependencies: str = None,
                   scalar_mix: ScalarMixWithDropoutBuilder = None,
                   use_raw_hidden_states=False,
                   lr=1e-3,
                   separate_optimizer=False,
                   cls_is_bos=True,
                   sep_is_eos=False,
                   char2concept_dim=128,
                   cnn_filters=((3, 256),),
                   concept_char_dim=32,
                   concept_dim=300,
                   dropout=0.2,
                   embed_dim=512,
                   eval_every=20,
                   ff_embed_dim=1024,
                   graph_layers=2,
                   inference_layers=4,
                   num_heads=8,
                   rel_dim=100,
                   snt_layers=4,
                   unk_rate=0.33,
                   vocab_min_freq=5,
                   beam_size=8,
                   alpha=0.6,
                   max_time_step=100,
                   amr_version='2.0',
                   **kwargs) -> None:
          super().__init__(**merge_locals_kwargs(locals(), kwargs))
          self.vocabs = VocabDict()
          utils_dir = get_resource(get_amr_utils(amr_version))
          self.sense_restore = NodeRestore(NodeUtilities.from_json(utils_dir))
  
      def build_dataloader(self,
                           data,
                           transform: Callable = None,
                           training=False,
                           device=None,
                           logger: logging.Logger = None,
                           cache=False,
                           gradient_accumulation=1,
                           **kwargs) -> DataLoader:
          if isinstance(data, list):
              data = GraphAbstractMeaningRepresentationParser.build_samples(self, data)
          dataset, lens = GraphAbstractMeaningRepresentationParser.build_dataset(self, data, logger=logger,
                                                                                 transform=transform, training=training)
          if self.vocabs.mutable:
              GraphAbstractMeaningRepresentationParser.build_vocabs(self, dataset, logger)
          dataloader = PrefetchDataLoader(
              DataLoader(batch_sampler=self.sampler_builder.build(lens, shuffle=training,
                                                                  gradient_accumulation=gradient_accumulation),
                         dataset=dataset,
                         collate_fn=merge_list_of_dict,
                         num_workers=0), batchify=self.build_batchify(device, training),
              prefetch=None)
          return dataloader
  
      def compute_loss(self,
                       batch: Dict[str, Any],
                       output: Union[torch.Tensor, Dict[str, torch.Tensor], Iterable[torch.Tensor], Any],
                       criterion) -> Union[torch.FloatTensor, Dict[str, torch.FloatTensor]]:
          concept_loss, arc_loss, rel_loss, graph_arc_loss = output
          concept_loss, concept_correct, concept_total = concept_loss
          rel_loss, rel_correct, rel_total = rel_loss
          loss = concept_loss + arc_loss + rel_loss
          return loss
  
      def decode_output(self,
                        output: Union[torch.Tensor, Dict[str, torch.Tensor], Iterable[torch.Tensor], Any],
                        mask: torch.BoolTensor,
                        batch: Dict[str, Any],
                        decoder: torch.nn.Module, **kwargs) -> Union[Dict[str, Any], Any]:
          return output
  
      def update_metrics(self,
                         batch: Dict[str, Any],
                         output: Union[torch.Tensor, Dict[str, torch.Tensor], Iterable[torch.Tensor], Any],
                         prediction: Dict[str, Any],
                         metric: Union[MetricDict, Metric]):
          pass
  
      def build_model(self, encoder_size, training=True, **kwargs) -> torch.nn.Module:
          return GraphAbstractMeaningRepresentationDecoder(vocabs=self.vocabs, encoder_size=encoder_size, **self.config)
  
      def build_metric(self, **kwargs):
          return SmatchScores({'Smatch': F1_(0, 0, 0)})
  
      def input_is_flat(self, data) -> bool:
          return GraphAbstractMeaningRepresentationParser.input_is_flat(self, data)
  
      def prediction_to_result(self, prediction: Dict[str, Any], batch: Dict[str, Any]) -> List:
          pp = PostProcessor(self.vocabs['rel'])
          for concept, relation, score in zip(prediction['concept'], prediction['relation'], prediction['score']):
              amr = pp.to_amr(concept, relation)
              amr_graph = AMRGraph(amr)
              self.sense_restore.restore_graph(amr_graph)
              yield amr_graph
  
      def evaluate_dataloader(self,
                              data: DataLoader,
                              criterion: Callable,
                              metric=None,
                              output=False,
                              input=None,
                              decoder=None,
                              h=None,
                              split=None,
                              **kwargs):
          # noinspection PyTypeChecker
          GraphAbstractMeaningRepresentationParser.evaluate_dataloader(self, data, logger=None, metric=metric,
                                                                       input=input, model=decoder, h=lambda x: h(x)[0],
                                                                       use_fast=True)
  
      def feed_batch(self,
                     h: torch.FloatTensor,
                     batch: Dict[str, torch.Tensor],
                     mask: torch.BoolTensor,
                     decoder: torch.nn.Module):
          if decoder.training:
              return super().feed_batch(h, batch, mask, decoder)
          beam_size = self.config.get('beam_size', 8)
          alpha = self.config.get('alpha', 0.6)
          max_time_step = self.config.get('max_time_step', 100)
          res = parse_batch(decoder, batch, beam_size, alpha, max_time_step, h=h)
          return res
  
      def transform_batch(self, batch: Dict[str, Any], results: Dict[str, Any] = None, cls_is_bos=False,
                          sep_is_eos=False) -> Dict[str, Any]:
          batch = super().transform_batch(batch, results, cls_is_bos, sep_is_eos)
          batch['lemma'] = [[CLS] + x for x in results['lem']]
          copy_seq = merge_list_of_dict(
              [get_concepts({'token': t[1:], 'lemma': l[1:]}, self.vocabs.predictable_concept) for t, l in
               zip(batch['token'], batch['lemma'])])
          copy_seq.pop('token')
          copy_seq.pop('lemma')
          batch.update(copy_seq)
          ret = batchify(batch, self.vocabs, device=batch['token_input_ids'].device)
          return ret
  