class CuBertDataset(plur_dataset.PlurDataset):
class JsonExtractor(beam.DoFn):
  _CUBERT_BUCKET = 'cubert'
  _VOCABULARY = '20200621_Python/github_python_minus_ethpy150open_deduplicated_vocabulary.txt'
  def __init__(
      self,
      stage_1_dir,
      configuration: plur_dataset.Configuration = plur_dataset.Configuration(),
      transformation_funcs=(),
      filter_funcs=(),
      user_defined_split_range=(),
      num_shards=1000,
      seed=0,
      deduplicate=False):
  def dataset_name(self) -> str:
  def dataset_description(self) -> str:
  def folder_path(self) -> str:
  def make_dirs(self, directory: str) -> None:
  def fetch_file(self, blob: storage.Blob, local_path: str) -> None:
  def download_dataset(self):
  def _initialize_tokenizers(self) -> None:
  def get_all_raw_data_paths(self):
  def raw_data_paths_to_raw_data_do_fn(self):
  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
  def raw_data_to_graph_to_output_example(
      self, raw_data: Mapping[str, Any]) -> Mapping[str, Any]:
  def single_classification_data_dictionary_to_graph_to_output_example(
      self, data: Mapping[str, Any], classes: AbstractSet[str],
      max_graph_size: int,
      split: str) -> GraphToOutputExample:
  def __init__(self, random_split_fn, use_random_split: bool) -> None:
  def _get_split(self, file_path):
  def process(self, file_path: str) -> Iterable[Mapping[str, Any]]: class Configuration():
