apache-beam 2.74.0rc2__tar.gz → 2.75.0__tar.gz

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-beam
-Version: 2.74.0rc2
+Version: 2.75.0
 Summary: Apache Beam SDK for Python
 Home-page: https://beam.apache.org
 Download-URL: https://pypi.python.org/pypi/apache-beam
@@ -27,7 +27,7 @@ Requires-Dist: fastavro<2,>=0.23.6
 Requires-Dist: fasteners<1.0,>=0.3
 Requires-Dist: grpcio!=1.48.0,!=1.59.*,!=1.60.*,!=1.61.*,!=1.62.0,!=1.62.1,!=1.66.*,!=1.67.*,!=1.68.*,!=1.69.*,!=1.70.*,<2,>=1.33.1
 Requires-Dist: httplib2<0.32.0,>=0.8
-Requires-Dist: jsonpickle<4.0.0,>=3.0.0
+Requires-Dist: jsonpickle<5.0.0,>=3.0.4
 Requires-Dist: numpy<2.5.0,>=1.14.3
 Requires-Dist: objsize<0.8.0,>=0.6.1
 Requires-Dist: packaging>=22.0
@@ -58,7 +58,7 @@ Requires-Dist: Sphinx<8.0,>=7.0.0; extra == "docs"
 Requires-Dist: docstring-parser<1.0,>=0.15; extra == "docs"
 Requires-Dist: docutils>=0.18.1; extra == "docs"
 Requires-Dist: markdown; extra == "docs"
-Requires-Dist: pandas<2.3.0; extra == "docs"
+Requires-Dist: pandas<2.4.0; extra == "docs"
 Requires-Dist: openai; extra == "docs"
 Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "docs"
 Provides-Extra: test
@@ -68,7 +68,7 @@ Requires-Dist: freezegun>=0.3.12; extra == "test"
 Requires-Dist: jinja2<3.2,>=3.0; extra == "test"
 Requires-Dist: joblib>=1.0.1; extra == "test"
 Requires-Dist: mock<6.0.0,>=1.0.1; extra == "test"
-Requires-Dist: pandas<2.3.0; extra == "test"
+Requires-Dist: pandas<2.4.0; extra == "test"
 Requires-Dist: parameterized<0.10.0,>=0.7.1; extra == "test"
 Requires-Dist: pydot<2,>=1.2.0; extra == "test"
 Requires-Dist: pyhamcrest!=1.10.0,<3.0.0,>=1.9; extra == "test"
@@ -80,7 +80,7 @@ Requires-Dist: pytest-timeout<3,>=2.1.0; extra == "test"
 Requires-Dist: scikit-learn<1.8.0,>=0.20.0; extra == "test"
 Requires-Dist: sqlalchemy<3.0,>=1.3; extra == "test"
 Requires-Dist: psycopg2-binary<3.0,>=2.8.5; extra == "test"
-Requires-Dist: testcontainers[kafka,milvus,mysql]<5.0.0,>=4.0.0; extra == "test"
+Requires-Dist: testcontainers[kafka,milvus,mysql,qdrant]<5.0.0,>=4.0.0; extra == "test"
 Requires-Dist: cryptography>=41.0.2; extra == "test"
 Requires-Dist: hypothesis<6.148.4,>5.0.0; extra == "test"
 Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "test"
@@ -94,11 +94,13 @@ Requires-Dist: cachetools<7,>=3.1.0; extra == "gcp"
 Requires-Dist: google-api-core<3,>=2.0.0; extra == "gcp"
 Requires-Dist: google-apitools<0.5.32,>=0.5.31; python_version < "3.13" and extra == "gcp"
 Requires-Dist: google-apitools>=0.5.35; python_version >= "3.13" and extra == "gcp"
