support various pipeline Interfaces and support arima with sktime package (#16861)

Author: RkGrit

iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py

@@ -30,6 +30,7 @@
 from iotdb.ainode.core.constant import INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE
 from iotdb.ainode.core.inference.batcher.basic_batcher import BasicBatcher
 from iotdb.ainode.core.inference.inference_request import InferenceRequest
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline, ClassificationPipeline, ChatPipeline
 from iotdb.ainode.core.inference.pipeline.pipeline_loader import load_pipeline
 from iotdb.ainode.core.inference.request_scheduler.basic_request_scheduler import (
     BasicRequestScheduler,
@@ -116,11 +117,24 @@ def _step(self):
 
         for requests in grouped_requests:
             batch_inputs = self._batcher.batch_request(requests).to(self.device)
-            batch_output = self._inference_pipeline.infer(
-                batch_inputs,
-                predict_length=requests[0].max_new_tokens,
-                revin=True,
-            )
+            if isinstance(self._inference_pipeline, ForecastPipeline):
+                batch_output = self._inference_pipeline.forecast(
+                    batch_inputs,
+                    predict_length=requests[0].max_new_tokens,
+                    revin=True,
+                )
+            elif isinstance(self._inference_pipeline, ClassificationPipeline):
+                batch_output = self._inference_pipeline.classify(
+                    batch_inputs,
+                    # more infer kwargs can be added here
+                )
+            elif isinstance(self._inference_pipeline, ChatPipeline):
+                batch_output = self._inference_pipeline.chat(
+                    batch_inputs,
+                    # more infer kwargs can be added here
+                )
+            else:
+                self._logger.error("[Inference] Unsupported pipeline type.")
             offset = 0
             for request in requests:
                 request.output_tensor = request.output_tensor.to(self.device)

iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/basic_pipeline.py

@@ -16,11 +16,10 @@
 # under the License.
 #
 
-from abc import ABC
+from abc import ABC, abstractmethod
 
 import torch
 
-from iotdb.ainode.core.exception import InferenceModelInternalError
 from iotdb.ainode.core.model.model_loader import load_model
 
 
@@ -48,12 +47,9 @@ def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
     def _preprocess(self, inputs):
-        if len(inputs.shape) != 2:
-            raise InferenceModelInternalError(
-                f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
-            )
         return inputs
 
+    @abstractmethod
     def forecast(self, inputs, **infer_kwargs):
         pass
 
@@ -68,6 +64,7 @@ def __init__(self, model_info, **model_kwargs):
     def _preprocess(self, inputs):
         return inputs
 
+    @abstractmethod
     def classify(self, inputs, **kwargs):
         pass
 
@@ -82,6 +79,7 @@ def __init__(self, model_info, **model_kwargs):
     def _preprocess(self, inputs):
         return inputs
 
+    @abstractmethod
     def chat(self, inputs, **kwargs):
         pass
 

iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/pipeline_loader.py

@@ -31,6 +31,7 @@
 def load_pipeline(model_info: ModelInfo, device: str, **model_kwargs):
     if model_info.model_type == "sktime":
         from iotdb.ainode.core.model.sktime.pipeline_sktime import SktimePipeline
+
         pipeline_cls = SktimePipeline
     elif model_info.category == ModelCategory.BUILTIN:
         module_name = (

iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py

@@ -35,6 +35,7 @@
     InferenceRequestProxy,
 )
 from iotdb.ainode.core.inference.pipeline.pipeline_loader import load_pipeline
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline, ClassificationPipeline, ChatPipeline
 from iotdb.ainode.core.inference.pool_controller import PoolController
 from iotdb.ainode.core.inference.utils import generate_req_id
 from iotdb.ainode.core.log import Logger
@@ -210,9 +211,16 @@ def _run(
             else:
                 model_info = self._model_manager.get_model_info(model_id)
                 inference_pipeline = load_pipeline(model_info, device="cpu")
-                outputs = inference_pipeline.infer(
-                    inputs, predict_length=predict_length, **inference_attrs
-                )
+                if isinstance(inference_pipeline, ForecastPipeline):
+                    outputs = inference_pipeline.forecast(
+                        inputs, predict_length=predict_length, **inference_attrs
+                    )
+                elif isinstance(inference_pipeline, ClassificationPipeline):
+                    outputs = inference_pipeline.classify(inputs)
+                elif isinstance(inference_pipeline, ChatPipeline):
+                    outputs = inference_pipeline.chat(inputs)
+                else:
+                    logger.error("[Inference] Unsupported pipeline type.")
                 outputs = convert_to_binary(pd.DataFrame(outputs[0]))
 
             # construct response

iotdb-core/ainode/iotdb/ainode/core/manager/utils.py

@@ -91,7 +91,9 @@ def estimate_pool_size(device: torch.device, model_id: str) -> int:
     system_res = evaluate_system_resources(device)
     free_mem = system_res["free_mem"]
 
-    mem_usage = MODEL_MEM_USAGE_MAP[model_info.model_type] * INFERENCE_EXTRA_MEMORY_RATIO
+    mem_usage = (
+        MODEL_MEM_USAGE_MAP[model_info.model_type] * INFERENCE_EXTRA_MEMORY_RATIO
+    )
     size = int((free_mem * INFERENCE_MEMORY_USAGE_RATIO) // mem_usage)
     if size <= 0:
         logger.error(

iotdb-core/ainode/iotdb/ainode/core/model/sktime/arima/config.json

@@ -3,11 +3,23 @@
   "model_id": "arima",
   "predict_length": 1,
   "order": [1, 0, 0],
-  "season_length": 1,
-  "seasonal_order": [0, 0, 0],
-  "include_mean": true,
-  "include_drift": false,
-  "biasadj": false,
-  "method": "CSS-ML"
+  "seasonal_order": [0, 0, 0, 0],
+  "start_params": null,
+  "method": "lbfgs",
+  "maxiter": 50,
+  "suppress_warnings": false,
+  "out_of_sample_size": 0,
+  "scoring": "mse",
+  "scoring_args": null,
+  "trend": null,
+  "with_intercept": true,
+  "time_varying_regression": false,
+  "enforce_stationarity": true,
+  "enforce_invertibility": true,
+  "simple_differencing": false,
+  "measurement_error": false,
+  "mle_regression": true,
+  "hamilton_representation": false,
+  "concentrate_scale": false
 }
 

iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py

@@ -176,21 +176,40 @@ def parse(self, string_value: str):
     "ARIMA": {
         "predict_length": AttributeConfig("predict_length", 1, "int", 1, 5000),
         "order": AttributeConfig("order", (1, 0, 0), "tuple", value_type=int),
-        "season_length": AttributeConfig("season_length", 1, "int", 1, 5000),
         "seasonal_order": AttributeConfig(
-            "seasonal_order", (0, 0, 0), "tuple", value_type=int
+            "seasonal_order", (0, 0, 0, 0), "tuple", value_type=int
         ),
-        "include_mean": AttributeConfig("include_mean", True, "bool"),
-        "include_drift": AttributeConfig("include_drift", False, "bool"),
-        "include_constant": AttributeConfig("include_constant", None, "bool"),
-        "blambda": AttributeConfig("blambda", None, "float"),
-        "biasadj": AttributeConfig("biasadj", False, "bool"),
+        "start_params": AttributeConfig("start_params", None, "str"),
         "method": AttributeConfig(
             "method",
-            "CSS-ML",
+            "lbfgs",
             "str",
-            choices=["CSS-ML", "ML", "CSS"],
+            choices=["lbfgs", "bfgs", "newton", "nm", "cg", "ncg", "powell"],
         ),
+        "maxiter": AttributeConfig("maxiter", 50, "int", 1, 5000),
+        "suppress_warnings": AttributeConfig("suppress_warnings", False, "bool"),
+        "out_of_sample_size": AttributeConfig("out_of_sample_size", 0, "int", 0, 5000),
+        "scoring": AttributeConfig(
+            "scoring",
+            "mse",
+            "str",
+            choices=["mse", "mae", "rmse", "mape", "smape", "rmsle", "r2"],
+        ),
+        "scoring_args": AttributeConfig("scoring_args", None, "str"),
+        "trend": AttributeConfig("trend", None, "str"),
+        "with_intercept": AttributeConfig("with_intercept", True, "bool"),
+        "time_varying_regression": AttributeConfig(
+            "time_varying_regression", False, "bool"
+        ),
+        "enforce_stationarity": AttributeConfig("enforce_stationarity", True, "bool"),
+        "enforce_invertibility": AttributeConfig("enforce_invertibility", True, "bool"),
+        "simple_differencing": AttributeConfig("simple_differencing", False, "bool"),
+        "measurement_error": AttributeConfig("measurement_error", False, "bool"),
+        "mle_regression": AttributeConfig("mle_regression", True, "bool"),
+        "hamilton_representation": AttributeConfig(
+            "hamilton_representation", False, "bool"
+        ),
+        "concentrate_scale": AttributeConfig("concentrate_scale", False, "bool"),
     },
     "STL_FORECASTER": {
         "predict_length": AttributeConfig("predict_length", 1, "int", 1, 5000),

iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py

@@ -24,7 +24,7 @@
 from sklearn.preprocessing import MinMaxScaler
 from sktime.detection.hmm_learn import GMMHMM, GaussianHMM
 from sktime.detection.stray import STRAY
-from statsforecast.models import ARIMA
+from sktime.forecasting.arima import ARIMA
 from sktime.forecasting.exp_smoothing import ExponentialSmoothing
 from sktime.forecasting.naive import NaiveForecaster
 from sktime.forecasting.trend import STLForecaster
@@ -59,12 +59,11 @@ class ForecastingModel(SktimeModel):
     def generate(self, data, **kwargs):
         """Execute forecasting"""
         try:
-            predict_length = kwargs.get("predict_length", self._attributes["predict_length"])
+            predict_length = kwargs.get(
+                "predict_length", self._attributes["predict_length"]
+            )
             self._model.fit(data)
-            if isinstance(self._model, ARIMA):
-                output = self._model.predict(h=predict_length)['mean']
-            else:
-                output = self._model.predict(fh=range(predict_length))
+            output = self._model.predict(fh=range(predict_length))
             return np.array(output, dtype=np.float64)
         except Exception as e:
             raise InferenceModelInternalError(str(e))
@@ -92,7 +91,7 @@ class ArimaModel(ForecastingModel):
     def __init__(self, attributes: Dict[str, Any]):
         super().__init__(attributes)
         self._model = ARIMA(
-            **{k: v for k, v in attributes.items() if k != "predict_length" and v is not None}
+            **{k: v for k, v in attributes.items() if k != "predict_length"}
         )
 
 
@@ -147,9 +146,7 @@ class STRAYModel(DetectionModel):
 
     def __init__(self, attributes: Dict[str, Any]):
         super().__init__(attributes)
-        self._model = STRAY(
-            **{k: v for k, v in attributes.items() if v is not None}
-        )
+        self._model = STRAY(**{k: v for k, v in attributes.items() if v is not None})
 
     def generate(self, data, **kwargs):
         """STRAY requires special handling: normalize first"""

iotdb-core/ainode/iotdb/ainode/core/model/sktime/pipeline_sktime.py

@@ -20,18 +20,18 @@
 import pandas as pd
 import torch
 
-from iotdb.ainode.core.inference.pipeline.basic_pipeline import BasicPipeline
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
 
 
-class SktimePipeline(BasicPipeline):
+class SktimePipeline(ForecastPipeline):
     def __init__(self, model_info, **model_kwargs):
+        model_kwargs.pop("device", None)  # sktime models run on CPU
         super().__init__(model_info, model_kwargs=model_kwargs)
-        model_kwargs.pop("device", None)
 
     def _preprocess(self, inputs):
-        return super()._preprocess(inputs)
+        return inputs
 
-    def infer(self, inputs, **infer_kwargs):
+    def forecast(self, inputs, **infer_kwargs):
         predict_length = infer_kwargs.get("predict_length", 96)
         input_ids = self._preprocess(inputs)
 

iotdb-core/ainode/iotdb/ainode/core/model/sundial/pipeline_sundial.py

@@ -19,16 +19,21 @@
 import torch
 
 from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
+from iotdb.ainode.core.exception import InferenceModelInternalError
 
 
 class SundialPipeline(ForecastPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
     def _preprocess(self, inputs):
-        return super()._preprocess(inputs)
+        if len(inputs.shape) != 2:
+            raise InferenceModelInternalError(
+                f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
+            )
+        return inputs
 
-    def infer(self, inputs, **infer_kwargs):
+    def forecast(self, inputs, **infer_kwargs):
         predict_length = infer_kwargs.get("predict_length", 96)
         num_samples = infer_kwargs.get("num_samples", 10)
         revin = infer_kwargs.get("revin", True)