class PlurDataset(abc.ABC):
  def __init__(self, dataset_name: str, urls: Dict[str, Dict[str, str]],
               git_url: Dict[str, str], dataset_description: str,
               stage_1_dir: str,
               configuration: Configuration = Configuration(),
               transformation_funcs=(), filter_funcs=(),
               user_defined_split_range=(), num_shards=1000, seed=0,
               deduplicate=False):
  def stage_1_mkdirs(self):
  def download_dataset(self):
  def download_dataset_using_git(self):
  def download_dataset_using_requests(self):
    def _get_confirm_token(response):
  def get_random_split(self):
  def get_all_raw_data_paths(self):
  def raw_data_paths_to_raw_data_do_fn(self):
  def raw_data_to_graph_to_output_example(self, raw_data):
  def exists_graph_to_output_example(self):
  def convert_and_write_pipeline(self, root):
  def update_dataset_metadata(self, metadata, split, raw_data_dicts,
                              graph_to_output_example_dicts):
  def _count_graph_to_output_example_with_max_num_nodes(
      self, graph_to_output_example_dicts, split, max_num_nodes):
  def run_pipeline(self): class CuBertExceptionClassificationDataset(cubert_dataset.CuBertDataset):
  """Converts CuBERT Exception Classification data to a PLUR dataset.

  The dataset is created by: Aditya Kanade, Petros Maniatis, Gogul Balakrishnan,
  Kensen Shi Proceedings of the 37th International Conference on Machine
  Learning, PMLR 119:5110-5121, 2020.

  The task is to predict the exception type (one of 20) that should be caught
  by an exception block in Python.

  The classes (i.e., available exception types) are:
  ValueError, KeyError, AttributeError, TypeError, OSError, IOError,
  ImportError, IndexError, DoesNotExist, KeyboardInterrupt, StopIteration,
  AssertionError, SystemExit, RuntimeError, HTTPError, UnicodeDecodeError,
  NotImplementedError, ValidationError, ObjectDoesNotExist, NameError.

  The context consists of the body of a Python function, where the exception
  type to predict has been replaced by a special __HOLE__ token. This context is
  tokenized using the CuBERT Python tokenizer, and encoded as WordPiece
  vocabulary IDs from the CuBERT-released Python vocabulary. The graph
  representation is as a chain of nodes, each holding a WordPiece subtoken. The
  output is one of the classification labels above.
  """

  ALL_CLASSES = frozenset((
      'AssertionError',
      'AttributeError',
      'DoesNotExist',
      'HTTPError',
      'IOError',
      'ImportError',
      'IndexError',
      'KeyError',
      'KeyboardInterrupt',
      'NameError',
      'NotImplementedError',
      'OSError',
      'ObjectDoesNotExist',
      'RuntimeError',
      'StopIteration',
      'SystemExit',
      'TypeError',
      'UnicodeDecodeError',
      'ValidationError',
      'ValueError',
  ))

  def folder_path(self) -> str:
    """As per superclass."""
    return '20200621_Python/exception_datasets/'

  def dataset_name(self) -> str:
    """As per superclass."""
    return 'cubert_exception_classification_dataset'

  def dataset_description(self) -> str:
    """As per superclass."""
    return """From CuBERT website:

  Here we describe the 6 Python benchmarks we created. All 6 benchmarks were
  derived from ETH Py150 Open. All examples are stored as sharded text files.
  Each text line corresponds to a separate example encoded as a JSON object. For
  each dataset, we release separate training/validation/testing splits along the
  same boundaries that ETH Py150 Open splits its files to the corresponding
  splits. The fine-tuned models are the checkpoints of each model with the
  highest validation accuracy.

  Combinations of functions where one exception type has been masked, along with
  a label indicating the masked exception type. The JSON fields are:
    function: string, the source code of a function as text, in which one
      exception type has been replaced with the special token “HOLE”
    label: string, one of (ValueError, KeyError, AttributeError, TypeError,
      OSError, IOError, ImportError, IndexError, DoesNotExist,
      KeyboardInterrupt,
      StopIteration, AssertionError, SystemExit, RuntimeError, HTTPError,
      UnicodeDecodeError, NotImplementedError, ValidationError,
      ObjectDoesNotExist, NameError, None), the masked exception type. Note that
      None never occurs in the data and will be removed in a future release.
    info: string, an unformatted description of how the example was constructed,
      including the source dataset (always “ETHPy150Open”), the repository and
      filepath, and the fully-qualified function name.
  """

  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
    """Convert data example to the unified GraphToOutputExample data structure.

    The input is a function string, and the output is the full name of the
    exception-type class.

    There are no edges in this task, since it was targeting BERT.

    Args:
      data: A dictionary with 'function', 'label', and 'info' as keys.
      max_graph_size: The maximum number of input nodes allowed for the example.
      split: The split of the example.

    Raises:
      GraphToOutputExampleNotValidError if the GraphToOutputExample is not
      valid.

    Returns:
      A GraphToOutputExample.
    """
    return (
        self.single_classification_data_dictionary_to_graph_to_output_example(
            data, self.ALL_CLASSES, max_graph_size, split)) class CuBertFunctionDocstringClassificationDataset(cubert_dataset.CuBertDataset
                                                  ):
  """Converts CuBERT Function Docstring Classification data to a PLUR dataset.

  The dataset is created by: Aditya Kanade, Petros Maniatis, Gogul Balakrishnan,
  Kensen Shi Proceedings of the 37th International Conference on Machine
  Learning, PMLR 119:5110-5121, 2020.

  The task is to predict whether a function and a docstring match, or if they
  come from distinct contexts.

  The context consists of the body of a Python function and a docstring. This
  context is tokenized using the CuBERT Python tokenizer, and encoded as
  WordPiece vocabulary IDs from the CuBERT-released Python vocabulary. The graph
  representation is as a chain of nodes, each holding a WordPiece subtoken. The
  output is one of the two classification labels. We use separate node types
  for docstrings and for function bodies.
  """

  ALL_CLASSES = frozenset((
      'Correct',
      'Incorrect',
  ))

  def folder_path(self) -> str:
    """As per superclass."""
    return '20200621_Python/function_docstring_datasets/'

  def dataset_name(self) -> str:
    """As per superclass."""
    return 'cubert_function_docstring_classification_dataset'

  def dataset_description(self) -> str:
    """As per superclass."""
    return """From CuBERT website:

  Here we describe the 6 Python benchmarks we created. All 6 benchmarks were
  derived from ETH Py150 Open. All examples are stored as sharded text files.
  Each text line corresponds to a separate example encoded as a JSON object. For
  each dataset, we release separate training/validation/testing splits along the
  same boundaries that ETH Py150 Open splits its files to the corresponding
  splits. The fine-tuned models are the checkpoints of each model with the
  highest validation accuracy.

  Combinations of functions with their correct or incorrect documentation
  string, used to train a classifier that can tell which pairs go together. The
  JSON fields are:

    function: string, the source code of a function as text

    docstring: string, the documentation string for that function

    label: string, one of (“Incorrect”, “Correct”), the label of the example.

    info: string, an unformatted description of how the example was constructed,
    including the source dataset (always “ETHPy150Open”), the repository and
    filepath, the function name and, for “Incorrect” examples, the function
    whose docstring was substituted.
  """

  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
    """Convert data example to the unified GraphToOutputExample data structure.

    The input is a function string, and the output is a class.

    There are no edges in this task, since it was targeting BERT.

    Args:
      data: A dictionary with 'function', 'docstring', 'label', and 'info' as
        keys.
      max_graph_size: The maximum number of input nodes allowed for the example.
      split: The split of the example.

    Raises:
      GraphToOutputExampleNotValidError if the GraphToOutputExample is not
      valid.

    Returns:
      A GraphToOutputExample.
    """
    del split  # Unused.
    function = data['function']
    docstring = data['docstring']
    label = data['label']
    assert label in self.ALL_CLASSES
    provenance = data['info']

    graph_to_output_example = GraphToOutputExample()

    # The input graph nodes are the source code tokens and docstring. We don't
    # filter any examples based on size. Instead, we trim the suffix of the
    # source-code sequence or docstring sequence, whichever is longest. This is
    # the same logic used by BERT for two-context examples. Note that we trim so
    # that the number of tokens plus the three delimiters (one extra between
    # function and docstring) is at most `max_graph_size`.
    sentences: List[List[str]] = (
        code_to_subtokenized_sentences.code_to_cubert_sentences(
            function, self.tokenizer, self.subword_text_encoder))
    docstring: List[List[str]] = (
        code_to_subtokenized_sentences.code_to_cubert_sentences(
            f'"""{docstring}"""', self.tokenizer, self.subword_text_encoder))
    docstring_tokens = sum(docstring, [])
    function_tokens = sum(sentences, [])
    # This updates `docstring_tokens` and `function_tokens` in place.
    _truncate_seq_pair(docstring_tokens, function_tokens, max_graph_size - 3)

    number_of_docstring_tokens = len(docstring_tokens)
    number_of_function_tokens = len(function_tokens)
    delimited_tokens = tuple(itertools.chain(('[CLS]_',), docstring_tokens,
                                             ('[SEP]_',), function_tokens,
                                             ('[SEP]_',)))
    types = tuple(['DOCSTRING'] * (number_of_docstring_tokens + 2) +
                  ['TOKEN'] * (number_of_function_tokens + 1))
    assert len(types) == len(delimited_tokens)
    assert len(delimited_tokens) <= max_graph_size
    for index, (token, token_type) in enumerate(zip(delimited_tokens, types)):
      graph_to_output_example.add_node(
          node_id=index, node_type=token_type, node_label=token)

    graph_to_output_example.add_class_output(label)
    graph_to_output_example.set_provenance(provenance)

    return graph_to_output_example class CuBertSwappedOperandClassificationDataset(cubert_dataset.CuBertDataset):
  """Converts CuBERT Swapped Operand Classification data to a PLUR dataset.

  The dataset is created by: Aditya Kanade, Petros Maniatis, Gogul Balakrishnan,
  Kensen Shi Proceedings of the 37th International Conference on Machine
  Learning, PMLR 119:5110-5121, 2020.

  The task is to predict whether a function contains swapped operands in a
  binary operator.

  The context consists of the body of a Python function. This context is
  tokenized using the CuBERT Python tokenizer, and encoded as WordPiece
  vocabulary IDs from the CuBERT-released Python vocabulary. The graph
  representation is as a chain of nodes, each holding a WordPiece subtoken. The
  output is one of the two classification labels.
  """

  ALL_CLASSES = frozenset((
      'Correct',
      'Swapped operands',
  ))

  def folder_path(self) -> str:
    """As per superclass."""
    return '20200621_Python/swapped_operands_datasets/'

  def dataset_name(self) -> str:
    """As per superclass."""
    return 'cubert_swapped_operand_classification_dataset'

  def dataset_description(self) -> str:
    """As per superclass."""
    return """From CuBERT website:

  Here we describe the 6 Python benchmarks we created. All 6 benchmarks were
  derived from ETH Py150 Open. All examples are stored as sharded text files.
  Each text line corresponds to a separate example encoded as a JSON object. For
  each dataset, we release separate training/validation/testing splits along the
  same boundaries that ETH Py150 Open splits its files to the corresponding
  splits. The fine-tuned models are the checkpoints of each model with the
  highest validation accuracy.

  Combinations of functions where one use binary operator’s arguments have been
  swapped, to create a buggy example, or left undisturbed, along with a label
  indicating if this bug-injection has occurred. The JSON fields are:

    function: string, the source code of a function as text.

    label: string, one of (“Correct”, “Swapped operands”) indicating if this is
      a buggy or bug-free example.

    info: string, an unformatted description of how the example was constructed,
      including the source dataset (always “ETHPy150Open”), the repository and
      filepath, the function, and whether the example is bugfree (marked
      “original”) or the operand swap has occurred (e.g., “swapped operands of
      not in”).
  """

  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
    """Convert data example to the unified GraphToOutputExample data structure.

    The input is a function string, and the output is a class.

    There are no edges in this task, since it was targeting BERT.

    Args:
      data: A dictionary with 'function', 'label', and 'info' as keys.
      max_graph_size: The maximum number of input nodes allowed for the example.
      split: The split of the example.

    Raises:
      GraphToOutputExampleNotValidError if the GraphToOutputExample is not
      valid.

    Returns:
      A GraphToOutputExample.
    """
    return (
        self.single_classification_data_dictionary_to_graph_to_output_example(
            data,
            classes=self.ALL_CLASSES,
            max_graph_size=max_graph_size,
            split=split)) class CuBertVariableMisuseClassificationDataset(cubert_dataset.CuBertDataset):
  """Converts CuBERT Variable Misuse Classification data to a PLUR dataset.

  The dataset is created by: Aditya Kanade, Petros Maniatis, Gogul Balakrishnan,
  Kensen Shi Proceedings of the 37th International Conference on Machine
  Learning, PMLR 119:5110-5121, 2020.

  The task is to predict whether a function contains a variable misuse bug.

  The context consists of the body of a Python function. This context is
  tokenized using the CuBERT Python tokenizer, and encoded as WordPiece
  vocabulary IDs from the CuBERT-released Python vocabulary. The graph
  representation is as a chain of nodes, each holding a WordPiece subtoken. The
  output is one of the two classification labels.
  """

  ALL_CLASSES = frozenset((
      'Correct',
      'Variable misuse',
  ))

  def folder_path(self) -> str:
    """As per superclass."""
    return '20200621_Python/variable_misuse_datasets/'

  def dataset_name(self) -> str:
    """As per superclass."""
    return 'cubert_variable_misuse_classification_dataset'

  def dataset_description(self) -> str:
    """As per superclass."""
    return """From CuBERT website:

  Here we describe the 6 Python benchmarks we created. All 6 benchmarks were
  derived from ETH Py150 Open. All examples are stored as sharded text files.
  Each text line corresponds to a separate example encoded as a JSON object. For
  each dataset, we release separate training/validation/testing splits along the
  same boundaries that ETH Py150 Open splits its files to the corresponding
  splits. The fine-tuned models are the checkpoints of each model with the
  highest validation accuracy.

  Combinations of functions where one use of a variable may have been replaced
  with another variable defined in the same context, along with a label
  indicating if this bug-injection has occurred. The JSON fields are:
    function: string, the source code of a function as text.

    label: string, one of (“Correct”, “Variable misuse”) indicating if this is a
      buggy or bug-free example.

    info: string, an unformatted description of how the example was constructed,
      including the source dataset (always “ETHPy150Open”), the repository and
      filepath, the function, and whether the example is bugfree (marked
      “original”) or the variable substitution that has occurred (e.g.,
      “correct_variable” → “incorrect_variable”).
  """

  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
    """Convert data example to the unified GraphToOutputExample data structure.

    The input is a function string, and the output is a class.

    There are no edges in this task, since it was targeting BERT.

    Args:
      data: A dictionary with 'function', 'label', and 'info' as keys.
      max_graph_size: The maximum number of input nodes allowed for the example.
      split: The split of the example.

    Raises:
      GraphToOutputExampleNotValidError if the GraphToOutputExample is not
      valid.

    Returns:
      A GraphToOutputExample.
    """
    return (
        self.single_classification_data_dictionary_to_graph_to_output_example(
            data, classes=self.ALL_CLASSES, max_graph_size=max_graph_size,
            split=split)) class CuBertVariableMisuseRepairDataset(cubert_dataset.CuBertDataset):
  """Converts CuBERT Variable Misuse Repair data to a PLUR dataset.

  The dataset is created by: Aditya Kanade, Petros Maniatis, Gogul Balakrishnan,
  Kensen Shi Proceedings of the 37th International Conference on Machine
  Learning, PMLR 119:5110-5121, 2020.

  The task is to predict where a variable misuse may be, and how to repair it,
  by pointing to the correct variable to use instead.

  The context consists of the body of a Python function. This context is
  tokenized using the CuBERT Python tokenizer, and encoded as WordPiece
  vocabulary IDs from the CuBERT-released Python vocabulary. The graph
  representation is as a chain of nodes, each holding a WordPiece subtoken. The
  output is one of the two classification labels.
  """

  ALL_CLASSES = frozenset((
      'Correct',
      'Variable misuse',
  ))

  _CONFOUNDING_STRING = '###DO_NOT_COPY_ME###'

  def __init__(self,
               stage_1_dir,
               configuration: plur_dataset.Configuration,
               *args,
               use_pointer_output: bool = True,
               allow_output_copy: bool = True,
               **kwargs) -> None:
    """As per superclass. It initializes variations (unpointed, no-copy)."""
    super().__init__(stage_1_dir, configuration, *args, **kwargs)
    self.use_pointer_output = use_pointer_output
    self.allow_output_copy = allow_output_copy

  def folder_path(self) -> str:
    """As per superclass."""
    return '20200621_Python/variable_misuse_repair_datasets/'

  def dataset_name(self) -> str:
    """As per superclass."""
    return 'cubert_variable_misuse_repair_dataset'

  def dataset_description(self) -> str:
    """As per superclass."""
    return """From CuBERT website:

  Here we describe the 6 Python benchmarks we created. All 6 benchmarks were
  derived from ETH Py150 Open. All examples are stored as sharded text files.
  Each text line corresponds to a separate example encoded as a JSON object. For
  each dataset, we release separate training/validation/testing splits along the
  same boundaries that ETH Py150 Open splits its files to the corresponding
  splits. The fine-tuned models are the checkpoints of each model with the
  highest validation accuracy.

  Variable-misuse localization and repair. Combinations of functions where one
  use of a variable may have been replaced with another variable defined in the
  same context, along with information that can be used to localize and repair
  the bug, as well as the location of the bug if such a bug exists. The JSON
  fields are:

  function: a list of strings, the source code of a function, tokenized with the
  vocabulary common vocabulary. Note that, unlike other task datasets, this
  dataset gives a tokenized function, rather than the code as a single string.

  target_mask: a list of integers (0 or 1). If the integer at some position is
  1, then the token at the corresponding position of the function token list is
  a correct repair for the introduced bug. If a variable has been split into
  multiple tokens, only the first subtoken is marked in this mask. If the
  example is bug-free, all integers are 0.

  error_location_mask: a list of integers (0 or 1). If the integer at some
  position is 1, then there is a variable-misuse bug at the corresponding
  location of the tokenized function. In a bug-free example, the first integer
  is 1. There is exactly one integer set to 1 for all examples. If a variable
  has been split into multiple tokens, only the first subtoken is marked in this
  mask.

  candidate_mask: a list of integers (0 or 1). If the integer at some position
  is 1, then the variable starting at that position in the tokenized function is
  a candidate to consider when repairing a bug. Candidates are all variables
  defined in the function parameters or via variable declarations in the
  function. If a variable has been split into multiple tokens, only the first
  subtoken is marked in this mask, for each candidate.

  provenance: string, an unformatted description of how the example was
  constructed, including the source dataset (always “ETHPy150Open”), the
  repository and filepath, the function, and whether the example is bugfree
  (marked “original”) or the buggy/repair token positions and variables (e.g.,
  “16/18 kwargs → self”). 16 is the position of the introduced error, 18 is the
  location of the repair.
  """

  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
    """Convert data example to the unified GraphToOutputExample data structure.

    The input is a function string, and the output is localization pointer and
    if the example is buggy, a sequence of subtokens for the repair.

    There are no edges in this task, since it was targeting BERT.

    Args:
      data: A dictionary with 'function', 'target_mask', 'candidate_mask',
        'error_location_mask', and 'provenance' as keys.
      max_graph_size: The maximum number of input nodes allowed for the example.
      split: The split of the example.

    Raises:
      GraphToOutputExampleNotValidError if the GraphToOutputExample is not
      valid.

    Returns:
      A GraphToOutputExample. None if this example cannot be created.
    """

    # The target mask points to the first subtoken of the repair. However, there
    # might be multiple subtokens in the repair identifier. Since we want to put
    # in the output the full repair token (i.e., all of its subtokens), we must
    # figure it out from the input. We start with the first target mask
    # occurrence and we decode the first whole token.
    target_mask = data['target_mask']
    function_tokens = data['function']
    target_indices = tuple(i for i, x in enumerate(target_mask) if x == 1)
    # We interpret all targets and require them to be identical, since we
    # have not done any trimming yet, and we assume the dataset to be correct.
    target_interpretations = []
    whole_tokens = []
    for target_index in target_indices:
      whole_token, end_target_index = (
          code_to_subtokenized_sentences.next_whole_token(
              function_tokens[target_index:], self.tokenizer,
              self.subword_text_encoder))
      target_subtokens = tuple(function_tokens[target_index:target_index +
                                               end_target_index])
      target_interpretations.append(target_subtokens)
      whole_tokens.append(whole_token)

    # Is this a valid example?
    original_error_location_mask = data['error_location_mask']
    original_is_bug_free = original_error_location_mask[0] == 1
    assert sum(original_error_location_mask) == 1, (
        'Exactly one error location mask must be set to 1, but instead got '
        f'{original_error_location_mask} in {data}')
    assert not target_indices or not original_is_bug_free, (
        # bug_free => not target_indices
        # not bug_free or not target_indices
        # Reminder: it's not possible for the first subtoken to be the site of
        # an error, so using the first subtoken as an indication of the absence
        # of bug is sound for this task, and in the absence of any prefix
        # trimming.
        'Bug-free (error mask on at first location) implies there must be '
        f'no targets in target mask, but instead got {data}')
    assert original_is_bug_free or target_indices, (
        # buggy => target_indices
        # not bug_free => target_indices
        # bug_free or target_indices
        'Buggy (error mask not on at first location) implies there must be '
        f'some targets in error mask, but instead got {data}')
    assert (
        len(target_indices) == len(target_interpretations) == sum(target_mask)
    ), (f'Expected target mask size {sum(target_mask)} to be the same as '
        f'target indices count {len(target_indices)}, and target '
        f'target interpretations count {len(target_interpretations)} '
        'but not quite: '
        f'{target_mask, target_indices, target_interpretations}')
    assert len(set(target_interpretations)) <= 1, (
        'Expected all target interpretations (if any) to be the same, but '
        f'instead got {target_interpretations}.')
    assert len(set(whole_tokens)) <= 1, (
        'Expected all whole tokens (if any) to be the same, but '
        f'instead got {whole_tokens}.')
    beam.metrics.Metrics.counter(_BEAM_METRIC_NAMESPACE,
                                 f'input.example_{split}.count').inc()
    if original_is_bug_free:
      beam.metrics.Metrics.counter(_BEAM_METRIC_NAMESPACE,
                                   f'input.is_bug_free_{split}.count').inc()
    else:
      beam.metrics.Metrics.counter(_BEAM_METRIC_NAMESPACE,
                                   f'input.is_buggy_{split}.count').inc()

    # Now that we have recovered a target string, we go back to graphinating the
    # input. First prune to size.
    pruned_tokens = function_tokens[:max_graph_size - 2]
    error_location_mask = original_error_location_mask[:max_graph_size - 2]
    is_bug_free = error_location_mask[0] == 1
    assert is_bug_free == original_is_bug_free
    candidate_mask = data['candidate_mask'][:max_graph_size - 2]
    pruned_target_mask = target_mask[:max_graph_size - 2]
    provenance = data['provenance']

    # Is this example impossible?
    if sum(error_location_mask) == 0:
      # There's an error location, but we pruned it.
      beam.metrics.Metrics.counter(_BEAM_METRIC_NAMESPACE,
                                   f'pruned_error_location_{split}.count').inc()
      return None
    if sum(pruned_target_mask) == 0 and not is_bug_free:
      # There's an error but the targets have been pruned.
      beam.metrics.Metrics.counter(
          _BEAM_METRIC_NAMESPACE,
          f'pruned_all_repair_targets_{split}.count').inc()
      return None

    graph_to_output_example = GraphToOutputExample()

    # The input graph nodes are the source code tokens. We don't filter any
    # examples based on size. Instead, we trim the suffix of the token sequence.
    # Note that we trim so that the number of tokens plus the two delimiters
    # is at most `max_graph_size`.
    delimited_tokens = tuple(
        itertools.chain(('[CLS]',), pruned_tokens, ('[SEP]',)))

    number_of_delimited_tokens = len(delimited_tokens)
    for index, token in enumerate(delimited_tokens):
      # index is in [0, number_of_delimited_tokens). We want it to index into
      # the candidate mask, which maps to [1, number_of_delimited_tokens - 1),
      # since the delimiters are at the two ends, after any trimming. We
      # definitely don't want to mark 0 and `number_of_delimited_tokens` - 1 as
      # repair candidates. For the indices in between, we want to look them up
      # into `candidate_mask` but offset by 1 to the left.
      if index == 0 or index == number_of_delimited_tokens - 1:
        is_repair_candidate = False
      else:
        original_index = index - 1
        is_repair_candidate = bool(candidate_mask[original_index])
      graph_to_output_example.add_node(
          node_id=index,
          node_type='TOKEN',
          node_label=token,
          is_repair_candidate=is_repair_candidate)

    # The first output is a pointer pointing to node 1 (for bug-free examples)
    # or some other element otherwise.
    error_location_index = error_location_mask.index(1) + 1
    if self.use_pointer_output:
      # We encode pointers as a Po entry (from ToCoPo).
      graph_to_output_example.add_pointer_output(error_location_index)
    else:
      # As an ablation, instead of using a Po entry for error localization, we
      # use a special token with the location.
      if is_bug_free:
        graph_to_output_example.add_token_output('BUG_FREE')
      else:
        graph_to_output_example.add_token_output(
            f'ERROR_LOCATION_{error_location_index}')

    # Now we output the repair, as a sequence of subtokens.
    if target_interpretations:
      confounder = '' if self.allow_output_copy else self._CONFOUNDING_STRING
      for repair_subtoken in target_interpretations[0]:
        graph_to_output_example.add_token_output(confounder + repair_subtoken)
      graph_to_output_example.add_additional_field('repair_text',
                                                   whole_tokens[0])

    graph_to_output_example.set_provenance(provenance)

    if original_is_bug_free:
      beam.metrics.Metrics.counter(_BEAM_METRIC_NAMESPACE,
                                   f'output.is_bug_free_{split}.count').inc()
    else:
      beam.metrics.Metrics.counter(_BEAM_METRIC_NAMESPACE,
                                   f'output.is_buggy_{split}.count').inc()
    return graph_to_output_example class CuBertWrongOperatorClassificationDataset(cubert_dataset.CuBertDataset):
  """Converts CuBERT Wrong Operator Classification data to a PLUR dataset.

  The dataset is created by: Aditya Kanade, Petros Maniatis, Gogul Balakrishnan,
  Kensen Shi Proceedings of the 37th International Conference on Machine
  Learning, PMLR 119:5110-5121, 2020.

  The task is to predict whether a function contains an incorrect operator.

  The context consists of the body of a Python function. This context is
  tokenized using the CuBERT Python tokenizer, and encoded as WordPiece
  vocabulary IDs from the CuBERT-released Python vocabulary. The graph
  representation is as a chain of nodes, each holding a WordPiece subtoken. The
  output is one of the two classification labels.
  """

  ALL_CLASSES = frozenset((
      'Correct',
      'Wrong binary operator',
  ))

  def folder_path(self) -> str:
    """As per superclass."""
    return '20200621_Python/wrong_binary_operator_datasets/'

  def dataset_name(self) -> str:
    """As per superclass."""
    return 'cubert_wrong_operator_classification_dataset'

  def dataset_description(self) -> str:
    """As per superclass."""
    return """From CuBERT website:

  Here we describe the 6 Python benchmarks we created. All 6 benchmarks were
  derived from ETH Py150 Open. All examples are stored as sharded text files.
  Each text line corresponds to a separate example encoded as a JSON object. For
  each dataset, we release separate training/validation/testing splits along the
  same boundaries that ETH Py150 Open splits its files to the corresponding
  splits. The fine-tuned models are the checkpoints of each model with the
  highest validation accuracy.

  Combinations of functions where one binary operator has been swapped with
  another, to create a buggy example, or left undisturbed, along with a label
  indicating if this bug-injection has occurred. The JSON fields are:

    function: string, the source code of a function as text.

    label: string, one of (“Correct”, “Wrong binary operator”) indicating if
      this is a buggy or bug-free example.

    info: string, an unformatted description of how the example was constructed,
      including the source dataset (always “ETHPy150Open”), the repository and
      filepath, the function, and whether the example is bugfree (marked
      “original”) or the operator replacement has occurred (e.g., “== -> !=”).
  """

  def data_to_graph_to_output_example(
      self, data: Mapping[str, Any],
      max_graph_size: int,
      split: str) -> Optional[GraphToOutputExample]:
    """Convert data example to the unified GraphToOutputExample data structure.

    The input is a function string, and the output is a class.

    There are no edges in this task, since it was targeting BERT.

    Args:
      data: A dictionary with 'function', 'label', and 'info' as keys.
      max_graph_size: The maximum number of input nodes allowed for the example.
      split: The split of the example.

    Raises:
      GraphToOutputExampleNotValidError if the GraphToOutputExample is not
      valid.

    Returns:
      A GraphToOutputExample.
    """
    return (
        self.single_classification_data_dictionary_to_graph_to_output_example(
            data, classes=self.ALL_CLASSES, max_graph_size=max_graph_size,
            split=split)) STAGE_1_DIRNAME = 'stage_1'