-Requires-Dist: google-auth<3,>=1.18.0; extra == "gcp"
+Requires-Dist: google-auth<3,>=2.0.0; extra == "gcp"
 Requires-Dist: google-auth-httplib2<0.3.0,>=0.1.0; extra == "gcp"
 Requires-Dist: google-cloud-datastore<3,>=2.0.0; extra == "gcp"
 Requires-Dist: google-cloud-pubsub<3,>=2.1.0; extra == "gcp"
 Requires-Dist: google-cloud-storage<4,>=2.18.2; extra == "gcp"
+Requires-Dist: google-cloud-resource-manager<2,>=1.12.0; extra == "gcp"
+Requires-Dist: google-cloud-dataflow-client<0.14.0,>=0.13.0; extra == "gcp"
 Requires-Dist: google-cloud-bigquery<4,>=2.0.0; extra == "gcp"
 Requires-Dist: google-cloud-bigquery-storage<3,>=2.6.3; extra == "gcp"
 Requires-Dist: google-cloud-core<3,>=2.0.0; extra == "gcp"
@@ -131,7 +133,7 @@ Requires-Dist: pydot<2,>=1.2.0; extra == "interactive"
 Requires-Dist: timeloop<2,>=1.0.2; extra == "interactive"
 Requires-Dist: nbformat<6,>=5.0.5; extra == "interactive"
 Requires-Dist: nbconvert<8,>=6.2.0; extra == "interactive"
-Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "interactive"
+Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.4,>=1.4.3; extra == "interactive"
 Provides-Extra: interactive-test
 Requires-Dist: needle<1,>=0.5.0; extra == "interactive-test"
 Requires-Dist: chromedriver-binary<118,>=117; extra == "interactive-test"
@@ -179,6 +181,7 @@ Requires-Dist: opentelemetry-api==1.37.0; extra == "p310-ml-test"
 Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p310-ml-test"
 Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p310-ml-test"
 Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p310-ml-test"
+Requires-Dist: qdrant-client>=1.15.0; extra == "p310-ml-test"
 Provides-Extra: p312-ml-test
 Requires-Dist: datatable; extra == "p312-ml-test"
 Requires-Dist: embeddings>=0.0.4; extra == "p312-ml-test"
@@ -201,6 +204,7 @@ Requires-Dist: opentelemetry-api==1.37.0; extra == "p312-ml-test"
 Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p312-ml-test"
 Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p312-ml-test"
 Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p312-ml-test"
+Requires-Dist: qdrant-client>=1.15.0; extra == "p312-ml-test"
 Provides-Extra: p313-ml-test
 Requires-Dist: embeddings>=0.0.4; extra == "p313-ml-test"
 Requires-Dist: onnxruntime; extra == "p313-ml-test"
@@ -223,6 +227,7 @@ Requires-Dist: opentelemetry-sdk==1.37.0; extra == "p313-ml-test"
 Requires-Dist: opentelemetry-exporter-otlp-proto-http==1.37.0; extra == "p313-ml-test"
 Requires-Dist: tf2onnx<1.18,>=1.17.0; extra == "p313-ml-test"
 Requires-Dist: pymilvus<3.0.0,>=2.5.10; extra == "p313-ml-test"
+Requires-Dist: qdrant-client>=1.15.0; extra == "p313-ml-test"
 Provides-Extra: aws
 Requires-Dist: boto3<2,>=1.9; extra == "aws"
 Provides-Extra: azure
@@ -230,7 +235,7 @@ Requires-Dist: azure-storage-blob<13,>=12.3.2; extra == "azure"
 Requires-Dist: azure-core<2,>=1.7.0; extra == "azure"
 Requires-Dist: azure-identity<2,>=1.12.0; extra == "azure"
 Provides-Extra: dataframe
-Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "dataframe"
+Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.4,>=1.4.3; extra == "dataframe"
 Provides-Extra: dask
 Requires-Dist: distributed>=2024.4.2; extra == "dask"
 Requires-Dist: dask>=2024.4.2; extra == "dask"
