# Path: flaxformer/param_conversion_util.py
# def load_tf_params(checkpoint_path: str) -> Dict[str, Any]:
# def convert_tf_params(
#     tf_params: Dict[str, Any],
#     mapping: Mapping[Tuple[str, ...], str],
#     tf_prefix: str = '',
# ) -> Dict[str, Any]:
# def get_int_regex_matches(pattern: str, state: Mapping[str, Any]) -> List[int]:


from typing import Any, Dict, Tuple
from flax.core import frozen_dict
from flaxformer import param_conversion_util
import flax
import tensorflow as tf





def load_params_from_tf_checkpoint(
    checkpoint_path: str) -> Tuple[flax.core.FrozenDict, flax.core.FrozenDict]:
  """Load Flax BERT Encoder params from a TF BERT checkpoint.

  The method loads the provided TF checkpoint, stores it in the dictionary
  with tf parameter names as keys and the weight tensors as values. Unlike Flax
  TF doesn't return the nested structure of the model weights, i.e. keys are
  recorded in the following format:
  "outer_layer_name/..middle_layer_name/...inner_tensor". Since we know the
  mapping between the TF keys and the Flax keys we recursively assemble
  the Flax params dictionary starting from the most outer layer.

  Args:
    checkpoint_path: Path to the TF checkpoint.

  Returns:
    A tuple of frozen dicts that are the parameters for the bert.BertEncoder and
    bert.heads.BertPooler.
  """
  ckpt_reader = tf.train.load_checkpoint(checkpoint_path)
  tf_params = {
      tf_name: ckpt_reader.get_tensor(tf_name)
      for tf_name in ckpt_reader.get_variable_to_dtype_map()
  }
  encoder_params = convert_bert_encoder_params(tf_params, 'bert/')

# Path: flaxformer/param_conversion_util.py
# def load_tf_params(checkpoint_path: str) -> Dict[str, Any]:
# def convert_tf_params(
#     tf_params: Dict[str, Any],
#     mapping: Mapping[Tuple[str, ...], str],
#     tf_prefix: str = '',
# ) -> Dict[str, Any]:
# def get_int_regex_matches(pattern: str, state: Mapping[str, Any]) -> List[int]:


# Path: flaxformer/architectures/bert/bert_checkpoint_converter.py
from typing import Any, Dict, Tuple
from flax.core import frozen_dict
from flaxformer import param_conversion_util
import flax
import tensorflow as tf






def load_params_from_tf_checkpoint(
    checkpoint_path: str) -> Tuple[flax.core.FrozenDict, flax.core.FrozenDict]:
  """Load Flax BERT Encoder params from a TF BERT checkpoint.

  The method loads the provided TF checkpoint, stores it in the dictionary
  with tf parameter names as keys and the weight tensors as values. Unlike Flax
  TF doesn't return the nested structure of the model weights, i.e. keys are
  recorded in the following format:
  "outer_layer_name/..middle_layer_name/...inner_tensor". Since we know the
  mapping between the TF keys and the Flax keys we recursively assemble
  the Flax params dictionary starting from the most outer layer.

  Args:
    checkpoint_path: Path to the TF checkpoint.

  Returns:
    A tuple of frozen dicts that are the parameters for the bert.BertEncoder and
    bert.heads.BertPooler.
  """
  ckpt_reader = tf.train.load_checkpoint(checkpoint_path)
  tf_params = {
      tf_name: ckpt_reader.get_tensor(tf_name)
      for tf_name in ckpt_reader.get_variable_to_dtype_map()
  }
  encoder_params = convert_bert_encoder_params(tf_params, 'bert/')

  pooler_params = param_conversion_util.convert_tf_params(