def tokenize_kb(kb_json):
  def tokenize_one_flight(flight_json):
def process_kb(raw_kb, word_map, stream=False):
def flatten_json(obj):
def apply_word_map(flattened_sent, word_map):
def format_tag(name, val):
def get_full_intent(intent_json):
def add_dash_to_name(original_name):
def tokenize_intent(intent_json):
def tokenize_action(action_json, first_name_cat, last_name_cat, flight_cat,
                    state_cat):
def process_main_data(raw_data,
                      sent_tok,
                      word_tok,
                      word_map,
                      input_type,
                      stream=False,
                      self_play_start_turn=None):
  def process_dialogue(dialogue):
    def do_tokenize(last, turn):
  def get_dialogue_boundary(start_token, flat_dialogue):
    def get_end_token(start, set_of_end_tokens, splitted_dialogues):
    def get_next_start_token(end_position, start_token, splitted_dialogues):
  def serialize_boundary(start, end):
def is_ascii(s, keep_non_ascii):
def write_vocabulary(output_file, output_all_vocab_file, word_frequency,
                     frequency_cutoff, keep_non_ascii):
def write_cat(files, cats):
def write_data(data, output_file_data, output_file_kb, alt_infer=False):
def write_completion(data, output_file_data_src, output_file_data_tar,
                     output_file_kb):