@@ -240,9 +245,9 @@ Provides-Extra: yaml
 Requires-Dist: docstring-parser<1.0,>=0.15; extra == "yaml"
 Requires-Dist: jinja2<3.2,>=3.0; extra == "yaml"
 Requires-Dist: virtualenv-clone<1.0,>=0.5; extra == "yaml"
-Requires-Dist: js2py<1,>=0.74; python_version < "3.12" and extra == "yaml"
+Requires-Dist: quickjs-ng<1.0.0,>=0.14.0; extra == "yaml"
 Requires-Dist: jsonschema<5.0.0,>=4.0.0; extra == "yaml"
-Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.3,>=1.4.3; extra == "yaml"
+Requires-Dist: pandas!=1.5.0,!=1.5.1,<2.4,>=1.4.3; extra == "yaml"
 Provides-Extra: torch
 Requires-Dist: torch<2.8.0,>=1.9.0; extra == "torch"
 Provides-Extra: tensorflow
@@ -251,7 +256,7 @@ Requires-Dist: absl-py>=0.12.0; extra == "tensorflow"
 Provides-Extra: transformers
 Requires-Dist: transformers<4.56.0,>=4.28.0; extra == "transformers"
 Requires-Dist: tensorflow>=2.12.0; extra == "transformers"
-Requires-Dist: torch>=1.9.0; extra == "transformers"
+Requires-Dist: torch<2.12.0,>=1.9.0; extra == "transformers"
 Provides-Extra: ml-cpu
 Requires-Dist: tensorflow>=2.12.0; extra == "ml-cpu"
 Requires-Dist: torch==2.8.0+cpu; extra == "ml-cpu"
@@ -280,6 +285,8 @@ Provides-Extra: tensorflow-hub
 Requires-Dist: tensorflow-hub<0.16.0,>=0.14.0; extra == "tensorflow-hub"
 Provides-Extra: milvus
 Requires-Dist: pymilvus<3.0.0,>=2.5.10; extra == "milvus"
+Provides-Extra: qdrant
+Requires-Dist: qdrant-client>=1.15.0; extra == "qdrant"
 Provides-Extra: vllm
 Requires-Dist: openai==1.107.1; extra == "vllm"
 Requires-Dist: vllm==0.10.1.1; extra == "vllm"

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/dataframe/frames.py

@@ -63,6 +63,10 @@ __all__ = [
 # Get major, minor version
 PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2]))
 
+# Find the name pandas uses for the NaN column in get_dummies().
+# Older pandas versions (<2.3.0) use 'nan', whereas newer versions (>=2.3.0) use 'NaN'.
+_DUMMY_NAN_COLUMN = 'NaN' if PD_VERSION >= (2, 3) else 'nan'
+
 
 def populate_not_implemented(pd_type):
   def wrapper(deferred_type):
@@ -5054,7 +5058,8 @@ class _DeferredStringMethods(frame_base.DeferredBase):
     # the data includes NaNs, which is not valid to be casted as a Category,
     # but nevertheless would be broadcasted as a column in get_dummies()
     columns = sorted(set().union(*split_cats))
-    columns = columns + ['nan'] if 'nan' not in columns else columns
+    if _DUMMY_NAN_COLUMN not in columns:
+      columns = columns + [_DUMMY_NAN_COLUMN]
 
     proxy = pd.DataFrame(columns=columns).astype(int)
 

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/dataframe/frames_test.py

@@ -37,6 +37,8 @@ from apache_beam.runners.interactive.testing.mock_env import isolated_env
 
 # Get major, minor version
 PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2]))
