add tfrecord_features

NREL · Sep 16, 2021 · eef032d · eef032d
1 parent 061c0ad
commit eef032d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/nfp/preprocessing/preprocessor.py b/nfp/preprocessing/preprocessor.py
@@ -179,6 +179,12 @@ def padding_values(self) -> {str: tf.constant}:
         return {key: tf.constant(0, dtype=self.output_dtype)
                 for key in self.output_signature.keys()}
 
+    @property
+    def tfrecord_features(self) -> {str: tf.io.FixedLenFeature}:
+        """For loading preprocessed inputs from a tf records file"""
+        return {key: tf.io.FixedLenFeature([], dtype=self.output_dtype if len(val.shape) == 0 else tf.string)
+                for key, val in self.output_signature.items()}
+
 
 def load_from_json(obj, data):
     for key, val in obj.__dict__.items():