Expose the processing interface of the pipeline

Author: RkGrit

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

@@ -123,6 +123,7 @@ def _step(self):
             batch_inputs = self._batcher.batch_request(requests).to(
                 "cpu"
             )  # The input data should first load to CPU in current version
+            batch_inputs = self._inference_pipeline.preprocess(batch_inputs)
             if isinstance(self._inference_pipeline, ForecastPipeline):
                 batch_output = self._inference_pipeline.forecast(
                     batch_inputs,
@@ -140,7 +141,10 @@ def _step(self):
                     # more infer kwargs can be added here
                 )
             else:
+                batch_output = None
                 self._logger.error("[Inference] Unsupported pipeline type.")
+            batch_output = self._inference_pipeline.postprocess(batch_output)
+
             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

@@ -29,59 +29,61 @@ def __init__(self, model_info, **model_kwargs):
         self.device = model_kwargs.get("device", "cpu")
         self.model = load_model(model_info, device_map=self.device, **model_kwargs)
 
-    def _preprocess(self, inputs):
+    @abstractmethod
+    def preprocess(self, inputs):
         """
         Preprocess the input before inference, including shape validation and value transformation.
         """
-        return inputs
+        raise NotImplementedError("preprocess not implemented")
 
-    def _postprocess(self, output: torch.Tensor):
+    @abstractmethod
+    def postprocess(self, output: torch.Tensor):
         """
         Post-process the outputs after the entire inference task.
         """
-        return output
+        raise NotImplementedError("postprocess not implemented")
 
 
 class ForecastPipeline(BasicPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         return inputs
 
     @abstractmethod
     def forecast(self, inputs, **infer_kwargs):
         pass
 
-    def _postprocess(self, output: torch.Tensor):
+    def postprocess(self, output: torch.Tensor):
         return output
 
 
 class ClassificationPipeline(BasicPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         return inputs
 
     @abstractmethod
     def classify(self, inputs, **kwargs):
         pass
 
-    def _postprocess(self, output: torch.Tensor):
+    def postprocess(self, output: torch.Tensor):
         return output
 
 
 class ChatPipeline(BasicPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         return inputs
 
     @abstractmethod
     def chat(self, inputs, **kwargs):
         pass
 
-    def _postprocess(self, output: torch.Tensor):
+    def postprocess(self, output: torch.Tensor):
         return output

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

@@ -215,6 +215,7 @@ def _run(
             else:
                 model_info = self._model_manager.get_model_info(model_id)
                 inference_pipeline = load_pipeline(model_info, device="cpu")
+                inputs = inference_pipeline.preprocess(inputs)
                 if isinstance(inference_pipeline, ForecastPipeline):
                     outputs = inference_pipeline.forecast(
                         inputs, predict_length=output_length, **inference_attrs
@@ -224,7 +225,9 @@ def _run(
                 elif isinstance(inference_pipeline, ChatPipeline):
                     outputs = inference_pipeline.chat(inputs)
                 else:
+                    outputs = None
                     logger.error("[Inference] Unsupported pipeline type.")
+                outputs = inference_pipeline.postprocess(outputs)
                 outputs = convert_to_binary(pd.DataFrame(outputs[0]))
 
             # construct response

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

@@ -28,41 +28,40 @@ 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)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         return inputs
 
     def forecast(self, inputs, **infer_kwargs):
         predict_length = infer_kwargs.get("predict_length", 96)
-        input_ids = self._preprocess(inputs)
 
         # Convert to pandas Series for sktime (sktime expects Series or DataFrame)
         # Handle batch dimension: if batch_size > 1, process each sample separately
-        if len(input_ids.shape) == 2 and input_ids.shape[0] > 1:
+        if len(inputs.shape) == 2 and inputs.shape[0] > 1:
             # Batch processing: convert each row to Series
             outputs = []
-            for i in range(input_ids.shape[0]):
+            for i in range(inputs.shape[0]):
                 series = pd.Series(
-                    input_ids[i].cpu().numpy()
-                    if isinstance(input_ids, torch.Tensor)
-                    else input_ids[i]
+                    inputs[i].cpu().numpy()
+                    if isinstance(inputs, torch.Tensor)
+                    else inputs[i]
                 )
                 output = self.model.generate(series, predict_length=predict_length)
                 outputs.append(output)
             output = np.array(outputs)
         else:
             # Single sample: convert to Series
-            if isinstance(input_ids, torch.Tensor):
-                series = pd.Series(input_ids.squeeze().cpu().numpy())
+            if isinstance(inputs, torch.Tensor):
+                series = pd.Series(inputs.squeeze().cpu().numpy())
             else:
-                series = pd.Series(input_ids.squeeze())
+                series = pd.Series(inputs.squeeze())
             output = self.model.generate(series, predict_length=predict_length)
             # Add batch dimension if needed
             if len(output.shape) == 1:
                 output = output[np.newaxis, :]
 
-        return self._postprocess(output)
+        return output
 
-    def _postprocess(self, output):
+    def postprocess(self, output):
         if isinstance(output, np.ndarray):
             return torch.from_numpy(output).float()
         return output

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

@@ -26,7 +26,7 @@ class SundialPipeline(ForecastPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         if len(inputs.shape) != 2:
             raise InferenceModelInternalException(
                 f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
@@ -38,14 +38,13 @@ def forecast(self, inputs, **infer_kwargs):
         num_samples = infer_kwargs.get("num_samples", 10)
         revin = infer_kwargs.get("revin", True)
 
-        input_ids = self._preprocess(inputs)
         output = self.model.generate(
-            input_ids,
+            inputs,
             max_new_tokens=predict_length,
             num_samples=num_samples,
             revin=revin,
         )
-        return self._postprocess(output)
+        return output
 
-    def _postprocess(self, output: torch.Tensor):
+    def postprocess(self, output: torch.Tensor):
         return output.mean(dim=1)

iotdb-core/ainode/iotdb/ainode/core/model/timer_xl/pipeline_timer.py

@@ -26,7 +26,7 @@ class TimerPipeline(ForecastPipeline):
     def __init__(self, model_info, **model_kwargs):
         super().__init__(model_info, model_kwargs=model_kwargs)
 
-    def _preprocess(self, inputs):
+    def preprocess(self, inputs):
         if len(inputs.shape) != 2:
             raise InferenceModelInternalException(
                 f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
@@ -37,11 +37,10 @@ def forecast(self, inputs, **infer_kwargs):
         predict_length = infer_kwargs.get("predict_length", 96)
         revin = infer_kwargs.get("revin", True)
 
-        input_ids = self._preprocess(inputs)
         output = self.model.generate(
-            input_ids, max_new_tokens=predict_length, revin=revin
+            inputs, max_new_tokens=predict_length, revin=revin
         )
-        return self._postprocess(output)
+        return output
 
-    def _postprocess(self, output: torch.Tensor):
+    def postprocess(self, output: torch.Tensor):
         return output