+# Get major, minor, patch version
+PD_FULL_VERSION = tuple(int(x) for x in re.findall(r'\d+', pd.__version__)[0:3])
 
 GROUPBY_DF = pd.DataFrame({
     'group': ['a' if i % 5 == 0 or i % 3 == 0 else 'b' for i in range(100)],
@@ -1437,7 +1439,7 @@ class DeferredFrameTest(_AbstractFrameTest):
     self._run_test(lambda s: s.unstack(level=0), s)
 
   @unittest.skipIf(
-      sys.version_info >= (3, 12) and PD_VERSION < (2, 3),
+      sys.version_info >= (3, 12) and PD_FULL_VERSION < (2, 3, 4),
       'https://github.com/pandas-dev/pandas/issues/58604')
   def test_unstack_pandas_example3(self):
     index = self._unstack_get_categorical_index()
@@ -2955,19 +2957,26 @@ class AllowNonParallelTest(unittest.TestCase):
       self._use_non_parallel_operation()
 
 
+_CONSTRUCTION_TIME_TEST_COLUMNS = [
+    'str_col', 'int_col', 'flt_col', 'cat_col', 'datetime_col'
+]
+
+
 class ConstructionTimeTest(unittest.TestCase):
   """Tests for operations that can be executed eagerly."""
-  DF = pd.DataFrame({
-      'str_col': ['foo', 'bar'] * 3,
-      'int_col': [1, 2] * 3,
-      'flt_col': [1.1, 2.2] * 3,
-      'cat_col': pd.Series(list('aabbca'), dtype="category"),
-      'datetime_col': pd.Series(
-          pd.date_range(
-              '1/1/2000', periods=6, freq='m', tz='America/Los_Angeles'))
-  })
-  DEFERRED_DF = frame_base.DeferredFrame.wrap(
-      expressions.PlaceholderExpression(DF.iloc[:0]))
+  @classmethod
+  def setUpClass(cls):
+    cls.DF = pd.DataFrame({
+        'str_col': ['foo', 'bar'] * 3,
+        'int_col': [1, 2] * 3,
+        'flt_col': [1.1, 2.2] * 3,
+        'cat_col': pd.Series(list('aabbca'), dtype="category"),
+        'datetime_col': pd.Series(
+            pd.date_range(
+                '1/1/2000', periods=6, freq='m', tz='America/Los_Angeles'))
+    })
+    cls.DEFERRED_DF = frame_base.DeferredFrame.wrap(
+        expressions.PlaceholderExpression(cls.DF.iloc[:0]))
 
   def _run_test(self, fn):
     expected = fn(self.DF)
@@ -2982,11 +2991,11 @@ class ConstructionTimeTest(unittest.TestCase):
     else:
       self.assertEqual(expected, actual)
 
-  @parameterized.expand(DF.columns)
+  @parameterized.expand(_CONSTRUCTION_TIME_TEST_COLUMNS)
   def test_series_name(self, col_name):
     self._run_test(lambda df: df[col_name].name)
 
-  @parameterized.expand(DF.columns)
+  @parameterized.expand(_CONSTRUCTION_TIME_TEST_COLUMNS)
   def test_series_dtype(self, col_name):
     self._run_test(lambda df: df[col_name].dtype)
     self._run_test(lambda df: df[col_name].dtypes)

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/dataframe/io.py

@@ -515,7 +515,7 @@ class _TruncatingFileHandle(object):
 
   @property
   def closed(self):
-    return False
+    return getattr(self._underlying, 'closed', False)
 
   def __iter__(self):
     # For pandas is_file_like.
@@ -584,7 +584,18 @@ class _TruncatingFileHandle(object):
     return res
 
   def flush(self):
-    self._underlying.flush()
+    if not self.closed:
+      try:
+        self._underlying.flush()
+      except ValueError:
+        pass
+
+  def close(self):
+    if not self.closed and hasattr(self._underlying, 'close'):
+      try:
+        self._underlying.close()
+      except (OSError, ValueError):
+        pass
 
 
 class _ReadFromPandasDoFn(beam.DoFn, beam.RestrictionProvider):
@@ -680,17 +691,27 @@ class _WriteToPandas(beam.PTransform):
     self.binary = binary
 
   def expand(self, pcoll):
-    if 'file_naming' in self.kwargs:
+    kwargs = dict(self.kwargs)
+    if 'file_naming' in kwargs:
       dir, name = self.path, ''
     else:
       dir, name = io.filesystems.FileSystems.split(self.path)
+    num_shards = kwargs.pop('num_shards', None)
+    max_writers_per_bundle = kwargs.pop('max_writers_per_bundle', None)
+    write_to_files_kwargs = {}
+    if num_shards is not None:
+      write_to_files_kwargs['shards'] = num_shards
+      write_to_files_kwargs['max_writers_per_bundle'] = 0
+    elif max_writers_per_bundle is not None:
+      write_to_files_kwargs['max_writers_per_bundle'] = max_writers_per_bundle
+
+    file_naming = kwargs.pop('file_naming', fileio.default_file_naming(name))
     return pcoll | fileio.WriteToFiles(
         path=dir,
-        shards=self.kwargs.pop('num_shards', None),
-        file_naming=self.kwargs.pop(
-            'file_naming', fileio.default_file_naming(name)),
+        file_naming=file_naming,
         sink=lambda _: _WriteToPandasFileSink(
-            self.writer, self.args, self.kwargs, self.incremental, self.binary))
+            self.writer, self.args, kwargs, self.incremental, self.binary),
+        **write_to_files_kwargs)
 
 
 class _WriteToPandasFileSink(fileio.FileSink):

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/dataframe/io_test.py

