Merge branch 'add-csv-deserializer' of github.com:bveeramani/sagemaker-python-sdk into add-csv-deserializer

Author: Balaji Veeramani

src/sagemaker/amazon/common.py

@@ -21,6 +21,7 @@
 import numpy as np
 
 from sagemaker.amazon.record_pb2 import Record
+from sagemaker.deserializers import BaseDeserializer
 from sagemaker.utils import DeferredError
 
 
@@ -48,26 +49,24 @@ def __call__(self, array):
         return buf
 
 
-class record_deserializer(object):
-    """Placeholder docstring"""
+class RecordDeserializer(BaseDeserializer):
+    """Deserialize RecordIO Protobuf data from an inference endpoint."""
 
-    def __init__(self, accept="application/x-recordio-protobuf"):
-        """
-        Args:
-            accept:
-        """
-        self.accept = accept
+    ACCEPT = "application/x-recordio-protobuf"
+
+    def deserialize(self, data, content_type):
+        """Deserialize RecordIO Protobuf data from an inference endpoint.
 
-    def __call__(self, stream, content_type):
-        """
         Args:
-            stream:
-            content_type:
+            data (object): The protobuf message to deserialize.
+            content_type (str): The MIME type of the data.
+        Returns:
+            list: A list of records.
         """
         try:
-            return read_records(stream)
+            return read_records(data)
         finally:
-            stream.close()
+            data.close()
 
 
 def _write_feature_tensor(resolved_type, record, vector):

src/sagemaker/amazon/factorization_machines.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge
 from sagemaker.predictor import Predictor
@@ -290,7 +290,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/kmeans.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin, ge, le
 from sagemaker.predictor import Predictor
@@ -223,7 +223,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/knn.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, isin
 from sagemaker.predictor import Predictor
@@ -211,7 +211,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/lda.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt
 from sagemaker.predictor import Predictor
@@ -195,7 +195,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/linear_learner.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import isin, gt, lt, ge, le
 from sagemaker.predictor import Predictor
@@ -454,7 +454,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/ntm.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le, isin
 from sagemaker.predictor import Predictor
@@ -225,7 +225,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/pca.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import gt, isin
 from sagemaker.predictor import Predictor
@@ -207,7 +207,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

src/sagemaker/amazon/randomcutforest.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.amazon.amazon_estimator import AmazonAlgorithmEstimatorBase, registry
-from sagemaker.amazon.common import numpy_to_record_serializer, record_deserializer
+from sagemaker.amazon.common import numpy_to_record_serializer, RecordDeserializer
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.amazon.validation import ge, le
 from sagemaker.predictor import Predictor
@@ -184,7 +184,7 @@ def __init__(self, endpoint_name, sagemaker_session=None):
             endpoint_name,
             sagemaker_session,
             serializer=numpy_to_record_serializer(),
-            deserializer=record_deserializer(),
+            deserializer=RecordDeserializer(),
         )
 
 

tests/unit/test_common.py

@@ -18,7 +18,7 @@
 import itertools
 from scipy.sparse import coo_matrix
 from sagemaker.amazon.common import (
-    record_deserializer,
+    RecordDeserializer,
     write_numpy_to_dense_tensor,
     read_recordio,
     numpy_to_record_serializer,
@@ -51,8 +51,8 @@ def test_deserializer():
     array_data = [[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]
     s = numpy_to_record_serializer()
     buf = s(np.array(array_data))
-    d = record_deserializer()
-    for record, expected in zip(d(buf, "who cares"), array_data):
+    d = RecordDeserializer()
+    for record, expected in zip(d.deserialize(buf, "who cares"), array_data):
         assert record.features["values"].float64_tensor.values == expected