runnable version

Author: CRZbulabula

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

@@ -209,9 +209,11 @@ def _run(
                 inference_pipeline = load_pipeline(model_info, device="cpu")
                 inputs = inference_pipeline.preprocess(inputs)
                 if isinstance(inference_pipeline, ForecastPipeline):
+                    inputs = inference_pipeline._preprocess(inputs)
                     outputs = inference_pipeline.forecast(
                         inputs, predict_length=output_length, **inference_attrs
                     )
+                    outputs = inference_pipeline._postprocess(outputs)
                 elif isinstance(inference_pipeline, ClassificationPipeline):
                     outputs = inference_pipeline.classify(inputs)
                 elif isinstance(inference_pipeline, ChatPipeline):

iotdb-core/ainode/iotdb/ainode/core/model/chronos2/pipeline_chronos2.py

@@ -42,6 +42,7 @@ def _preprocess(self, inputs):
             raise InferenceModelInternalException(
                 f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
             )
+        inputs = inputs.unsqueeze(0)
         return inputs
 
     @property
@@ -239,7 +240,7 @@ def forecast(self, inputs, **infer_kwargs):
             )
             logger.warning(msg)
 
-        context_length = len(inputs)
+        context_length = inputs.shape[-1]
         if context_length > self.model_context_length:
             logger.warning(
                 f"The specified context_length {context_length} is greater than the model's default context length {self.model_context_length}. "
@@ -391,4 +392,4 @@ def _predict_step(
         return prediction
 
     def _postprocess(self, output: torch.Tensor):
-        return output
+        return output[0].mean(dim=1)

iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py

@@ -43,7 +43,7 @@
 
 
 def load_model(model_info: ModelInfo, **model_kwargs) -> Any:
-    if model_info.auto_map is not None:
+    if model_info.category == ModelCategory.BUILTIN or model_info.auto_map is not None:
         model = load_model_from_transformers(model_info, **model_kwargs)
     else:
         if model_info.model_type == "sktime":