RAW_DATA_DIRNAME = 'raw_data'
GIT_REPO_DIRNAME = 'git_repo'
GRAPH_TO_OUTPUT_EXAMPLE_DIRNAME = 'graph_to_output_example'
TRAIN_SPLIT_NAME = 'train'
VALIDATION_SPLIT_NAME = 'valid'
TEST_SPLIT_NAME = 'test'
METADATA_FILENAME = 'dataset_metadata.json'
STAGE_2_DIRNAME = 'stage_2'
VOCAB_FILES_DIRNAME = 'vocab_files'
NODE_TYPE_VOCAB_FILENAME = 'node_type_vocab.txt'
NODE_LABEL_VOCAB_FILENAME = 'node_label_vocab.txt'
EDGE_TYPE_VOCAB_FILENAME = 'edge_type_vocab.txt'
OUTPUT_TOKEN_VOCAB_FILENAME = 'output_token_vocab.txt'
TFRECORD_DIRNAME = 'tfrecords'
POINTER_TOKEN = 'POINTER'
CLASS_TOKEN_PREFIX = 'CLASS_'
DONE_TOKEN = 'DONE'
PAD_TOKEN = '<pad>'
OOV_TOKEN = '<OOV>'
RESERVED_TOKENS = [OOV_TOKEN, PAD_TOKEN]
NUM_RESERVED_TOKENS = len(RESERVED_TOKENS)
EXAMPLE_GUID = 'graph/global_id'  # For future use.
KEY_PROVENANCE = 'metadata/provenance'
KEY_OUTPUT_TOKEN_IDS = 'targets/token_ids'
KEY_OUTPUT_TOKEN_TEXTS = 'targets/texts'
KEY_COPY_INPUT_INDICES = 'nodes/copy_indices'
KEY_COPY_OUTPUT_INDICES = 'targets/copy_indices'
KEY_POINTER_INPUT_INDICES = 'nodes/pointer_indices'
KEY_POINTER_OUTPUT_INDICES = 'targets/pointer_indices'
KEY_NODE_TOKEN_IDS = 'nodes/node_id_sequences'
KEY_NODE_TEXTS = 'nodes/node_text_sequences'
KEY_NODE_TOKEN_LENGTHS = 'nodes/node_sequence_lengths'
KEY_NODE_GUIDS = 'nodes/global_node_id'  # For future use.
KEY_NODE_TYPE_IDS = 'nodes/node_types'
KEY_EDGE_TYPE_IDS = 'edges/types'
KEY_EDGE_SOURCE_INDICES = 'edges/source_indices'
KEY_EDGE_DESTINATION_INDICES = 'edges/dest_indices'
KEY_MASKING_CANDIDATE_INDICES = 'masking/candidate_indices'
PROVENANCE_TENSOR_NAME = 'provenance'
TARGET_TOKEN_IDS_TENSOR_NAME = 'target_token_ids'
COPY_INDICES_TENSOR_NAME = 'copy_indices'
POINTER_INDICES_TENSOR_NAME = 'pointer_indices'
NODE_ID_SEQUENCES_TENSOR_NAME = 'node_id_sequences'
NODE_TYPES_TENSOR_NAME = 'node_types'
EDGE_INDICATORS_TENSOR_NAME = 'edge_indicators'
NODE_TEXT_SEQUENCES_TENSOR_NAME = 'node_text_sequences'
TARGET_TEXTS_TENSOR_NAME = 'target_texts'
MASKING_CANDIDATE_TENSOR_NAME = 'candidate_indices'
HOPPITY_ADD_NODE_OP_NAME = 'add_node'
HOPPITY_DEL_NODE_OP_NAME = 'del_node'
HOPPITY_REPLACE_VAL_OP_NAME = 'replace_val'
HOPPITY_REPLACE_TYPE_OP_NAME = 'replace_type'
HOPPITY_REPLACE_NOOP_OP_NAME = 'NoOp'
HOPPITY_OUTPUT_UNKNOWN_TOKEN = 'UNKNOWN'
NATURALIZE_CONTEXT_TOKEN_NODE_NAME = 'CONTEXT_TOKEN'
NATURALIZE_NEXT_TOKEN_EDGE_NAME = 'NEXT_TOKEN'
NATURALIZE_CONTEXT_ROOT_NODE_NAME = 'CONTEXT_ROOT'
NATURALIZE_CONTEXT_TOKEN_EDGE_NAME = 'CONTEXT_TOKEN'
NATURALIZE_FEATURE_NODE_NODE_NAME = 'FEATURE_NODE'
NATURALIZE_NEXT_FEATURE_EDGE_NAME = 'NEXT_FEATURE'
NATURALIZE_FEATURE_ROOT_NODE_NAME = 'FEATURE_ROOT'
NATURALIZE_FEATURE_EDGE_NAME = 'FEATURE'
NATURALIZE_NATURALIZE_ROOT_NODE_NAME = 'NATURALIZE_ROOT'
NATURALIZE_LOCAL_CONTEXT_EDGE_NAME = 'LOCAL_CONTEXT'
NATURALIZE_FEATURES_EDGE_NAME = 'FEATURES' class GraphToOutputExampleNotValidError(Exception):
  """Exception to be thrown if check_if_valid fails."""
  pass
