text: import tqdm
  
  from .task_oriented_eda import eda
  from .types import MultiwozSampleType, SentenceType, MultiwozDatasetType
  from .util import AugmentationRecorder, iter_dialogues, Helper, choice, is_span_info_consistent_with_text, p_str
  from .tokenize_util import tokenize, convert_tokens_to_string, convert_sentence_to_tokens
  from .db.slot_value_replace import replace_slot_values_in_turn, MultiSourceDBLoader, assert_correct_turn
  
  
  class MultiwozEDA:
      def __init__(self, multiwoz: MultiwozDatasetType,
                   db_loader: MultiSourceDBLoader,
                   inform_intents=('inform',),
                   slot_value_replacement_probability=0.25,
                   alpha_sr=0.1, alpha_ri=0.1, alpha_rs=0.1, p_rd=0.1, num_aug=2):
          # attributes for slot value replacement
          self.db_loader = db_loader
          self.inform_intents = inform_intents
          self.slot_value_replacement_probability = slot_value_replacement_probability
  
          # attributes for EDA.
          self.eda_config = dict(alpha_sr=alpha_sr, alpha_ri=alpha_ri, alpha_rs=alpha_rs, p_rd=p_rd, num_aug=num_aug)
          self.multiwoz = multiwoz
          self.helper = Helper(multiwoz)
  
      def _get_excluding_indexes(self, words, span_info, dialog_act):
          return self.helper._get_excluding_indexes(words, span_info, dialog_act)
  
      def _augment_sentence_only(self, sentence: SentenceType, span_info, dialog_act):
          """don't change DA (span indexes may change)"""
          words = convert_sentence_to_tokens(sentence)
          excluding_indexes = self._get_excluding_indexes(words, span_info, dialog_act)
  
          for new_words, index_map in eda(words, **self.eda_config, excluding_indexes=excluding_indexes):
              new_span_info = []
              for x in span_info:
                  new_span_info.append([*x[:3], index_map[x[3]], index_map[x[4]]])
              yield convert_tokens_to_string(new_words), new_span_info, dialog_act
  
      def augment_sentence_only(self, sentence: SentenceType, span_info, dialog_act):
          return list(self._augment_sentence_only(sentence, span_info, dialog_act))
  
      def _augment_sample(self, sample: MultiwozSampleType, mode='usr') -> AugmentationRecorder:
          recorder = AugmentationRecorder(sample)
  
          for turn_index, turn in iter_dialogues(sample, mode=mode):
              if not is_span_info_consistent_with_text(turn['text'], turn['span_info']):
                  continue
              try:
                  assert_correct_turn(turn)
              except:
                  continue
              from copy import deepcopy
              orig_turn = deepcopy(turn)
              new_turn = replace_slot_values_in_turn(
                  turn,
                  self.db_loader,
                  p=self.slot_value_replacement_probability,
                  inform_intents=self.inform_intents
              )
              augmented = new_turn != turn
              turn = new_turn
  
              try:
                  text = turn['text']
                  span_info = turn['span_info']
                  dialog_act = turn['dialog_act']
                  tokens = tokenize(text)
                  augmented_sentence, augmented_span_info, augmented_dialog_act = choice(
                      self._augment_sentence_only(tokens, span_info, dialog_act)
                  )
              except (ValueError, IndexError):
                  pass
              else:
                  assert is_span_info_consistent_with_text(augmented_sentence, augmented_span_info), p_str(
                      [orig_turn, turn])
                  augmented = True
                  turn = {
                      'text': augmented_sentence,
                      'span_info': augmented_span_info,
                      'dialog_act': augmented_dialog_act,
                      **{k: v for k, v in turn.items() if k not in ('text', 'span_info', 'dialog_act')}
                  }
  
              if augmented:
                  recorder.add_augmented_dialog(turn_index, turn)
          return recorder
  
      def augment_sample(self, sample: MultiwozSampleType, mode='usr') -> MultiwozSampleType:
          return self._augment_sample(sample, mode=mode).get_augmented_sample()
  
      __call__ = augment_sample
  
      def augment_multiwoz_dataset(self, mode='usr', progress_bar=True):
          assert mode in ('usr', 'user', 'sys', 'all')
          res = {}
          if progress_bar:
              items = tqdm.tqdm(self.multiwoz.items(), total=len(self.multiwoz))
          else:
              items = self.multiwoz.items()
          for sample_id, sample in items:
              res[sample_id] = self.augment_sample(sample, mode=mode)
          return res
  