# 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

  If split is `train`, shuffle data.

  Args:
    name: Name of dataset, such as `pfam/unaligned_encoded`,
    split: Split name. One of train/valid/test.
    field: Field to return from each example.
    num_epochs: Number of times to repeat dataset. None for infinite.
    shuffle_buffer: Shuffle buffer size. Only applied to train split.
    batch_size: Number of examples per batch.
    max_length: Max length of sequences to return.
    sample_length: If True, sample a random crop if sequence is too long.
    data_dir: Data directory containing tfds files.
    drop_remainder: Whether to drop the last batch when having less than
      batch_size elements.

  Returns:
    Tuple of (batched dataset, protein_vocab.Vocabulary)
  """

  shuffle = split == 'train'
  ds = tfds.load(
      name,
      split=split,
      with_info=False,
      data_dir=data_dir,
      shuffle_files=shuffle)

  # Construct vocab from stored metadata.
  # TODO(ddohan): Regenerate dataset with new vocab.

# 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


  If split is `train`, shuffle data.

  Args:
    name: Name of dataset, such as `pfam/unaligned_encoded`,
    split: Split name. One of train/valid/test.
    field: Field to return from each example.
    num_epochs: Number of times to repeat dataset. None for infinite.
    shuffle_buffer: Shuffle buffer size. Only applied to train split.
    batch_size: Number of examples per batch.
    max_length: Max length of sequences to return.
    sample_length: If True, sample a random crop if sequence is too long.
    data_dir: Data directory containing tfds files.
    drop_remainder: Whether to drop the last batch when having less than
      batch_size elements.

  Returns:
    Tuple of (batched dataset, protein_vocab.Vocabulary)
  """

  shuffle = split == 'train'
  ds = tfds.load(
      name,
      split=split,
      with_info=False,
      data_dir=data_dir,
      shuffle_files=shuffle)

  # Construct vocab from stored metadata.
  # TODO(ddohan): Regenerate dataset with new vocab.

  vocab = protein_vocab.make_protein_vocab(include_align_tokens=True)