text: # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  from __future__ import division
  from __future__ import print_function
  import numpy as np
  import string
  import paddle
  from paddle.nn import functional as F
  
  class BeamEntry:
      "information about one single beam at specific time-step"
      def __init__(self):
          self.prTotal = 0 # blank and non-blank
          self.prNonBlank = 0 # non-blank
          self.prBlank = 0 # blank
          self.prText = 1 # LM score
          self.lmApplied = False # flag if LM was already applied to this beam
          self.labeling = () # beam-labeling
  
  
  class BeamState:
      "information about the beams at specific time-step"
      def __init__(self):
          self.entries = {}
  
      def norm(self):
          "length-normalise LM score"
          for (k, _) in self.entries.items():
              labelingLen = len(self.entries[k].labeling)
              self.entries[k].prText = self.entries[k].prText ** (1.0 / (labelingLen if labelingLen else 1.0))
  
      def sort(self):
          "return beam-labelings, sorted by probability"
          beams = [v for (_, v) in self.entries.items()]
          sortedBeams = sorted(beams, reverse=True, key=lambda x: x.prTotal*x.prText)
          return [x.labeling for x in sortedBeams]
  
  
  def applyLM(parentBeam, childBeam, classes, lm):
      "calculate LM score of child beam by taking score from parent beam and bigram probability of last two chars"
      if lm and not childBeam.lmApplied:
          c1 = classes[parentBeam.labeling[-1] if parentBeam.labeling else classes.index(' ')] # first char
          c2 = classes[childBeam.labeling[-1]] # second char
          lmFactor = 0.01 # influence of language model
          bigramProb = lm.getCharBigram(c1, c2) ** lmFactor # probability of seeing first and second char next to each other
          childBeam.prText = parentBeam.prText * bigramProb # probability of char sequence
          childBeam.lmApplied = True # only apply LM once per beam entry
  
  
  def addBeam(beamState, labeling):
      "add beam if it does not yet exist"
      if labeling not in beamState.entries:
          beamState.entries[labeling] = BeamEntry()
  
  def testBeamSearch():
      "test decoder"
      classes = 'ab'
      mat = np.array([[0.4, 0, 0.6], [0.4, 0, 0.6]])
      print('Test beam search')
      expected = 'a'
      actual = ctcBeamSearch(mat, classes, None)
      print('Expected: "' + expected + '"')
      print('Actual: "' + actual + '"')
      print('OK' if expected == actual else 'ERROR')
  
  class BaseRecLabelDecode(object):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False):
          support_character_type = [
              'ch', 'en', 'EN_symbol', 'french', 'german', 'japan', 'korean',
              'it', 'xi', 'pu', 'ru', 'ar', 'ta', 'ug', 'fa', 'ur', 'rs', 'oc',
              'rsc', 'bg', 'uk', 'be', 'te', 'ka', 'chinese_cht', 'hi', 'mr',
              'ne', 'EN', 'latin', 'arabic', 'cyrillic', 'devanagari'
          ]
          assert character_type in support_character_type, "Only {} are supported now but get {}".format(
              support_character_type, character_type)
  
          self.beg_str = "sos"
          self.end_str = "eos"
  
          if character_type == "en":
              self.character_str = "0123456789abcdefghijklmnopqrstuvwxyz"
              dict_character = list(self.character_str)
          elif character_type == "EN_symbol":
              # same with ASTER setting (use 94 char).
              self.character_str = string.printable[:-6]
              dict_character = list(self.character_str)
          elif character_type in support_character_type:
              self.character_str = ""
              assert character_dict_path is not None, "character_dict_path should not be None when character_type is {}".format(
                  character_type)
              with open(character_dict_path, "rb") as fin:
                  lines = fin.readlines()
                  for line in lines:
                      line = line.decode('utf-8').strip("\n").strip("\r\n")
                      self.character_str += line
              if use_space_char:
                  self.character_str += " "
              dict_character = list(self.character_str)
  
          else:
              raise NotImplementedError
          self.character_type = character_type
          dict_character = self.add_special_char(dict_character)
          self.dict = {}
          for i, char in enumerate(dict_character):
              self.dict[char] = i
          self.character = dict_character
  
      def add_special_char(self, dict_character):
          return dict_character
  
      def beamdecode(self,text_prob,lm=None,beamWidth=5):
          "beam search as described by the paper of Hwang et al. and the paper of Graves et al."
          blankIdx = self.get_ignored_tokens()
          label = None
          char_list = []
          conf_list = []
          maxT, maxC = text_prob.shape
  
          # initialise beam state
          last = BeamState()
          labeling = ()
          last.entries[labeling] = BeamEntry()
          last.entries[labeling].prBlank = 1
          last.entries[labeling].prTotal = 1
  
          # go over all time-steps
          for t in range(maxT):
              curr = BeamState()
  
              # get beam-labelings of best beams
              bestLabelings = last.sort()[0:beamWidth]
  
              # go over best beams
              for labeling in bestLabelings:
  
                  # probability of paths ending with a non-blank
                  prNonBlank = 0
                  # in case of non-empty beam
                  if labeling:
                      # probability of paths with repeated last char at the end
                      prNonBlank = last.entries[labeling].prNonBlank * text_prob[t, labeling[-1]]
  
                  # probability of paths ending with a blank
                  prBlank = (last.entries[labeling].prTotal) * text_prob[t, blankIdx]
  
                  # add beam at current time-step if needed
                  addBeam(curr, labeling)
  
                  # fill in data
                  curr.entries[labeling].labeling = labeling
                  curr.entries[labeling].prNonBlank += prNonBlank
                  curr.entries[labeling].prBlank += prBlank
                  curr.entries[labeling].prTotal += prBlank + prNonBlank
                  curr.entries[labeling].prText = last.entries[labeling].prText # beam-labeling not changed, therefore also LM score unchanged from
                  curr.entries[labeling].lmApplied = True # LM already applied at previous time-step for this beam-labeling
  
                  # extend current beam-labeling
                  for c in range(maxC - 1):
                      # add new char to current beam-labeling
                      if text_prob[t,c] < 0.001:
                          continue
                      newLabeling = labeling + (c,)
  
                      # if new labeling contains duplicate char at the end, only consider paths ending with a blank
                      if labeling and labeling[-1] == c:
                          prNonBlank = text_prob[t, c] * last.entries[labeling].prBlank
                      else:
                          prNonBlank = text_prob[t, c] * last.entries[labeling].prTotal
  
                      # add beam at current time-step if needed
                      addBeam(curr, newLabeling)
  
                      # fill in data
                      curr.entries[newLabeling].labeling = newLabeling
                      curr.entries[newLabeling].prNonBlank += prNonBlank
                      curr.entries[newLabeling].prTotal += prNonBlank
  
                      # apply LM
                      applyLM(curr.entries[labeling], curr.entries[newLabeling], label, lm)
  
              # set new beam state
              last = curr
  
          # normalise LM scores according to beam-labeling-length
          last.norm()
  
          # sort by probability
          bestLabeling = last.sort()[0] # get most probable labeling
          print(bestLabeling)
          for i,idx in enumerate(bestLabeling):
              if idx not in blankIdx:
                  char_list.append(self.character[int(idx)])
                  conf_list.append(text_prob[i][idx])
          text = ''.join(char_list)
          # map labels to chars
          print(text,conf_list)
          return (text, np.mean(conf_list))
  
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          ignored_tokens = self.get_ignored_tokens()
          batch_size = len(text_index)
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] in ignored_tokens:
                      continue
                  if is_remove_duplicate:
                      # only for predict
                      if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                              batch_idx][idx]:
                          continue
                  char_list.append(self.character[int(text_index[batch_idx][
                      idx])])
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
              text = ''.join(char_list)
              result_list.append((text, np.mean(conf_list)))
          return result_list
  
      def get_ignored_tokens(self):
          return [0]  # for ctc blank
  
  
  class CTCLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False,
                   **kwargs):
          super(CTCLabelDecode, self).__init__(character_dict_path,
                                               character_type, use_space_char)
      def beamsearch(self, preds_idx, *args, **kwargs):
          batch_size = len(preds_idx)
          result_list = []
          for batch_idx in range(batch_size):
              result = self.beamdecode(preds_idx[batch_idx])
              result_list.append(result)
          return result_list
  
      def __call__(self, preds, label=None, *args, **kwargs):
          if isinstance(preds, paddle.Tensor):
              preds = preds.numpy()
          # preds(batch_size = 200, max_length = 80,char_num = 3977)
          # label(batch_size = 200, max_length = 30)
          flag = False
          if label is None and flag == True:
              text = self.beamsearch(preds)
          else:
              preds_idx = preds.argmax(axis=2)
              preds_prob = preds.max(axis=2)
              text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
          if label is None:
              return text
          label = self.decode(label)
          return text, label
  
      def add_special_char(self, dict_character):
          dict_character = ['blank'] + dict_character
          return dict_character
  
  
  class AttnLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False,
                   **kwargs):
          super(AttnLabelDecode, self).__init__(character_dict_path,
                                                character_type, use_space_char)
  
      def add_special_char(self, dict_character):
          self.beg_str = "sos"
          self.end_str = "eos"
          dict_character = dict_character
          dict_character = [self.beg_str] + dict_character + [self.end_str]
          return dict_character
  
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          ignored_tokens = self.get_ignored_tokens()
          [beg_idx, end_idx] = self.get_ignored_tokens()
          batch_size = len(text_index)
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] in ignored_tokens:
                      continue
                  if int(text_index[batch_idx][idx]) == int(end_idx):
                      break
                  if is_remove_duplicate:
                      # only for predict
                      if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                              batch_idx][idx]:
                          continue
                  char_list.append(self.character[int(text_index[batch_idx][
                      idx])])
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
              text = ''.join(char_list)
              result_list.append((text, np.mean(conf_list)))
          return result_list
  
      def __call__(self, preds, label=None, *args, **kwargs):
          """
          text = self.decode(text)
          if label is None:
              return text
          else:
              label = self.decode(label, is_remove_duplicate=False)
              return text, label
          """
          if isinstance(preds, paddle.Tensor):
              preds = preds.numpy()
  
          preds_idx = preds.argmax(axis=2)
          preds_prob = preds.max(axis=2)
          text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
          if label is None:
              return text
          label = self.decode(label, is_remove_duplicate=False)
          return text, label
  
      def get_ignored_tokens(self):
          beg_idx = self.get_beg_end_flag_idx("beg")
          end_idx = self.get_beg_end_flag_idx("end")
          return [beg_idx, end_idx]
  
      def get_beg_end_flag_idx(self, beg_or_end):
          if beg_or_end == "beg":
              idx = np.array(self.dict[self.beg_str])
          elif beg_or_end == "end":
              idx = np.array(self.dict[self.end_str])
          else:
              assert False, "unsupport type %s in get_beg_end_flag_idx" \
                            % beg_or_end
          return idx
  
  
  class SRNLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='en',
                   use_space_char=False,
                   **kwargs):
          super(SRNLabelDecode, self).__init__(character_dict_path,
                                               character_type, use_space_char)
  
      def __call__(self, preds, label=None, *args, **kwargs):
          pred = preds['predict']
          char_num = len(self.character_str) + 2
          if isinstance(pred, paddle.Tensor):
              pred = pred.numpy()
          pred = np.reshape(pred, [-1, char_num])
  
          preds_idx = np.argmax(pred, axis=1)
          preds_prob = np.max(pred, axis=1)
  
          preds_idx = np.reshape(preds_idx, [-1, 25])
  
          preds_prob = np.reshape(preds_prob, [-1, 25])
  
          text = self.decode(preds_idx, preds_prob)
  
          if label is None:
              text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
              return text
          label = self.decode(label)
          return text, label
  
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          ignored_tokens = self.get_ignored_tokens()
          batch_size = len(text_index)
  
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] in ignored_tokens:
                      continue
                  if is_remove_duplicate:
                      # only for predict
                      if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                              batch_idx][idx]:
                          continue
                  char_list.append(self.character[int(text_index[batch_idx][
                      idx])])
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
  
              text = ''.join(char_list)
              result_list.append((text, np.mean(conf_list)))
          return result_list
  
      def add_special_char(self, dict_character):
          dict_character = dict_character + [self.beg_str, self.end_str]
          return dict_character
  
      def get_ignored_tokens(self):
          beg_idx = self.get_beg_end_flag_idx("beg")
          end_idx = self.get_beg_end_flag_idx("end")
          return [beg_idx, end_idx]
  
      def get_beg_end_flag_idx(self, beg_or_end):
          if beg_or_end == "beg":
              idx = np.array(self.dict[self.beg_str])
          elif beg_or_end == "end":
              idx = np.array(self.dict[self.end_str])
          else:
              assert False, "unsupport type %s in get_beg_end_flag_idx" \
                            % beg_or_end
          return idx
  