black

autogluon · Feb 6, 2023 · 355ebfe · 355ebfe
1 parent 71f84fc
commit 355ebfe
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 41 deletions.
diff --git a/src/autogluon/cloud/predictor/timeseries_cloud_predictor.py b/src/autogluon/cloud/predictor/timeseries_cloud_predictor.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
+
 import copy
 import logging
-from typing import Optional, Union, Dict, Any
+from typing import Any, Dict, Optional, Union
 
 import pandas as pd
 
@@ -27,7 +28,7 @@ def _get_local_predictor_cls(self):
 
         predictor_cls = TimeSeriesPredictor
         return predictor_cls
-    
+
     def _preprocess_data(
         self,
         data,
@@ -50,15 +51,15 @@ def _preprocess_data(
         target_index = cols.index(target)
         cols.append(cols.pop(target_index))
         data = data[cols]
-        
+
         if static_features is not None:
             # Merge static features so only one dataframe needs to be sent to remote container
             if isinstance(static_features, str):
                 static_features = load_pd.load(static_features)
             data = pd.merge(data, static_features, how="left", on=id_column)
-        
+
         return data
-    
+
     def fit(
         self,
         *,
@@ -135,15 +136,15 @@ def fit(
             id_column=id_column,
             timestamp_column=timestamp_column,
             target=target,
-            static_features=static_features
+            static_features=static_features,
         )
         if tuning_data is not None:
             tuning_data = self._preprocess_data(
                 data=tuning_data,
                 id_column=id_column,
                 timestamp_column=timestamp_column,
                 target=target,
-                static_features=static_features
+                static_features=static_features,
             )
         predictor_fit_args["train_data"] = train_data
         predictor_fit_args["tuning_data"] = tuning_data
@@ -161,7 +162,7 @@ def fit(
             autogluon_sagemaker_estimator_kwargs=autogluon_sagemaker_estimator_kwargs,
             **kwargs,
         )
-        
+
     def predict_real_time(
         self,
         test_data: Union[str, pd.DataFrame],
@@ -206,17 +207,14 @@ def predict_real_time(
             id_column=id_column,
             timestamp_column=timestamp_column,
             target=target,
-            static_features=static_features
+            static_features=static_features,
         )
         pred, _ = self._predict_real_time(test_data=test_data, accept=accept)
         return pred
-
-    def predict_proba_real_time(
-        self,
-        **kwargs
-    ) -> pd.DataFrame:
+
+    def predict_proba_real_time(self, **kwargs) -> pd.DataFrame:
         raise ValueError(f"{self.__class__.__name__} does not support predict_proba operation.")
-    
+
     def predict(
         self,
         test_data: Union[str, pd.DataFrame],
@@ -256,13 +254,13 @@ def predict(
             id_column=id_column,
             timestamp_column=timestamp_column,
             target=target,
-            static_features=static_features
+            static_features=static_features,
         )
         return super().predict(
             test_data,
             **kwargs,
         )
-        
+
     def predict_proba(
         self,
         **kwargs,

diff --git a/src/autogluon/cloud/scripts/script_manager.py b/src/autogluon/cloud/scripts/script_manager.py
@@ -12,7 +12,7 @@ class ScriptManager:
     _SERVE_SCRIPT_MAP = dict(
         tabular=TABULAR_SERVE_SCRIPT_PATH,
         multimodal=MULTIMODAL_SERVE_SCRIPT_PATH,
-        timeseries=TIMESERIES_SERVE_SCRIPT_PATH
+        timeseries=TIMESERIES_SERVE_SCRIPT_PATH,
     )
 
     @classmethod

diff --git a/src/autogluon/cloud/scripts/timeseries_serve.py b/src/autogluon/cloud/scripts/timeseries_serve.py
@@ -6,7 +6,6 @@
 from autogluon.timeseries import TimeSeriesDataFrame, TimeSeriesPredictor
 
 
-
 def model_fn(model_dir):
     """loads model from previously saved artifact"""
     model = TimeSeriesPredictor.load(model_dir)
@@ -23,15 +22,11 @@ def prepare_timeseries_dataframe(df, predictor):
     if target != cols[-1]:
         # target is not the last column, then there are static features being merged in
         target_index = cols.index(target)
-        static_columns = cols[target_index+1:]
-        static_features = df[[id_column]+static_columns].groupby([id_column], sort=False).head(1)
+        static_columns = cols[target_index + 1 :]
+        static_features = df[[id_column] + static_columns].groupby([id_column], sort=False).head(1)
         static_features.set_index(id_column, inplace=True)
         df.drop(columns=static_columns, inplace=True)
-    df = TimeSeriesDataFrame.from_data_frame(
-        df,
-        id_column=id_column,
-        timestamp_column=timestamp_column
-    )
+    df = TimeSeriesDataFrame.from_data_frame(df, id_column=id_column, timestamp_column=timestamp_column)
     if static_features is not None:
         print(static_features)
         df.static_features = static_features

diff --git a/src/autogluon/cloud/scripts/train.py b/src/autogluon/cloud/scripts/train.py
@@ -31,7 +31,7 @@ def get_env_if_present(name):
 
 
 def prepare_timeseries_dataframe(df, predictor_init_args):
-    target = predictor_init_args['target']
+    target = predictor_init_args["target"]
     cols = df.columns.to_list()
     id_column = cols[0]
     timestamp_column = cols[1]
@@ -40,15 +40,11 @@ def prepare_timeseries_dataframe(df, predictor_init_args):
     if target != cols[-1]:
         # target is not the last column, then there are static features being merged in
         target_index = cols.index(target)
-        static_columns = cols[target_index+1:]
-        static_features = df[[id_column]+static_columns].groupby([id_column], sort=False).head(1)
+        static_columns = cols[target_index + 1 :]
+        static_features = df[[id_column] + static_columns].groupby([id_column], sort=False).head(1)
         static_features.set_index(id_column, inplace=True)
         df.drop(columns=static_columns, inplace=True)
-    df = TimeSeriesDataFrame.from_data_frame(
-        df,
-        id_column=id_column,
-        timestamp_column=timestamp_column
-    )
+    df = TimeSeriesDataFrame.from_data_frame(df, id_column=id_column, timestamp_column=timestamp_column)
     if static_features is not None:
         print(static_features)
         df.static_features = static_features
@@ -120,9 +116,11 @@ def prepare_data(data_file, predictor_type, predictor_init_args):
             predictor_fit_args["feature_meatadata"] = FeatureMetadata(**predictor_fit_args["feature_meatadata"])
     elif predictor_type == "multimodal":
         from autogluon.multimodal import MultiModalPredictor
+
         predictor_cls = MultiModalPredictor
     elif predictor_type == "timeseries":
         from autogluon.timeseries import TimeSeriesPredictor, TimeSeriesDataFrame
+
         predictor_cls = TimeSeriesPredictor
 
     train_file = get_input_path(args.train_dir)

diff --git a/tests/unittests/timeseries/test_timeseries.py b/tests/unittests/timeseries/test_timeseries.py
@@ -7,18 +7,16 @@
 def test_timeseries(test_helper, framework_version):
     train_data = "timeseries_train.csv"
     static_features = "timeseries_static_features.csv"
-    id_column="item_id"
-    timestamp_column="timestamp"
-    target="target"
+    id_column = "item_id"
+    timestamp_column = "timestamp"
+    target = "target"
     timestamp = test_helper.get_utc_timestamp_now()
     with tempfile.TemporaryDirectory() as temp_dir:
         os.chdir(temp_dir)
         test_helper.prepare_data(train_data, static_features)
         time_limit = 60
 
-        predictor_init_args = dict(
-            target=target
-        )
+        predictor_init_args = dict(target=target)
 
         predictor_fit_args = dict(
             train_data=train_data,
@@ -50,7 +48,7 @@ def test_timeseries(test_helper, framework_version):
                 target=target,
                 static_features=static_features,
                 framework_version=framework_version,
-                custom_image_uri=inference_custom_image_uri
+                custom_image_uri=inference_custom_image_uri,
             ),
             predict_real_time_kwargs=dict(
                 id_column=id_column,