n/a

PiperOrigin-RevId: 421538469
tensorflow · Jan 13, 2022 · 963f2c4 · 963f2c4
1 parent fef8d73
commit 963f2c4
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/tensorflow_transform/experimental/analyzers.py b/tensorflow_transform/experimental/analyzers.py
@@ -41,15 +41,7 @@
 # `collections.namedtuple` or `typing.NamedTuple` once the Spark issue is
 # resolved.
 from tfx_bsl.types import tfx_namedtuple
-
-try:
-  import apache_beam as beam  # pylint: disable=g-import-not-at-top. # pytype: disable=import-error
-  _BeamPTransform = beam.PTransform
-except ModuleNotFoundError:
-  _BeamPTransform = Any
-
-_APPROXIMATE_VOCAB_FILENAME_PREFIX = 'approx_vocab_'
-_APPROXIMATE_VOCAB_FREQUENCY_FILENAME_PREFIX = 'approx_vocab_frequency_'
+from typing_extensions import Protocol
 
 __all__ = [
     'PTransformAnalyzerCacheCoder',
@@ -62,6 +54,20 @@
 PTransformAnalyzerCacheCoder = analyzer_nodes.CacheCoder
 SimpleJsonPTransformAnalyzerCacheCoder = analyzer_nodes.JsonNumpyCacheCoder
 
+_APPROXIMATE_VOCAB_FILENAME_PREFIX = 'approx_vocab_'
+_APPROXIMATE_VOCAB_FREQUENCY_FILENAME_PREFIX = 'approx_vocab_frequency_'
+
+
+class _BeamPTransform(Protocol):
+  """Pytype for `beam.PTransform` without depending on beam in this module.
+  """
+
+  def expand(self, pcol: Any) -> Any:
+    ...
+
+  def default_label(self) -> str:
+    ...
+
 
 # TODO(zoyahav): Add an example for using this API.
 class CacheablePTransformAnalyzer(