# Path: tensorflow_transform/annotators.py
# _ASSET_KEY_COLLECTION = 'tft_asset_key_collection'
# _ASSET_FILENAME_COLLECTION = 'tft_asset_filename_collection'
# _OBJECT_TRACKER = None
# VOCABULARY_SIZE_BY_NAME_COLLECTION = 'tft_vocabulary_size_by_name_collection'
#   _OBJECT_TRACKER = object_tracker
#     _OBJECT_TRACKER = None
# class ObjectTracker:
#   def __init__(self):
#   def trackable_objects(self) -> List[base.Trackable]:
#   def add_trackable_object(self, trackable_object: base.Trackable,
#                            name: Optional[str]):
# def object_tracker_scope(object_tracker: ObjectTracker):
# def _get_object(name: str) -> Optional[base.Trackable]:
# def track_object(trackable: base.Trackable, name: Optional[str]):
# def make_and_track_object(trackable_factory_callable: Callable[[],
#                                                                base.Trackable],
#                           name: Optional[str] = None) -> base.Trackable:
# def get_asset_annotations(graph: tf.Graph):
# def clear_asset_annotations(graph: tf.Graph):
# def annotate_asset(asset_key: str, asset_filename: str):
# def annotate_vocab_size(vocab_filename: str, vocab_size: tf.Tensor):
# 
# Path: tensorflow_transform/common_types.py
# def is_ragged_feature_available() -> bool:
# def is_ragged_feature(spec: FeatureSpecType) -> bool:


import contextlib
import enum
import tensorflow as tf
from typing import Callable, Optional, Tuple, Union
from tensorflow_transform import annotators
from tensorflow_transform import common_types
from tfx_bsl.types import tfx_namedtuple
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import func_graph
from tensorflow.python.framework import ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.util import object_identity
    def convert_dtype_and_swap(v, k):
      return key_dtype_fn(k), tf.cast(v, value_dtype)

    return dataset.enumerate().map(convert_dtype_and_swap)


def make_tfrecord_vocabulary_lookup_initializer(filename_tensor,
                                                key_dtype=tf.string,
                                                value_dtype=tf.int64,
                                                return_indicator_as_value=False,
                                                has_indicator=False):
  """Makes a lookup table initializer from a compressed tfrecord file."""
  graph = ops.get_default_graph()
  with contextlib.ExitStack() as stack:
    # TODO(b/165884902): Use tf.inside_function after dropping TF 2.3 support.
    # If filename_tensor is a graph tensor (e.g. temporary analyzer output), the
    # following operation cannot be lifted to init scope. Hence, check it is an
    # eager tensor or a string constant.
    if isinstance(graph, func_graph.FuncGraph) and isinstance(
        filename_tensor, (ops.EagerTensor, str)):
      # Lift the dataset creation out of graph construction to avoid
      # repeated initialization in TF2.
      stack.enter_context(tf.init_scope())

    dataset = _make_tfrecord_vocabulary_dataset(filename_tensor, key_dtype,
                                                value_dtype,
                                                return_indicator_as_value,
                                                has_indicator)
    # TODO(b/165884902): Use tf.inside_function after dropping TF 2.3 support.
    if isinstance(graph, func_graph.FuncGraph):

# Path: tensorflow_transform/annotators.py
# _ASSET_KEY_COLLECTION = 'tft_asset_key_collection'
# _ASSET_FILENAME_COLLECTION = 'tft_asset_filename_collection'
# _OBJECT_TRACKER = None
# VOCABULARY_SIZE_BY_NAME_COLLECTION = 'tft_vocabulary_size_by_name_collection'
#   _OBJECT_TRACKER = object_tracker
#     _OBJECT_TRACKER = None
# class ObjectTracker:
#   def __init__(self):
#   def trackable_objects(self) -> List[base.Trackable]:
#   def add_trackable_object(self, trackable_object: base.Trackable,
#                            name: Optional[str]):
# def object_tracker_scope(object_tracker: ObjectTracker):
# def _get_object(name: str) -> Optional[base.Trackable]:
# def track_object(trackable: base.Trackable, name: Optional[str]):
# def make_and_track_object(trackable_factory_callable: Callable[[],
#                                                                base.Trackable],
#                           name: Optional[str] = None) -> base.Trackable:
# def get_asset_annotations(graph: tf.Graph):
# def clear_asset_annotations(graph: tf.Graph):
# def annotate_asset(asset_key: str, asset_filename: str):
# def annotate_vocab_size(vocab_filename: str, vocab_size: tf.Tensor):
# 
# Path: tensorflow_transform/common_types.py
# def is_ragged_feature_available() -> bool:
# def is_ragged_feature(spec: FeatureSpecType) -> bool:


# Path: tensorflow_transform/tf_utils.py
import contextlib
import enum
import tensorflow as tf
from typing import Callable, Optional, Tuple, Union
from tensorflow_transform import annotators
from tensorflow_transform import common_types
from tfx_bsl.types import tfx_namedtuple
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import func_graph
from tensorflow.python.framework import ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.util import object_identity

    def convert_dtype_and_swap(v, k):
      return key_dtype_fn(k), tf.cast(v, value_dtype)

    return dataset.enumerate().map(convert_dtype_and_swap)


def make_tfrecord_vocabulary_lookup_initializer(filename_tensor,
                                                key_dtype=tf.string,
                                                value_dtype=tf.int64,
                                                return_indicator_as_value=False,
                                                has_indicator=False):
  """Makes a lookup table initializer from a compressed tfrecord file."""
  graph = ops.get_default_graph()
  with contextlib.ExitStack() as stack:
    # TODO(b/165884902): Use tf.inside_function after dropping TF 2.3 support.
    # If filename_tensor is a graph tensor (e.g. temporary analyzer output), the
    # following operation cannot be lifted to init scope. Hence, check it is an
    # eager tensor or a string constant.
    if isinstance(graph, func_graph.FuncGraph) and isinstance(
        filename_tensor, (ops.EagerTensor, str)):
      # Lift the dataset creation out of graph construction to avoid
      # repeated initialization in TF2.
      stack.enter_context(tf.init_scope())

    dataset = _make_tfrecord_vocabulary_dataset(filename_tensor, key_dtype,
                                                value_dtype,
                                                return_indicator_as_value,
                                                has_indicator)
    # TODO(b/165884902): Use tf.inside_function after dropping TF 2.3 support.
    if isinstance(graph, func_graph.FuncGraph):

      annotators.track_object(dataset, name=None)