def write_self_play(data, output_file_data, output_file_kb):
def word_tokenize(tokens): def process_kb(raw_kb, word_map, stream=False):
  """main entry to process kb."""
  processed_data = []
  d = raw_kb
  if not stream:
    d = tqdm(raw_kb, desc='process kb')
  for kb_object in d:
    # the database will be flattened into a single sequence of tokens.
    flattened = tokenize_kb(kb_object)
    processed_data.append(flattened)
    word_map = apply_word_map(flattened, word_map)
  return processed_data, word_map def process_main_data(raw_data,
                      sent_tok,
                      word_tok,
                      word_map,
                      input_type,
                      stream=False,
                      self_play_start_turn=None):
  """This function processes the main data."""

  def process_dialogue(dialogue):
    """This function processes dialogues."""

    def do_tokenize(last, turn):
      """This function tokenizes the diloagues of a specific turn."""
      if turn.startswith('customer: '):
        # agent:
        turn = turn[10:].strip()
        sot = start_of_turn1
        eot = start_of_turn2
      else:  # agent:
        turn = turn[7:].strip()
        sot = start_of_turn2
        eot = start_of_turn1
      sentences = sent_tok(turn)
      tokenized_sents = []
      for s in sentences:
        words = word_tok(s)
        tokenized_sents.append(' '.join(words))
      flat_content = ' '.join(tokenized_sents)
      if last:
        return sot + ' ' + flat_content + ' ' + end_of_dialogue + ' ' + eot
      else:
        return sot + ' ' + flat_content

    tokenized_dialogue = []
    max_turn_length = 0
    for i, turn in enumerate(dialogue):
      tokenized_turn = do_tokenize(i == len(dialogue) - 1, turn)
      max_turn_length = max(max_turn_length, len(tokenized_turn.split(' ')))
      tokenized_dialogue.append(tokenized_turn)

    return ' '.join(tokenized_dialogue), max_turn_length

  def get_dialogue_boundary(start_token, flat_dialogue):
    """This function gets the boundary array of the dialogues."""

    def get_end_token(start, set_of_end_tokens, splitted_dialogues):
      for i in range(start, len(splitted_dialogues)):
        if splitted_dialogues[i] in set_of_end_tokens:
          return i
      assert False, 'end token not found : ' + ' '.join(
          flat_dialogue) + 'start=' + str(start) + '/' + str(
              len(splitted_dialogues))

    def get_next_start_token(end_position, start_token, splitted_dialogues):
      for i in range(end_position, len(splitted_dialogues)):
        if splitted_dialogues[i] == start_token:
          return i
      return len(splitted_dialogues)

    set_of_end_tokens = set([
        start_of_turn1, start_of_turn2
    ])  # taking out end_of_dialogue token because of dynamic rnn decoder
    # set_of_end_tokens=set([start_of_turn1,start_of_turn2,end_of_dialogue])
    splitted_dialogue = flat_dialogue.split(' ')
    i = get_next_start_token(0, start_token, splitted_dialogue)
    all_starts = []
    all_ends = []
    while i < len(splitted_dialogue
                 ) - 1:  # we don't find the end token for the last turn change.
      end_position = get_end_token(i + 1, set_of_end_tokens, splitted_dialogue)
      err_msg = 'start token appeared twice: ' + flat_dialogue
      assert splitted_dialogue[end_position] != start_token, err_msg
      # all_res.append((i,end_position))
      all_starts.append(i)
      all_ends.append(end_position)
      i = get_next_start_token(i + 1, start_token, splitted_dialogue)
    return (all_starts, all_ends)

  def serialize_boundary(start, end):
    all_res = []
    for s in start:
      all_res.append(str(s))
    for e in end:
      all_res.append(str(e))
    return ' '.join(all_res)

  intents = []
  actions = []
  expected_actions = []
  dialogues = []
  boundaries1 = []
  boundaries2 = []
  lengths = []
  max_diag_length = 0
  max_turn1 = 0
  first_name_cat = set([])
  last_name_cat = set([])
  flight_cat = set([])
  state_cat = set([])

  d = raw_data
  if not stream:
    d = tqdm(raw_data, desc='process raw data')

  for loaded_json in d:
    # loaded_json = json.loads(delete_non_ascii(line))
    # input_type
    if 'intent' in loaded_json:
      intent = loaded_json['intent']
      processed_intent = tokenize_intent(get_full_intent(intent))
      intents.append(processed_intent)
      word_map = apply_word_map(processed_intent, word_map)
    else:
      intents.append('')
    if 'action' in loaded_json and loaded_json['action']:
      action = loaded_json['action']
      processed_action = tokenize_action(
          action,
          first_name_cat,
          last_name_cat,
          flight_cat,
          state_cat,
      )
      actions.append(processed_action)
      word_map = apply_word_map(processed_action, word_map)
    else:
      actions.append('')
    if 'expected_action' in loaded_json:
      expected_action = loaded_json['expected_action']
      # print "processed_intent", processed_intent
      processed_expected_action = tokenize_action(
          expected_action,
          first_name_cat,
          last_name_cat,
          flight_cat,
          state_cat,
      )
      # print "processed_action", processed_action
      expected_actions.append(processed_expected_action)

    # NB for word map updates above:
    # word map should contain everything in the supervised eval and training
    # however, it should not contain expected_action since this is for self-play
    # it contains tokens that has _, on flights.

    # process dialogue only when input_type is dialogue
    if input_type == 'dialogue':
      dialogue = loaded_json['dialogue']
      processed_dialogue, max_diag_len_this = process_dialogue(dialogue)
      max_diag_length = max(max_diag_length, max_diag_len_this)
      t1_boundaries = get_dialogue_boundary(start_of_turn1, processed_dialogue)
      t2_boundaries = get_dialogue_boundary(start_of_turn2, processed_dialogue)
      max_turn1 = max(max_turn1, len(t1_boundaries[0]) + len(t2_boundaries[0]))
      boundaries1.append(serialize_boundary(*t1_boundaries))
      boundaries2.append(serialize_boundary(
          *t2_boundaries))  # this is used only for inference files generation
      if processed_dialogue == '' and \
          self_play_start_turn in ['agent', 'customer']:
        if self_play_start_turn == 'agent':
          processed_dialogue = '<t2>'
        else:
          processed_dialogue = '<t1>'

      dialogues.append(processed_dialogue)
      length = len(processed_dialogue.split(' '))
      lengths.append(length)
      word_map = apply_word_map(processed_dialogue, word_map)

  if input_type == 'dialogue' and not stream:  #  output stats only when input is dialogue
    min_length, mean_length, max_length = np.min(lengths), np.mean(
        lengths), np.max(lengths)
    print(
        ('min_len: ${0}, mean_len: {1}, max_len: {2}'
         'max_sent_len: {3}, max_turn: {4}').format(min_length, mean_length,
                                                    max_length, max_diag_length,
                                                    max_turn1))

  # return all the processed data. Some of them will be empty arrays when in
  # context mode.
  all_cat = [first_name_cat, last_name_cat, flight_cat, state_cat]
  return intents, actions, expected_actions, dialogues, word_map, boundaries1, boundaries2, all_cat def word_tokenize(tokens):
  return [
      token.replace("''", '"').replace('``', '"')
      for token in nltk.word_tokenize(tokens)
  ] def write_cat(files, cats):
  for file, category in zip(files, cats):
    with gfile.Open(file, 'w') as f:
      for cat in category:
        f.write(str(cat) + '\n') def write_completion(data, output_file_data_src, output_file_data_tar,
                     output_file_kb):
  """This function write both kb and main data into the files."""
  f_data_src = gfile.Open(output_file_data_src, 'w')
  f_data_tar = gfile.Open(output_file_data_tar, 'w')
  f_kb = gfile.Open(output_file_kb, 'w')
  for entry in data:
    bd1 = entry['boundaries1'].split(' ')
    bd2 = entry['boundaries2'].split(' ')
    start = bd1[0:len(bd1) // 2] + bd2[0:len(bd2) // 2]
    end = bd1[len(bd1) // 2:] + bd2[len(bd2) // 2:]
    # random_turn = random.randint(0, len(start) - 1)
    for random_turn in range(len(start)):
      f_kb.write(flatten_json(entry['kb']) + '\n')
      # print len(start),len(end),len(bd),random_turn
      turn_start = int(start[random_turn])
      turn_end = int(end[random_turn])
      dialogue_split = entry['dialogue'].split(' ')
      # print turn_start,turn_end
      dialogue_src = dialogue_split[0:turn_start + 1]
      dialogue_tar = dialogue_split[turn_start + 1:turn_end + 1]
      src_arr = [entry['intent'], ' '.join(dialogue_src)]
      f_data_src.write('|'.join(src_arr) + '\n')
      tar_arr = [entry['action'], ' '.join(dialogue_tar)]
      f_data_tar.write('|'.join(tar_arr) + '\n')

  f_data_src.close()
  f_data_tar.close()
  f_kb.close() def write_data(data, output_file_data, output_file_kb, alt_infer=False):
  """This function writes data into a text file."""
  f_data = gfile.Open(output_file_data, 'w')
  f_kb = gfile.Open(output_file_kb, 'w')
  for entry in data:
    f_kb.write(flatten_json(entry['kb']) + '\n')
    new_arr = []
    if alt_infer:
      new_arr = [entry['intent'], entry['dialogue'].replace('<eod> ', '')]
    else:
      new_arr = [
          entry['intent'], entry['action'], entry['dialogue'],
          entry['boundaries1']
      ]
      # only boundary1 is used but not 2 because it's not necessary.
    f_data.write('|'.join(new_arr) + '\n')
  f_data.close()
  f_kb.close() def write_self_play(data, output_file_data, output_file_kb):
  f_data = gfile.Open(output_file_data, 'w')
  f_kb = gfile.Open(output_file_kb, 'w')
  for entry in data:
    f_kb.write(flatten_json(entry['kb']) + '\n')
    new_arr = [entry['intent'],
               entry['expected_action']]  # intent and action are both needed
    f_data.write('|'.join(new_arr) + '\n')
  f_data.close()
  f_kb.close() def write_vocabulary(output_file, output_all_vocab_file, word_frequency,
                     frequency_cutoff, keep_non_ascii):
  special_chars = [unk_token, start_of_turn1, start_of_turn2, end_of_dialogue]
  new_word_frequency = set([])
  # if output_filei not None, we write to file, otherwise we don't.
  if output_file:
    f = gfile.Open(output_file, 'w')
  else:
    f = None

  # first one should always be unktoken
  for special_char in special_chars:
    if f:
      f.write(special_char + '\n')
  for key in word_frequency:
    # We write to the vocabulary only when the key is not empty.
    # Otherwise tensorflow will complain.
    if word_frequency[key] >= frequency_cutoff and (
        key not in special_chars) and is_ascii(key, keep_non_ascii) and key:
      if f:
        f.write(key + '\n')
      new_word_frequency.add(key)

  if f:
    f.close()
  # all vocab
  if output_all_vocab_file:
    with gfile.Open(output_all_vocab_file, 'w') as f2:
      f2.write(str(word_frequency))
  return new_word_frequency def standardize_and_drop(data_file,
                         kb_file,
                         drop_incorrect=True,
                         verbose=False):
  """ this function filter incorrect samples and standardize them

   the same time.
  """
  loaded_data, loaded_kb = load_and_drop(data_file, kb_file, drop_incorrect,
                                         verbose)
  for data_obj in tqdm(loaded_data, desc='standardizing data'):
    org_time = data_obj['timestamps'] if 'timestamps' in data_obj else None
    org_diag = data_obj['dialogue'] if 'dialogue' in data_obj else None
    if org_diag:
      new_diag, new_time = standardize_message(org_diag, org_time)
      data_obj['dialogue'] = new_diag
      if new_time:
        data_obj['timestamps'] = new_time
        assert len(data_obj['dialogue']) == len(data_obj['timestamps'])
  return loaded_data, loaded_kb def load_and_drop(data_file, kb_file, drop_incorrect=True, verbose=False):
  """ this function filter incorrect samples without standardization."""
  fin_data = gfile.GFile(data_file)
  fin_kb = gfile.GFile(kb_file)
  total_in_file = 0
  loaded_data = []
  loaded_kb = []
  for line1 in tqdm(fin_data, desc='loading data'):
    if len(line1.strip()) < 10:
      continue
    line2 = fin_kb.readline()
    if len(line2.strip()) < 10:
      continue
    line1 = delete_non_ascii(line1)
    line2 = delete_non_ascii(line2)

    data_obj = json.loads(line1)
    kb_obj = json.loads(line2)
    if (not drop_incorrect) or (
        'correct_sample' not in data_obj) or data_obj['correct_sample']:
      loaded_data.append(data_obj)
      loaded_kb.append(kb_obj)
    total_in_file += 1

  if verbose:
    print(('loaded: ', len(loaded_data), '/', total_in_file, '=',
           len(loaded_data) * 1.0 / total_in_file))
  return loaded_data, loaded_kb def load_and_drop_stream(data_file,
                         kb_file,
                         drop_incorrect=True,
                         verbose=False):
  """ this function filter incorrect samples without standardization."""
  if verbose:
    print('loading stream')
  fin_data = gfile.GFile(data_file)
  if gfile.exists(kb_file):
    fin_kb = gfile.GFile(kb_file)
  else:
    fin_kb = None
  if verbose:
    print('gfile loaded: ', fin_data)
  for line1 in fin_data:
    if verbose:
      print(line1)
    if len(line1.strip()) < 10:
      continue
    line1 = delete_non_ascii(line1)
    data_obj = json.loads(line1)

    if fin_kb:
      line2 = fin_kb.readline()
      if len(line2.strip()) < 10:
        continue
      line2 = delete_non_ascii(line2)
      kb_obj = json.loads(line2)
    else:
      kb_obj = None
    if (not drop_incorrect) or (
        'correct_sample' not in data_obj) or data_obj['correct_sample']:
      yield data_obj, kb_obj
import argparse
import os
import nltk
import tensorflow.compat.v1 as tf
import sys
from tensorflow.compat.v1 import gfile
from tqdm import tqdm
from airdialogue.prepro.tokenize_lib import list_of_action_tokens_except_name
from airdialogue.prepro.tokenize_lib import process_kb
from airdialogue.prepro.tokenize_lib import process_main_data
from airdialogue.prepro.tokenize_lib import word_tokenize
from airdialogue.prepro.tokenize_lib import write_cat
from airdialogue.prepro.tokenize_lib import write_completion
from airdialogue.prepro.tokenize_lib import write_data
from airdialogue.prepro.tokenize_lib import write_self_play
from airdialogue.prepro.tokenize_lib import write_vocabulary
from airdialogue.prepro.standardize_data_lib import standardize_and_drop
from airdialogue.prepro.standardize_data_lib import load_and_drop, load_and_drop_stream

def generate_entry(intents, actions, expected_actions, dialogues, kbs,
                   boundaries1, boundaries2, i):
  """generate entry based on the available of data."""
  # required fields
  entry = {'intent': intents[i], 'action': actions[i], 'kb': kbs[i]}
  # optional fields depending on the input type
  if dialogues:
    entry['dialogue'] = dialogues[i]
  if boundaries1:
    entry['boundaries1'] = boundaries1[i]
  if boundaries2:
    entry['boundaries2'] = boundaries2[i]
  if expected_actions:
    entry['expected_action'] = expected_actions[i]
  return entry


def reorganize_data(intents, actions, expected_actions, dialogues, kbs,
                    boundaries1, boundaries2):
  """reorganize data into compact form.

  This is mostly because we want to reuse
  code for the random split version.
  """
  data = []
  for i in range(len(intents)):
    entry = generate_entry(intents, actions, expected_actions, dialogues, kbs,
                           boundaries1, boundaries2, i)
    data.append(entry)
  return data


def process_job_type(job_type_str, input_type):
  """process job type string."""
  all_jobs = []
  job_type_arr = job_type_str.split('|')
  assert len(job_type_arr) == 5
  for i in range(5):
    if int(job_type_arr[i]) == 1:
      all_jobs.append('train|eval|infer|sp-train|sp-eval'.split('|')[i])
  assert input_type in ['dialogue', 'context']
  # if input type is context, we can't do train/eval/infer
  if input_type == 'context':
    assert 'train' not in all_jobs
    assert 'eval' not in all_jobs
    assert 'infer' not in all_jobs
  return all_jobs


def load_data_from_jsons(FLAGS, input_data_file, input_kb_file, output_vab,
                         output_all_vab, gen_cat, cat_files):
  vocal_map = {}
  sent_tokenize = nltk.sent_tokenize

  raw_data, raw_kb = load_and_drop(
      input_data_file,
      input_kb_file,
      drop_incorrect=not FLAGS.keep_incorrect,
      verbose=FLAGS.verbose)
  # has to be there no matter what
  if FLAGS.verbose:
    print('processing kb')
  processed_kb, vocal_map = process_kb(raw_kb, vocal_map)
  # if dialogue, everything will be there.
  # if context, only intents, actions, vocal_map will be there
  if FLAGS.verbose:
    print('processing data')
  result = process_main_data(
      raw_data,
      sent_tokenize,
      word_tokenize,
      vocal_map,
      input_type=FLAGS.input_type)
  intents, actions, expected_actions, dialogues, vocal_map, boundaries1, boundaries2, cats = result
  frequency_cutoff = FLAGS.word_cutoff
  # 3 is the number of special tokens
  if FLAGS.verbose:
    print('vocabulary before cutoff', len(vocal_map) + 3)
  vocal_map = write_vocabulary(output_vab, output_all_vab, vocal_map,
                               frequency_cutoff, FLAGS.keep_non_ascii)
  if gen_cat:
    if FLAGS.verbose:
      print('writing category')
    write_cat(cat_files, cats)

  if FLAGS.verbose:
    print(
        'frequency_cutoff= {0}, vocabulary after cutoff'.format(
            frequency_cutoff), len(vocal_map))
  data = reorganize_data(intents, actions, expected_actions, dialogues,
                         processed_kb, boundaries1, boundaries2)
  return data


def load_data_from_jsons_stream(FLAGS,
                                input_data_file,
                                input_kb_file,
                                output_vab,
                                output_all_vab,
                                gen_cat,
                                cat_files,
                                self_play_start_turn=None):
  vocal_map = {}
  sent_tokenize = nltk.sent_tokenize

  for raw_data, raw_kb in tqdm(
      load_and_drop_stream(
          input_data_file,
          input_kb_file,
          drop_incorrect=not FLAGS.keep_incorrect,
          verbose=FLAGS.verbose),
      desc='processing stream'):
    # has to be there no matter what
    if raw_kb is not None:
      processed_kb, vocal_map = process_kb([raw_kb], vocal_map, stream=True)
    else:
      processed_kb = [['no_res']]
    # if dialogue, everything will be there.
    # if context, only intents, actions, vocal_map will be there

    result = process_main_data([raw_data],