# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):
# 
# Path: init2winit/dataset_lib/protein_vocab.py
# _AA_TOKENS = 'ACDEFGHIKLMNPQRSTVWY'
# _ANOMALOUS_AA_TOKENS = 'BOUXZ'
# GAP = '-'  # Default gap character.
# INSERT_GAP = '.'  # Gap character in insert states of an MSA.
# ALIGN_TOKENS = (INSERT_GAP, GAP)
# BOS_TOKEN = '<'  # Beginning of sequence token.
# EOS_TOKEN = '>'  # End of sequence token.
# PAD_TOKEN = '_'  # Padding token.
# MASK_TOKEN = '*'  # Mask token.
# SEP_TOKEN = '|'  # A special token for separating tokens for serialization.
# _DEFAULT_VOCAB_NAME = 'vocab'
# class Vocabulary(object):
#   def __init__(self,
#                tokens,
#                name=None,
#                include_bos=False,
#                include_eos=False,
#                include_pad=False,
#                include_mask=False,
#                bos_token=BOS_TOKEN,
#                eos_token=EOS_TOKEN,
#                pad_token=PAD_TOKEN,
#                mask_token=MASK_TOKEN,
#                disallow_sep_token=True):
#   def __eq__(self, other):
#   def _set_tokens(self, tokens):
#   def __setstate__(self, state):
#   def __getstate__(self):
#   def as_dict(self):
#   def from_dict(cls, state):
#   def copy(self, **kwargs):
#   def __len__(self):
#   def name(self):
#   def tokens(self):
#   def token_ids(self):
#   def bos(self):
#   def eos(self):
#   def mask(self):
#   def pad(self):
#   def is_valid(self, value):
#   def are_valid(self, values):
#   def encode_token(self, token):
#   def encode(self, tokens):
#   def decode_token(self, int_token):
#   def decode(self, values, stop_at_eos=False, as_str=False):
# def make_protein_vocab(include_anomalous_amino_acids=True,
#                        include_bos=True,
#                        include_eos=True,
#                        include_pad=True,
#                        include_mask=True,
#                        include_align_tokens=False):
# 
# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


import enum
import functools
import itertools
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from absl import logging
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import protein_vocab
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict
  masker = BertMasker(vocab=vocab)

  def train_iterator_fn():
    for batch_index, batch in enumerate(iter(train_ds)):
      batch_rng = jax.random.fold_in(data_rng, batch_index)
      yield _batch_to_dict(batch, masker, 'train', batch_rng)

  def eval_train_epoch(num_batches=None):
    eval_train_iter = iter(train_ds)
    for batch_index, batch in enumerate(
        itertools.islice(eval_train_iter, num_batches)):
      batch_rng = jax.random.fold_in(data_rng, batch_index)
      yield _batch_to_dict(batch, masker, 'eval', batch_rng)

  def valid_epoch(num_batches=None):
    valid_iter = iter(eval_ds)
    for batch_index, batch in enumerate(
        itertools.islice(valid_iter, num_batches)):
      batch_rng = jax.random.fold_in(data_rng, batch_index)
      yield _batch_to_dict(batch, masker, 'eval', batch_rng)

  # pylint: disable=unreachable
  def test_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.

  # pylint: enable=unreachable


# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):
# 
# Path: init2winit/dataset_lib/protein_vocab.py
# _AA_TOKENS = 'ACDEFGHIKLMNPQRSTVWY'
# _ANOMALOUS_AA_TOKENS = 'BOUXZ'
# GAP = '-'  # Default gap character.
# INSERT_GAP = '.'  # Gap character in insert states of an MSA.
# ALIGN_TOKENS = (INSERT_GAP, GAP)
# BOS_TOKEN = '<'  # Beginning of sequence token.
# EOS_TOKEN = '>'  # End of sequence token.
# PAD_TOKEN = '_'  # Padding token.
# MASK_TOKEN = '*'  # Mask token.
# SEP_TOKEN = '|'  # A special token for separating tokens for serialization.
# _DEFAULT_VOCAB_NAME = 'vocab'
# class Vocabulary(object):
#   def __init__(self,
#                tokens,
#                name=None,
#                include_bos=False,
#                include_eos=False,
#                include_pad=False,
#                include_mask=False,
#                bos_token=BOS_TOKEN,
#                eos_token=EOS_TOKEN,
#                pad_token=PAD_TOKEN,
#                mask_token=MASK_TOKEN,
#                disallow_sep_token=True):
#   def __eq__(self, other):
#   def _set_tokens(self, tokens):
#   def __setstate__(self, state):
#   def __getstate__(self):
#   def as_dict(self):
#   def from_dict(cls, state):
#   def copy(self, **kwargs):
#   def __len__(self):
#   def name(self):
#   def tokens(self):
#   def token_ids(self):
#   def bos(self):
#   def eos(self):
#   def mask(self):
#   def pad(self):
#   def is_valid(self, value):
#   def are_valid(self, values):
#   def encode_token(self, token):
#   def encode(self, tokens):
#   def decode_token(self, int_token):
#   def decode(self, values, stop_at_eos=False, as_str=False):
# def make_protein_vocab(include_anomalous_amino_acids=True,
#                        include_bos=True,
#                        include_eos=True,
#                        include_pad=True,
#                        include_mask=True,
#                        include_align_tokens=False):
# 
# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


# Path: init2winit/dataset_lib/proteins.py
import enum
import functools
import itertools
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from absl import logging
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import protein_vocab
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict

  masker = BertMasker(vocab=vocab)

  def train_iterator_fn():
    for batch_index, batch in enumerate(iter(train_ds)):
      batch_rng = jax.random.fold_in(data_rng, batch_index)
      yield _batch_to_dict(batch, masker, 'train', batch_rng)

  def eval_train_epoch(num_batches=None):
    eval_train_iter = iter(train_ds)
    for batch_index, batch in enumerate(
        itertools.islice(eval_train_iter, num_batches)):
      batch_rng = jax.random.fold_in(data_rng, batch_index)
      yield _batch_to_dict(batch, masker, 'eval', batch_rng)

  def valid_epoch(num_batches=None):
    valid_iter = iter(eval_ds)
    for batch_index, batch in enumerate(
        itertools.islice(valid_iter, num_batches)):
      batch_rng = jax.random.fold_in(data_rng, batch_index)
      yield _batch_to_dict(batch, masker, 'eval', batch_rng)

  # pylint: disable=unreachable
  def test_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.

  # pylint: enable=unreachable


  return Dataset(train_iterator_fn, eval_train_epoch, valid_epoch, test_epoch)