@@ -296,6 +296,51 @@ A     B
     self._run_truncating_file_handle_iter_test('aaa b cccccccccccccccccccc')
     self._run_truncating_file_handle_iter_test('aaa b ccccccccccccccccc ')
 
+  def test_truncating_filehandle_flush_on_closed_stream(self):
+    class ClosedFlushingStream(StringIO):
+      def flush(self):
+        if self.closed:
+          raise ValueError("I/O operation on closed file.")
+        super().flush()
+
+    s = 'a b c'
+    tracker = restriction_trackers.OffsetRestrictionTracker(
+        restriction_trackers.OffsetRange(0, len(s)))
+    underlying = ClosedFlushingStream(s)
+    handle = io._TruncatingFileHandle(
+        underlying, tracker, splitter=io._DelimSplitter(' ', 10))
+
+    # Verify that calling flush() when the underlying stream is closed
+    # succeeds without raising ValueError.
+    underlying.close()
+    handle.flush()
+    handle.close()
+
+  def test_truncating_filehandle_exception_suppression(self):
+    class FaultyStream(StringIO):
+      @property
+      def closed(self):
+        return False
+
+      def flush(self):
+        raise ValueError("Simulated flush error")
+
+      def close(self):
+        raise OSError("Simulated close error")
+
+    s = 'a b c'
+    tracker = restriction_trackers.OffsetRestrictionTracker(
+        restriction_trackers.OffsetRange(0, len(s)))
+    underlying = FaultyStream(s)
+    handle = io._TruncatingFileHandle(
+        underlying, tracker, splitter=io._DelimSplitter(' ', 10))
+
+    # Verify that ValueError raised during flush() is safely suppressed.
+    handle.flush()
+
+    # Verify that OSError raised during close() is safely suppressed.
+    handle.close()
+
   @parameterized.expand([
       ('defaults', {}),
       ('header', dict(header=1)),

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/examples/cookbook/ordered_window_elements/streaming.py

@@ -135,7 +135,7 @@ class OrderedWindowElementsDoFn(beam.DoFn):
 
   def _state_clear_range(self, buffer_state, range_lo, range_hi):
     """Clears a specified range of elements from the buffer state."""
-    # TODO: Dataflow runner v2 gets stuck when MIN_TIMESTAMP is used
+    # TODO: Dataflow Portable Runner gets stuck when MIN_TIMESTAMP is used
     # as the lower bound for clear_range. Investigate this further.
     buffer_state.clear_range(range_lo, range_hi)
 

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/examples/inference/gemini_image_generation.py

@@ -17,8 +17,8 @@
 
 """ A sample pipeline using the RunInference API to classify text using an LLM.
 This pipeline creates a set of prompts and sends it to a Gemini service then
-returns the predictions from the classifier model. This example uses the
-gemini-2.0-flash-001 model.
+returns generated images from the model. This example uses the
+gemini-2.5-flash-image model.
 """
 
 import argparse

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/examples/inference/gemini_text_classification.py

@@ -18,7 +18,7 @@
 """ A sample pipeline using the RunInference API to classify text using an LLM.
 This pipeline creates a set of prompts and sends it to a Gemini service then
 returns the predictions from the classifier model. This example uses the
-gemini-2.0-flash-001 model.
+gemini-2.5-flash model.
 """
 
 import argparse
@@ -88,7 +88,7 @@ def run(
   pipeline_options = PipelineOptions(pipeline_args)
   pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
   model_handler = GeminiModelHandler(
-      model_name='gemini-2.0-flash-001',
+      model_name='gemini-2.5-flash',
       request_fn=generate_from_string,
       api_key=known_args.api_key,
       project=known_args.project,

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/examples/inference/pytorch_sentiment.py

@@ -47,9 +47,6 @@ from transformers import DistilBertTokenizerFast
 
 class SentimentPostProcessor(beam.DoFn):
   """Processes PredictionResult to extract sentiment label and confidence."""
-  def __init__(self, tokenizer: DistilBertTokenizerFast):
-    self.tokenizer = tokenizer
-
   def process(self, element: tuple[str, PredictionResult]) -> Iterable[dict]:
     text, prediction_result = element
     logits = prediction_result.inference['logits']
@@ -62,16 +59,34 @@ class SentimentPostProcessor(beam.DoFn):
     }
 
 
-def tokenize_text(text: str,
-                  tokenizer: DistilBertTokenizerFast) -> tuple[str, dict]:
-  """Tokenizes input text using the specified tokenizer."""
-  tokenized = tokenizer(
-      text,
-      padding='max_length',
-      truncation=True,
-      max_length=128,
-      return_tensors="pt")
-  return text, {k: torch.squeeze(v) for k, v in tokenized.items()}
+class TokenizeTextDoFn(beam.DoFn):
+  """Initializes tokenizer per worker and tokenizes input text."""
+  def __init__(self, model_path: str):
+    self.model_path = model_path
+    self.tokenizer = None
+
+  def setup(self):
+    self.tokenizer = DistilBertTokenizerFast.from_pretrained(self.model_path)
+    if self.tokenizer.pad_token is None:
+      self.tokenizer.pad_token = '[PAD]'
+
+  def process(self, text: str) -> Iterable[tuple[str, dict]]:
+    tokenized = self.tokenizer(
+        text,
+        padding='max_length',
+        truncation=True,
+        max_length=128,
+        return_tensors="pt")
+    yield text, {k: torch.squeeze(v, 0) for k, v in tokenized.items()}
+
+
+class DistilBertForSequenceClassificationCompat(
+    DistilBertForSequenceClassification):
+  """Builds config in worker runtime to avoid cross-env config drift."""
+  def __init__(self, model_name: str, num_labels: int = 2):
+    config = _ensure_transformers_config_compat(
+        DistilBertConfig.from_pretrained(model_name, num_labels=num_labels))
+    super().__init__(config)
 
 
 class RateLimitDoFn(beam.DoFn):
@@ -83,6 +98,31 @@ class RateLimitDoFn(beam.DoFn):
     yield element
 
 
+def _ensure_transformers_config_compat(
+    config: DistilBertConfig) -> DistilBertConfig:
+  """Adds missing config attributes for cross-version transformers compatibility.
+
+  The benchmark can run with container images whose transformers version differs
+  from the launcher environment. Some versions assume these attributes exist.
+  """
+  # Use a default config instance as the source of canonical attributes for the
+  # transformers version available on the worker. This avoids chasing one
+  # missing field at a time (e.g. torchscript, output_attentions).
+  default_config = DistilBertConfig()
+  for key, value in default_config.to_dict().items():
+    if not hasattr(config, key):
+      setattr(config, key, value)
+
+  # Keep non-serialized fields explicitly for older/newer transformers mixes.
+  if not hasattr(config, 'pruned_heads'):
+    config.pruned_heads = {}
+  if not hasattr(config, 'torchscript'):
+    config.torchscript = False
+  if not hasattr(config, 'return_dict'):
+    config.return_dict = True
+  return config
+
+
 def parse_known_args(argv):
   """Parses command-line arguments for pipeline execution."""
   parser = argparse.ArgumentParser()
@@ -235,13 +275,14 @@ def run(
     pipeline_options.view_as(StandardOptions).streaming = True
 
   model_handler = PytorchModelHandlerKeyedTensor(
-      model_class=DistilBertForSequenceClassification,
-      model_params={'config': DistilBertConfig(num_labels=2)},
+      model_class=DistilBertForSequenceClassificationCompat,
+      model_params={
+          'model_name': known_args.model_path,
+          'num_labels': 2,
+      },
       state_dict_path=known_args.model_state_dict_path,
       device='GPU')
 
-  tokenizer = DistilBertTokenizerFast.from_pretrained(known_args.model_path)
-
   pipeline = test_pipeline or beam.Pipeline(options=pipeline_options)
 
   # Main pipeline: read, process, write result to BigQuery output table
@@ -264,9 +305,9 @@ def run(
 
   _ = (
       input
-      | 'Tokenize' >> beam.Map(lambda text: tokenize_text(text, tokenizer))
+      | 'Tokenize' >> beam.ParDo(TokenizeTextDoFn(known_args.model_path))
       | 'RunInference' >> RunInference(KeyedModelHandler(model_handler))
-      | 'PostProcess' >> beam.ParDo(SentimentPostProcessor(tokenizer))
+      | 'PostProcess' >> beam.ParDo(SentimentPostProcessor())
       | 'WriteToBigQuery' >> beam.io.WriteToBigQuery(
           known_args.output_table,
           schema='text:STRING, sentiment:STRING, confidence:FLOAT',
@@ -277,6 +318,7 @@ def run(
   result = pipeline.run()
   result.wait_until_finish(duration=1800000)  # 30 min
   result.cancel()
+  result.wait_until_finish(duration=600000)  # up to 10 min to settle cancel
 
   cleanup_pubsub_resources(
       project=known_args.project,

{apache_beam-2.74.0rc2 → apache_beam-2.75.0}/apache_beam/examples/matrix_power_test.py

@@ -17,6 +17,7 @@
 
 """Test for the matrix power integration test."""
 
+import glob
 import logging
 import tempfile
 import unittest
@@ -51,9 +52,15 @@ class MatrixPowerTest(unittest.TestCase):
         '--input_matrix=%s --input_vector=%s --exponent=%d --output=%s.result' %
         (matrix_path, vector_path, self.EXPONENT, vector_path)).split())
     # Parse result file and compare.
-    with open(vector_path + '.result-00000-of-00001') as result_file:
-      results = result_file.read()
-      self.assertEqual(sorted(self.EXPECTED_OUTPUT), sorted(results))
+    shard_paths = glob.glob(vector_path + '.result*')
+    self.assertTrue(
+        shard_paths,
+        'No output shards found matching prefix: %s.result' % vector_path)
+    results = []
+    for path in shard_paths:
+      with open(path) as result_file:
+        results.append(result_file.read())
+    self.assertEqual(sorted(self.EXPECTED_OUTPUT), sorted(''.join(results)))
 
 
 if __name__ == '__main__':