import json
import os
import random
import apache_beam as beam
from typing import Any, Iterable, Mapping
from plur.stage_1 import cubert_dataset
from plur.stage_1 import plur_dataset
from plur.stage_1.cubert_exception_classification_dataset import CuBertExceptionClassificationDataset
from plur.stage_1.cubert_function_docstring_classification_dataset import CuBertFunctionDocstringClassificationDataset
from plur.stage_1.cubert_swapped_operand_classification_dataset import CuBertSwappedOperandClassificationDataset
from plur.stage_1.cubert_variable_misuse_classification_dataset import CuBertVariableMisuseClassificationDataset
from plur.stage_1.cubert_variable_misuse_repair_dataset import CuBertVariableMisuseRepairDataset
from plur.stage_1.cubert_wrong_operator_classification_dataset import CuBertWrongOperatorClassificationDataset
from plur.utils import constants
from plur.utils.graph_to_output_example import GraphToOutputExampleNotValidError
# Copyright 2021 Google LLC
#
# 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.
"""Converts multiple CuBERT datasets into a single PLUR task."""



# We add a new field name in GraphToOutputExample, to remember the unitask
# from which the example came. This will help downstream to generate
# per-unitask test datasets.
TASK_TYPE_FIELD_NAME = 'task'
PATH_FIELD_NAME = 'actual_path'


