# Path: alphafold/model/tf/input_pipeline.py
# NUM_RES = shape_placeholders.NUM_RES
# NUM_MSA_SEQ = shape_placeholders.NUM_MSA_SEQ
# NUM_EXTRA_SEQ = shape_placeholders.NUM_EXTRA_SEQ
# NUM_TEMPLATES = shape_placeholders.NUM_TEMPLATES
# def nonensembled_map_fns(data_config):
# def ensembled_map_fns(data_config):
# def process_tensors_from_config(tensors, data_config):
#   def wrap_ensemble_fn(data, i):
# def compose(x, fs):
# 
# Path: alphafold/model/tf/proteins_dataset.py
# def parse_tfexample(
#     raw_data: bytes,
#     features: protein_features.FeaturesMetadata,
#     key: Optional[str] = None) -> Dict[str, tf.train.Feature]:
# def _first(tensor: tf.Tensor) -> tf.Tensor:
# def parse_reshape_logic(
#     parsed_features: TensorDict,
#     features: protein_features.FeaturesMetadata,
#     key: Optional[str] = None) -> TensorDict:
# def _make_features_metadata(
#     feature_names: Sequence[str]) -> protein_features.FeaturesMetadata:
# def create_tensor_dict(
#     raw_data: bytes,
#     features: Sequence[str],
#     key: Optional[str] = None,
#     ) -> TensorDict:
# def np_to_tensor_dict(
#     np_example: Mapping[str, np.ndarray],
#     features: Sequence[str],
#     ) -> TensorDict:


import copy
import ml_collections
import numpy as np
import tensorflow.compat.v1 as tf
from typing import List, Mapping, Tuple
from alphafold.model.tf import input_pipeline
from alphafold.model.tf import proteins_dataset
  cfg = copy.deepcopy(config.data)

  feature_names = cfg.common.unsupervised_features
  if cfg.common.use_templates:
    feature_names += cfg.common.template_features

  with cfg.unlocked():
    cfg.eval.crop_size = num_res

  return cfg, feature_names


def tf_example_to_features(tf_example: tf.train.Example,
                           config: ml_collections.ConfigDict,
                           random_seed: int = 0) -> FeatureDict:
  """Converts tf_example to numpy feature dictionary."""
  num_res = int(tf_example.features.feature['seq_length'].int64_list.value[0])
  cfg, feature_names = make_data_config(config, num_res=num_res)

  if 'deletion_matrix_int' in set(tf_example.features.feature):
    deletion_matrix_int = (
        tf_example.features.feature['deletion_matrix_int'].int64_list.value)
    feat = tf.train.Feature(float_list=tf.train.FloatList(
        value=map(float, deletion_matrix_int)))
    tf_example.features.feature['deletion_matrix'].CopyFrom(feat)
    del tf_example.features.feature['deletion_matrix_int']

  tf_graph = tf.Graph()
  with tf_graph.as_default(), tf.device('/device:CPU:0'):
    tf.compat.v1.set_random_seed(random_seed)

# Path: alphafold/model/tf/input_pipeline.py
# NUM_RES = shape_placeholders.NUM_RES
# NUM_MSA_SEQ = shape_placeholders.NUM_MSA_SEQ
# NUM_EXTRA_SEQ = shape_placeholders.NUM_EXTRA_SEQ
# NUM_TEMPLATES = shape_placeholders.NUM_TEMPLATES
# def nonensembled_map_fns(data_config):
# def ensembled_map_fns(data_config):
# def process_tensors_from_config(tensors, data_config):
#   def wrap_ensemble_fn(data, i):
# def compose(x, fs):
# 
# Path: alphafold/model/tf/proteins_dataset.py
# def parse_tfexample(
#     raw_data: bytes,
#     features: protein_features.FeaturesMetadata,
#     key: Optional[str] = None) -> Dict[str, tf.train.Feature]:
# def _first(tensor: tf.Tensor) -> tf.Tensor:
# def parse_reshape_logic(
#     parsed_features: TensorDict,
#     features: protein_features.FeaturesMetadata,
#     key: Optional[str] = None) -> TensorDict:
# def _make_features_metadata(
#     feature_names: Sequence[str]) -> protein_features.FeaturesMetadata:
# def create_tensor_dict(
#     raw_data: bytes,
#     features: Sequence[str],
#     key: Optional[str] = None,
#     ) -> TensorDict:
# def np_to_tensor_dict(
#     np_example: Mapping[str, np.ndarray],
#     features: Sequence[str],
#     ) -> TensorDict:


# Path: alphafold/model/features.py
import copy
import ml_collections
import numpy as np
import tensorflow.compat.v1 as tf
from typing import List, Mapping, Tuple
from alphafold.model.tf import input_pipeline
from alphafold.model.tf import proteins_dataset

  cfg = copy.deepcopy(config.data)

  feature_names = cfg.common.unsupervised_features
  if cfg.common.use_templates:
    feature_names += cfg.common.template_features

  with cfg.unlocked():
    cfg.eval.crop_size = num_res

  return cfg, feature_names


def tf_example_to_features(tf_example: tf.train.Example,
                           config: ml_collections.ConfigDict,
                           random_seed: int = 0) -> FeatureDict:
  """Converts tf_example to numpy feature dictionary."""
  num_res = int(tf_example.features.feature['seq_length'].int64_list.value[0])
  cfg, feature_names = make_data_config(config, num_res=num_res)

  if 'deletion_matrix_int' in set(tf_example.features.feature):
    deletion_matrix_int = (
        tf_example.features.feature['deletion_matrix_int'].int64_list.value)
    feat = tf.train.Feature(float_list=tf.train.FloatList(
        value=map(float, deletion_matrix_int)))
    tf_example.features.feature['deletion_matrix'].CopyFrom(feat)
    del tf_example.features.feature['deletion_matrix_int']

  tf_graph = tf.Graph()
  with tf_graph.as_default(), tf.device('/device:CPU:0'):
    tf.compat.v1.set_random_seed(random_seed)

    tensor_dict = proteins_dataset.create_tensor_dict(