Parsing the bAbi dataset (train and test).

  Args:
    tmp_dir: temp directory to download and extract the dataset
    babi_task_id: babi task id
    subset: babi subset
    dataset_split: dataset split (train or eval)
    joint_training: if training the model on all tasks.

  Returns:
     babi_instances: set of training examples, each a dict containing a story,
     a question and an answer.
     babi_lines: all the texts in the data separated based on their
     appearance in the stories, questions, or answers.

def _babi_parser(tmp_dir,
                 babi_task_id,
                 subset,
                 dataset_split,
                 joint_training=True):
  """Parsing the bAbi dataset (train and test).

  Args:
    tmp_dir: temp directory to download and extract the dataset
    babi_task_id: babi task id
    subset: babi subset
    dataset_split: dataset split (train or eval)
    joint_training: if training the model on all tasks.

  Returns:
     babi_instances: set of training examples, each a dict containing a story,
     a question and an answer.
     babi_lines: all the texts in the data separated based on their
     appearance in the stories, questions, or answers.
  """

  def _data_file(mode, task_id):
    """Generates the path to the data file for the given mode(train/test).

    Args:
      mode: either train or test for bAbi dataset
      task_id: babi task id

    Returns:
      data file path
    """
    file_name = (_TASKS[task_id] + "_{}.txt")
    return os.path.join(_DIR_NAME, subset, file_name.format(mode))

  def _all_task_raw_data_generator(tmp_dir, data_file, dataset_split):
    """Prepares raw data for all tasks to gether..

    Args:
      tmp_dir: temp directory
      data_file: data file
      dataset_split: dataset split
    """

    tf.logging.info("Preparing dataset of all task together")
    globe_name = ("*_{}.txt")
    mode_name = "test"
    if dataset_split == problem.DatasetSplit.TRAIN:
      mode_name = "train"
    files_name = os.path.join(
        tmp_dir, _DIR_NAME, subset,
        globe_name.format(mode_name))
    with tf.gfile.GFile(data_file, "wb") as outfile:
      for filename in tf.gfile.Glob(files_name):
        if filename == data_file:
          # don"t want to copy the output into the output
          continue
        with tf.gfile.GFile(filename, "rb") as readfile:
          shutil.copyfileobj(readfile, outfile)

  def _parse_answer(answer):
    if (joint_training or babi_task_id in ["qa8", "qa19", "qa0"
                                          ]):  # "lists-sets" or "path finding"
      return "".join([d for d in answer.split(",")])  # as a single token!
    else:
      return answer

  if dataset_split == problem.DatasetSplit.TRAIN:
    babi_train_task_id = "qa0" if joint_training else babi_task_id
    data_file = os.path.join(tmp_dir, _data_file("train", babi_train_task_id))
  else:
    data_file = os.path.join(tmp_dir, _data_file("test", babi_task_id))

  if ((babi_task_id == "qa0" or joint_training) and
      not tf.gfile.Exists(os.path.join(tmp_dir, data_file))):
    _all_task_raw_data_generator(tmp_dir, data_file, dataset_split)

  tf.logging.info("Parsing %s into training/testing instances...", data_file)

  babi_instances = []
  with tf.gfile.GFile(data_file, mode="r") as f:
    story = []
    for line in f:
      line_num, line = line.strip().split(" ", 1)
      if int(line_num) == 1:
        story = []
      if "\t" in line:
        question, answer, _ = line.split("\t")
        question = _normalize_string(question)
        substories = [s for s in story if s]
        answer = _parse_answer(answer)
        instance = {
            FeatureNames.STORY: substories,
            FeatureNames.QUESTION: question,
            FeatureNames.ANSWER: answer
        }
        babi_instances.append(instance)

        story.append("")
      else:
        story.append(_normalize_string(line))

  return babi_instances