class CuBertMultitaskDataset(plur_dataset.PlurDataset):
  """Converts CuBERT data from different tasks into a PLUR dataset.

  The datasets are created by: Aditya Kanade, Petros Maniatis, Gogul
  Balakrishnan, Kensen Shi Proceedings of the 37th International Conference on
  Machine Learning, PMLR 119:5110-5121, 2020.

  This class combines multiple CuBERT task classes into a single multi-task
  dataset/task.
  """

  _UNITASK_DATSETS = {
      'EC': CuBertExceptionClassificationDataset,
      'FD': CuBertFunctionDocstringClassificationDataset,
      'SO': CuBertSwappedOperandClassificationDataset,
      'VM': CuBertVariableMisuseClassificationDataset,
      'VMR': CuBertVariableMisuseRepairDataset,
      'WB': CuBertWrongOperatorClassificationDataset,
  }

  # This is drawn from Table 1 in the paper.
  _UNITASK_TRAIN_SIZES = {
      'EC': 18480,
      'FD': 340846,
      'SO': 236246,
      'VM': 700708,
      'VMR': 700708,
      'WB': 459400,
  }

  def __init__(
      self,
      stage_1_dir,
      *args,
      configuration: plur_dataset.Configuration = plur_dataset.Configuration(),
      rebalance_skew: bool = False,
      **kwargs):
    """Initializes with and without rebalance skew.

    Args:
      stage_1_dir: As per superclass.
      *args: args.
      configuration: Any configuration options.
      rebalance_skew: If True, duplicate examples from each task to make the
        sizes of training datasets similar. If false, examples are left as
        input.
      **kwargs: kwargs.
    """

    # We initialize `unitask_datasets` first, because doing in later confuses
    # the type checker into thinking it may not be defined.
    if rebalance_skew:
      max_train_size = max(self._UNITASK_TRAIN_SIZES.values())
      self.unitask_duplication = {k: max_train_size/s
                                  for k, s in self._UNITASK_TRAIN_SIZES.items()}
    else:
      self.unitask_duplication = {}
    assert all((d >= 1.0 for d in self.unitask_duplication.values()))


    self.unitask_datasets: Mapping[str, cubert_dataset.CuBertDataset] = {}