From 203601ebacc85d395eb02276de20c00b12b91141 Mon Sep 17 00:00:00 2001 From: hanzalaareeb Date: Thu, 20 Aug 2026 12:54:31 +0530 Subject: [PATCH 1/2] eager tensorflow imports block non-tfrecord pipeline paths --- .../creating_data_recorder_converter.py | 4 +++- dpsynth/pipeline_transformations/input_output.py | 13 ++++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/dpsynth/dataset_descriptors/creating_data_recorder_converter.py b/dpsynth/dataset_descriptors/creating_data_recorder_converter.py index 2030d05..91dc3ba 100644 --- a/dpsynth/dataset_descriptors/creating_data_recorder_converter.py +++ b/dpsynth/dataset_descriptors/creating_data_recorder_converter.py @@ -19,7 +19,7 @@ from dpsynth.dataset_descriptors import csv_descriptor from dpsynth.dataset_descriptors import dataset_descriptor -from dpsynth.dataset_descriptors import tfrecord_descriptor +# from dpsynth.dataset_descriptors import tfrecord_descriptor from dpsynth.pipeline_transformations import types @@ -35,6 +35,8 @@ def create_data_record_converter( } return csv_descriptor.CSVConverter(attributes_dict) elif data_format == types.DataFormat.TFRECORD: + from dpsynth.dataset_descriptors import tfrecord_descriptor + attributes_dict = { attr.name: attr.data_type for attr in dataset_desc.attributes } diff --git a/dpsynth/pipeline_transformations/input_output.py b/dpsynth/pipeline_transformations/input_output.py index 1cac62f..6c5e497 100644 --- a/dpsynth/pipeline_transformations/input_output.py +++ b/dpsynth/pipeline_transformations/input_output.py @@ -13,19 +13,20 @@ # limitations under the License. """Library for input/output transformations for data generation pipelines.""" - +from __future__ import annotations from collections.abc import Iterable, Mapping import csv import glob import os import pickle -from typing import Any +from typing import TYPE_CHECKING, Any import apache_beam as beam from dpsynth.pipeline_transformations import types import pandas as pd -import tensorflow as tf +if TYPE_CHECKING: + import tensorflow as tf def load_csv(pipeline: beam.Pipeline, path: str) -> beam.PCollection: """Loads the data to generate synthetic data for.""" @@ -79,6 +80,8 @@ def save_data_local( case types.DataFormat.CSV: save_csv(data, path, attributes) # pytype: disable=wrong-arg-types case types.DataFormat.TFRECORD: + import tensorflow as tf + os.makedirs(os.path.dirname(path), exist_ok=True) with tf.io.TFRecordWriter(path) as writer: for record in data: @@ -168,6 +171,8 @@ def load_data_for_beam( case types.DataFormat.CSV: return load_csv(pipeline, path) case types.DataFormat.TFRECORD: + import tensorflow as tf + return pipeline | 'ReadTFRecord' >> beam.io.ReadFromTFRecord( path, coder=beam.coders.ProtoCoder(tf.train.Example) ) @@ -184,6 +189,8 @@ def save_beam_data( """Saves the synthetic data to a file(s) using Beam.""" match data_format: case types.DataFormat.TFRECORD: + import tensorflow as tf + _ = data | 'WriteTFRecord' >> beam.io.WriteToTFRecord( path, coder=beam.coders.ProtoCoder(tf.train.Example) ) From 7f00c934499ccff76059f492e8ddca13aa46925f Mon Sep 17 00:00:00 2001 From: hanzalaareeb Date: Thu, 20 Aug 2026 14:54:24 +0530 Subject: [PATCH 2/2] remooving commented import --- .../dataset_descriptors/creating_data_recorder_converter.py | 2 -- dpsynth/pipeline_transformations/input_output.py | 3 --- 2 files changed, 5 deletions(-) diff --git a/dpsynth/dataset_descriptors/creating_data_recorder_converter.py b/dpsynth/dataset_descriptors/creating_data_recorder_converter.py index 91dc3ba..9c94927 100644 --- a/dpsynth/dataset_descriptors/creating_data_recorder_converter.py +++ b/dpsynth/dataset_descriptors/creating_data_recorder_converter.py @@ -19,7 +19,6 @@ from dpsynth.dataset_descriptors import csv_descriptor from dpsynth.dataset_descriptors import dataset_descriptor -# from dpsynth.dataset_descriptors import tfrecord_descriptor from dpsynth.pipeline_transformations import types @@ -36,7 +35,6 @@ def create_data_record_converter( return csv_descriptor.CSVConverter(attributes_dict) elif data_format == types.DataFormat.TFRECORD: from dpsynth.dataset_descriptors import tfrecord_descriptor - attributes_dict = { attr.name: attr.data_type for attr in dataset_desc.attributes } diff --git a/dpsynth/pipeline_transformations/input_output.py b/dpsynth/pipeline_transformations/input_output.py index 6c5e497..ab77e8f 100644 --- a/dpsynth/pipeline_transformations/input_output.py +++ b/dpsynth/pipeline_transformations/input_output.py @@ -81,7 +81,6 @@ def save_data_local( save_csv(data, path, attributes) # pytype: disable=wrong-arg-types case types.DataFormat.TFRECORD: import tensorflow as tf - os.makedirs(os.path.dirname(path), exist_ok=True) with tf.io.TFRecordWriter(path) as writer: for record in data: @@ -172,7 +171,6 @@ def load_data_for_beam( return load_csv(pipeline, path) case types.DataFormat.TFRECORD: import tensorflow as tf - return pipeline | 'ReadTFRecord' >> beam.io.ReadFromTFRecord( path, coder=beam.coders.ProtoCoder(tf.train.Example) ) @@ -190,7 +188,6 @@ def save_beam_data( match data_format: case types.DataFormat.TFRECORD: import tensorflow as tf - _ = data | 'WriteTFRecord' >> beam.io.WriteToTFRecord( path, coder=beam.coders.ProtoCoder(tf.train.Example) )