<commit_before># coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# 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.

# Lint as: python3
"""Utils for preprocessing the CFQ dataset."""

import collections
import json
import os
import string
from typing import Any, Dict, List, Tuple

from absl import logging

from tensorflow.compat.v1.io import gfile

Dataset = Dict[str, List[Tuple[str, str]]]


def load_json(path):
  logging.info(f'Reading json from {path} into memory...')
  with gfile.GFile(path) as f:
    data = json.load(f)
  logging.info(f'Successfully loaded json data from {path} into memory.')
  return data


def load_scan(path):
  """Read original scan task data and convert into CFQ-style json format."""
  logging.info(f'Reading SCAN tasks from {path}.')
  def parse(infile):
    for line in infile.read().split('\n'):
      if not line.startswith('IN: '):
        continue
      commands, actions = line[len('IN: '):].strip().split(' OUT: ', 1)
      yield {'questionPatternModEntities': commands,
             'sparqlPatternModEntities': actions}
  return list(parse(gfile.GFile(path)))


def load_dataset(path):
  """Load dataset from .json or SCAN task format."""
  if path[-5:] == '.json':
    return load_json(path)
  else:
    return load_scan(path)
<commit_msg>Reduce memory requirements of CFQ preprocessing to aid reproducibility.<commit_after># coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# 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.

# Lint as: python3
"""Utils for preprocessing the CFQ dataset."""

import collections
import json
import os
import re
import string
from typing import Any, Dict, List, Tuple

from absl import logging

from tensorflow.compat.v1.io import gfile

Dataset = Dict[str, List[Tuple[str, str]]]

_QUESTION_FIELD = 'questionPatternModEntities'
_QUERY_FIELD = 'sparqlPatternModEntities'


def _scrub_json(content):
  """Reduce JSON by filtering out only the fields of interest."""
  # Loading of json data with the standard Python library is very inefficient:
  # For the 4GB dataset file it requires more than 40GB of RAM and takes 3min.
  # There are more efficient libraries but in order to avoid additional
  # dependencies we use a simple (perhaps somewhat brittle) regexp to reduce
  # the content to only what is needed. This takes 1min to execute but
  # afterwards loading requires only 500MB or RAM and is done in 2s.
  regex = re.compile(
      r'("%s":\s*"[^"]*").*?("%s":\s*"[^"]*")' %
      (_QUESTION_FIELD, _QUERY_FIELD), re.DOTALL)
  return '[' + ','.join([
      '{' + m.group(1) + ',' + m.group(2) + '}' for m in regex.finditer(content)
  ]) + ']'


def load_json(path, scrub = False):
  logging.info(f'Reading json from {path} into memory...')
  with gfile.GFile(path) as f:
    if scrub:
      data = json.loads(_scrub_json(f.read()))
    else:
      data = json.load(f)
  logging.info(f'Successfully loaded json data from {path} into memory.')
  return data


def load_scan(path):
  """Read original scan task data and convert into CFQ-style json format."""
  logging.info(f'Reading SCAN tasks from {path}.')

  def parse(infile):
    for line in infile.read().split('\n'):
      if not line.startswith('IN: '):
        continue
      commands, actions = line[len('IN: '):].strip().split(' OUT: ', 1)
      yield {_QUESTION_FIELD: commands, _QUERY_FIELD: actions}

  return list(parse(gfile.GFile(path)))


def load_dataset(path):
  """Load dataset from .json or SCAN task format."""
  if path[-5:] == '.json':
    return load_json(path, scrub=True)
  else:
    return load_scan(path)

{'fraction_non_alphanumeric': 0.11331877729257642, 'fraction_numerical': 0.008078602620087336, 'mean_word_length': 3.9524324324324325, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5785327', 'n_tokens_mistral': 1398, 'n_tokens_neox': 1359, 'n_words': 532}