You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by yo...@apache.org on 2025/11/26 09:40:35 UTC
(iotdb) 03/04: Reconstruct model management and model loading
This is an automated email from the ASF dual-hosted git repository.
yongzao pushed a commit to branch model-management
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 3448a16859605d18a1a847fcc51886b67802ea77
Author: RkGrit <G1...@gmail.com>
AuthorDate: Mon Nov 24 21:40:59 2025 +0800
Reconstruct model management and model loading
---
iotdb-core/ainode/iotdb/ainode/core/config.py | 16 +
iotdb-core/ainode/iotdb/ainode/core/constant.py | 143 +----
.../ainode/core/inference/inference_request.py | 6 +-
.../core/inference/inference_request_pool.py | 110 ++--
.../inference/{strategy => pipeline}/__init__.py | 12 +
.../core/inference/pipeline/basic_pipeline.py | 111 ++++
.../core/inference/pipeline/sktime_pipeline.py | 62 ++
.../core/inference/pipeline/sundial_pipeline.py | 48 ++
.../core/inference/pipeline/timerxl_pipeline.py | 44 ++
.../iotdb/ainode/core/inference/pool_controller.py | 26 +-
.../pool_scheduler/basic_pool_scheduler.py | 7 +-
.../strategy/abstract_inference_pipeline.py | 60 --
.../strategy/timer_sundial_inference_pipeline.py | 51 --
.../strategy/timerxl_inference_pipeline.py | 51 --
.../iotdb/ainode/core/manager/inference_manager.py | 141 +----
.../iotdb/ainode/core/manager/model_manager.py | 121 ++++
.../ainode/iotdb/ainode/core/manager/utils.py | 13 +-
.../core/{inference/strategy => model}/__init__.py | 0
.../ainode/iotdb/ainode/core/model/model_enums.py | 58 ++
.../ainode/iotdb/ainode/core/model/model_info.py | 129 ++++
.../ainode/iotdb/ainode/core/model/model_loader.py | 170 ++++++
.../iotdb/ainode/core/model/model_storage.py | 576 ++++++++++++++++++
.../strategy => model/sktime}/__init__.py | 0
.../ainode/core/model/sktime/arima/config.json | 22 +
.../core/model/sktime/configuration_sktime.py | 379 ++++++++++++
.../model/sktime/exponential_smoothing/config.json | 11 +
.../core/model/sktime/gaussian_hmm/config.json | 20 +
.../ainode/core/model/sktime/gmm_hmm/config.json | 20 +
.../ainode/core/model/sktime/modeling_sktime.py | 178 ++++++
.../core/model/sktime/naive_forecaster/config.json | 8 +
.../core/model/sktime/stl_forecaster/config.json | 14 +
.../ainode/core/model/sktime/stray/config.json | 11 +
.../strategy => model/sundial}/__init__.py | 0
.../core/model/sundial/configuration_sundial.py | 65 ++
.../iotdb/ainode/core/model/sundial/flow_loss.py | 255 ++++++++
.../ainode/core/model/sundial/modeling_sundial.py | 651 +++++++++++++++++++++
.../core/model/sundial/ts_generation_mixin.py | 383 ++++++++++++
.../strategy => model/timerxl}/__init__.py | 0
.../core/model/timerxl/configuration_timer.py | 59 ++
.../ainode/core/model/timerxl/modeling_timer.py | 640 ++++++++++++++++++++
.../core/model/timerxl/ts_generation_mixin.py | 370 ++++++++++++
iotdb-core/ainode/iotdb/ainode/core/model/utils.py | 94 +++
iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py | 105 ++--
43 files changed, 4654 insertions(+), 586 deletions(-)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/config.py b/iotdb-core/ainode/iotdb/ainode/core/config.py
index afcf0683d7d..04ec3ee68c1 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/config.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/config.py
@@ -31,6 +31,7 @@ from iotdb.ainode.core.constant import (
AINODE_CONF_FILE_NAME,
AINODE_CONF_GIT_FILE_NAME,
AINODE_CONF_POM_FILE_NAME,
+ AINODE_FINETUNE_MODELS_DIR,
AINODE_INFERENCE_BATCH_INTERVAL_IN_MS,
AINODE_INFERENCE_EXTRA_MEMORY_RATIO,
AINODE_INFERENCE_MAX_PREDICT_LENGTH,
@@ -44,6 +45,7 @@ from iotdb.ainode.core.constant import (
AINODE_SYSTEM_FILE_NAME,
AINODE_TARGET_CONFIG_NODE_LIST,
AINODE_THRIFT_COMPRESSION_ENABLED,
+ AINODE_USER_DEFINED_MODELS_DIR,
AINODE_VERSION_INFO,
)
from iotdb.ainode.core.exception import BadNodeUrlError
@@ -96,6 +98,8 @@ class AINodeConfig(object):
# Directory to save models
self._ain_models_dir = AINODE_MODELS_DIR
self._ain_builtin_models_dir = AINODE_BUILTIN_MODELS_DIR
+ self._ain_finetune_models_dir = AINODE_FINETUNE_MODELS_DIR
+ self._ain_user_defined_models_dir = AINODE_USER_DEFINED_MODELS_DIR
self._ain_system_dir = AINODE_SYSTEM_DIR
# Whether to enable compression for thrift
@@ -210,6 +214,18 @@ class AINodeConfig(object):
def set_ain_builtin_models_dir(self, ain_builtin_models_dir: str) -> None:
self._ain_builtin_models_dir = ain_builtin_models_dir
+ def get_ain_finetune_models_dir(self) -> str:
+ return self._ain_finetune_models_dir
+
+ def set_ain_finetune_models_dir(self, ain_finetune_models_dir: str) -> None:
+ self._ain_finetune_models_dir = ain_finetune_models_dir
+
+ def get_ain_user_defined_models_dir(self) -> str:
+ return self._ain_user_defined_models_dir
+
+ def set_ain_user_defined_models_dir(self, ain_user_defined_models_dir: str) -> None:
+ self._ain_user_defined_models_dir = ain_user_defined_models_dir
+
def get_ain_system_dir(self) -> str:
return self._ain_system_dir
diff --git a/iotdb-core/ainode/iotdb/ainode/core/constant.py b/iotdb-core/ainode/iotdb/ainode/core/constant.py
index b9923d3e3ee..74decf9e88a 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/constant.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/constant.py
@@ -18,9 +18,7 @@
import logging
import os
from enum import Enum
-from typing import List
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
from iotdb.thrift.common.ttypes import TEndPoint
IOTDB_AINODE_HOME = os.getenv("IOTDB_AINODE_HOME", "")
@@ -52,21 +50,26 @@ DEFAULT_RECONNECT_TIMES = 3
AINODE_INFERENCE_BATCH_INTERVAL_IN_MS = 15
AINODE_INFERENCE_MAX_PREDICT_LENGTH = 2880
AINODE_INFERENCE_MODEL_MEM_USAGE_MAP = {
- BuiltInModelType.SUNDIAL.value: 1036 * 1024**2, # 1036 MiB
- BuiltInModelType.TIMER_XL.value: 856 * 1024**2, # 856 MiB
+ "sundial": 1036 * 1024**2, # 1036 MiB
+ "timerxl": 856 * 1024**2, # 856 MiB
} # the memory usage of each model in bytes
AINODE_INFERENCE_MEMORY_USAGE_RATIO = 0.4 # the device space allocated for inference
AINODE_INFERENCE_EXTRA_MEMORY_RATIO = (
1.2 # the overhead ratio for inference, used to estimate the pool size
)
-# AINode folder structure
AINODE_MODELS_DIR = os.path.join(IOTDB_AINODE_HOME, "data/ainode/models")
AINODE_BUILTIN_MODELS_DIR = os.path.join(
- IOTDB_AINODE_HOME, "data/ainode/models/weights"
+ IOTDB_AINODE_HOME, "data/ainode/models/builtin"
) # For built-in models, we only need to store their weights and config.
-AINODE_SYSTEM_DIR = os.path.join(IOTDB_AINODE_HOME, "data/ainode/system")
-AINODE_LOG_DIR = os.path.join(IOTDB_AINODE_HOME, "logs")
+AINODE_FINETUNE_MODELS_DIR = os.path.join(
+ IOTDB_AINODE_HOME, "data/ainode/models/finetune"
+)
+AINODE_USER_DEFINED_MODELS_DIR = os.path.join(
+ IOTDB_AINODE_HOME, "data/ainode/models/user_defined"
+)
+AINODE_SYSTEM_DIR = "data/ainode/system"
+AINODE_LOG_DIR = "logs"
# AINode log
LOG_FILE_TYPE = ["all", "info", "warn", "error"]
@@ -141,132 +144,8 @@ class HyperparameterName(Enum):
return self.value
-class ForecastModelType(Enum):
- DLINEAR = "dlinear"
- DLINEAR_INDIVIDUAL = "dlinear_individual"
- NBEATS = "nbeats"
-
- @classmethod
- def values(cls) -> List[str]:
- values = []
- for item in list(cls):
- values.append(item.value)
- return values
-
-
class ModelInputName(Enum):
DATA_X = "data_x"
TIME_STAMP_X = "time_stamp_x"
TIME_STAMP_Y = "time_stamp_y"
DEC_INP = "dec_inp"
-
-
-class AttributeName(Enum):
- # forecast Attribute
- PREDICT_LENGTH = "predict_length"
-
- # NaiveForecaster
- STRATEGY = "strategy"
- SP = "sp"
-
- # STLForecaster
- # SP = 'sp'
- SEASONAL = "seasonal"
- SEASONAL_DEG = "seasonal_deg"
- TREND_DEG = "trend_deg"
- LOW_PASS_DEG = "low_pass_deg"
- SEASONAL_JUMP = "seasonal_jump"
- TREND_JUMP = "trend_jump"
- LOSS_PASS_JUMP = "low_pass_jump"
-
- # ExponentialSmoothing
- DAMPED_TREND = "damped_trend"
- INITIALIZATION_METHOD = "initialization_method"
- OPTIMIZED = "optimized"
- REMOVE_BIAS = "remove_bias"
- USE_BRUTE = "use_brute"
-
- # Arima
- ORDER = "order"
- SEASONAL_ORDER = "seasonal_order"
- METHOD = "method"
- MAXITER = "maxiter"
- SUPPRESS_WARNINGS = "suppress_warnings"
- OUT_OF_SAMPLE_SIZE = "out_of_sample_size"
- SCORING = "scoring"
- WITH_INTERCEPT = "with_intercept"
- TIME_VARYING_REGRESSION = "time_varying_regression"
- ENFORCE_STATIONARITY = "enforce_stationarity"
- ENFORCE_INVERTIBILITY = "enforce_invertibility"
- SIMPLE_DIFFERENCING = "simple_differencing"
- MEASUREMENT_ERROR = "measurement_error"
- MLE_REGRESSION = "mle_regression"
- HAMILTON_REPRESENTATION = "hamilton_representation"
- CONCENTRATE_SCALE = "concentrate_scale"
-
- # GAUSSIAN_HMM
- N_COMPONENTS = "n_components"
- COVARIANCE_TYPE = "covariance_type"
- MIN_COVAR = "min_covar"
- STARTPROB_PRIOR = "startprob_prior"
- TRANSMAT_PRIOR = "transmat_prior"
- MEANS_PRIOR = "means_prior"
- MEANS_WEIGHT = "means_weight"
- COVARS_PRIOR = "covars_prior"
- COVARS_WEIGHT = "covars_weight"
- ALGORITHM = "algorithm"
- N_ITER = "n_iter"
- TOL = "tol"
- PARAMS = "params"
- INIT_PARAMS = "init_params"
- IMPLEMENTATION = "implementation"
-
- # GMMHMM
- # N_COMPONENTS = "n_components"
- N_MIX = "n_mix"
- # MIN_COVAR = "min_covar"
- # STARTPROB_PRIOR = "startprob_prior"
- # TRANSMAT_PRIOR = "transmat_prior"
- WEIGHTS_PRIOR = "weights_prior"
-
- # MEANS_PRIOR = "means_prior"
- # MEANS_WEIGHT = "means_weight"
- # ALGORITHM = "algorithm"
- # COVARIANCE_TYPE = "covariance_type"
- # N_ITER = "n_iter"
- # TOL = "tol"
- # INIT_PARAMS = "init_params"
- # PARAMS = "params"
- # IMPLEMENTATION = "implementation"
-
- # STRAY
- ALPHA = "alpha"
- K = "k"
- KNN_ALGORITHM = "knn_algorithm"
- P = "p"
- SIZE_THRESHOLD = "size_threshold"
- OUTLIER_TAIL = "outlier_tail"
-
- # timerxl
- INPUT_TOKEN_LEN = "input_token_len"
- HIDDEN_SIZE = "hidden_size"
- INTERMEDIATE_SIZE = "intermediate_size"
- OUTPUT_TOKEN_LENS = "output_token_lens"
- NUM_HIDDEN_LAYERS = "num_hidden_layers"
- NUM_ATTENTION_HEADS = "num_attention_heads"
- HIDDEN_ACT = "hidden_act"
- USE_CACHE = "use_cache"
- ROPE_THETA = "rope_theta"
- ATTENTION_DROPOUT = "attention_dropout"
- INITIALIZER_RANGE = "initializer_range"
- MAX_POSITION_EMBEDDINGS = "max_position_embeddings"
- CKPT_PATH = "ckpt_path"
-
- # sundial
- DROPOUT_RATE = "dropout_rate"
- FLOW_LOSS_DEPTH = "flow_loss_depth"
- NUM_SAMPLING_STEPS = "num_sampling_steps"
- DIFFUSION_BATCH_MUL = "diffusion_batch_mul"
-
- def name(self) -> str:
- return self.value
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py
index 82c72cc37ab..a70445c5efd 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request.py
@@ -15,14 +15,12 @@
# specific language governing permissions and limitations
# under the License.
#
+
import threading
from typing import Any
import torch
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
- AbstractInferencePipeline,
-)
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.util.atmoic_int import AtomicInt
@@ -41,7 +39,6 @@ class InferenceRequest:
req_id: str,
model_id: str,
inputs: torch.Tensor,
- inference_pipeline: AbstractInferencePipeline,
max_new_tokens: int = 96,
**infer_kwargs,
):
@@ -52,7 +49,6 @@ class InferenceRequest:
self.model_id = model_id
self.inputs = inputs
self.infer_kwargs = infer_kwargs
- self.inference_pipeline = inference_pipeline
self.max_new_tokens = (
max_new_tokens # Number of time series data points to generate
)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
index 6b054c91fe3..fcaa4c7a754 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/inference_request_pool.py
@@ -25,19 +25,17 @@ from enum import Enum
import numpy as np
import torch
import torch.multiprocessing as mp
-from transformers import PretrainedConfig
from iotdb.ainode.core.config import AINodeDescriptor
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 import get_pipeline
from iotdb.ainode.core.inference.request_scheduler.basic_request_scheduler import (
BasicRequestScheduler,
)
from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.manager.model_manager import ModelManager
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
-from iotdb.ainode.core.model.model_info import ModelInfo
+from iotdb.ainode.core.model.model_storage import ModelInfo
from iotdb.ainode.core.util.gpu_mapping import convert_device_id_to_torch_device
@@ -62,7 +60,6 @@ class InferenceRequestPool(mp.Process):
pool_id: int,
model_info: ModelInfo,
device: str,
- config: PretrainedConfig,
request_queue: mp.Queue,
result_queue: mp.Queue,
ready_event,
@@ -71,7 +68,6 @@ class InferenceRequestPool(mp.Process):
super().__init__()
self.pool_id = pool_id
self.model_info = model_info
- self.config = config
self.pool_kwargs = pool_kwargs
self.ready_event = ready_event
self.device = convert_device_id_to_torch_device(device)
@@ -86,8 +82,8 @@ class InferenceRequestPool(mp.Process):
self._batcher = BasicBatcher()
self._stop_event = mp.Event()
- self._model = None
- self._model_manager = None
+ # self._inference_pipeline = get_pipeline(self.model_info.model_id, self.device)
+
self._logger = None
# Fix inference seed
@@ -98,9 +94,6 @@ class InferenceRequestPool(mp.Process):
def _activate_requests(self):
requests = self._request_scheduler.schedule_activate()
for request in requests:
- request.inputs = request.inference_pipeline.preprocess_inputs(
- request.inputs
- )
request.mark_running()
self._running_queue.put(request)
self._logger.debug(
@@ -123,66 +116,36 @@ class InferenceRequestPool(mp.Process):
for requests in grouped_requests:
batch_inputs = self._batcher.batch_request(requests).to(self.device)
- if self.model_info.model_type == BuiltInModelType.SUNDIAL.value:
- batch_output = self._model.generate(
- batch_inputs,
- max_new_tokens=requests[0].max_new_tokens,
- num_samples=10,
- revin=True,
- )
-
- offset = 0
- for request in requests:
- request.output_tensor = request.output_tensor.to(self.device)
- cur_batch_size = request.batch_size
- cur_output = batch_output[offset : offset + cur_batch_size]
- offset += cur_batch_size
- request.write_step_output(cur_output.mean(dim=1))
-
- request.inference_pipeline.post_decode()
- if request.is_finished():
- request.inference_pipeline.post_inference()
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
- )
- # ensure the output tensor is on CPU before sending to result queue
- request.output_tensor = request.output_tensor.cpu()
- self._finished_queue.put(request)
- else:
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
- )
- self._waiting_queue.put(request)
-
- elif self.model_info.model_type == BuiltInModelType.TIMER_XL.value:
- batch_output = self._model.generate(
- batch_inputs,
- max_new_tokens=requests[0].max_new_tokens,
- revin=True,
- )
-
- offset = 0
- for request in requests:
- request.output_tensor = request.output_tensor.to(self.device)
- cur_batch_size = request.batch_size
- cur_output = batch_output[offset : offset + cur_batch_size]
- offset += cur_batch_size
- request.write_step_output(cur_output)
-
- request.inference_pipeline.post_decode()
- if request.is_finished():
- request.inference_pipeline.post_inference()
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
- )
- # ensure the output tensor is on CPU before sending to result queue
- request.output_tensor = request.output_tensor.cpu()
- self._finished_queue.put(request)
- else:
- self._logger.debug(
- f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
- )
- self._waiting_queue.put(request)
+ batch_output = self._inference_pipeline.infer(
+ batch_inputs,
+ predict_length=requests[0].max_new_tokens,
+ # num_samples=10,
+ revin=True,
+ )
+ offset = 0
+ for request in requests:
+ request.output_tensor = request.output_tensor.to(self.device)
+ cur_batch_size = request.batch_size
+ cur_output = batch_output[offset : offset + cur_batch_size]
+ offset += cur_batch_size
+ # request.write_step_output(cur_output.mean(dim=1))
+ request.write_step_output(cur_output)
+
+ # self._inference_pipeline.post_decode()
+ if request.is_finished():
+ # self._inference_pipeline.post_inference()
+ # ensure the output tensor is on CPU before sending to result queue
+ request.output_tensor = request.output_tensor.cpu()
+ self._finished_queue.put(request)
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
+ )
+ else:
+ self._waiting_queue.put(request)
+ self._logger.debug(
+ f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is not finished, re-queueing"
+ )
+ return
def _requests_execute_loop(self):
while not self._stop_event.is_set():
@@ -193,11 +156,8 @@ class InferenceRequestPool(mp.Process):
self._logger = Logger(
INFERENCE_LOG_FILE_NAME_PREFIX_TEMPLATE.format(self.device)
)
- self._model_manager = ModelManager()
self._request_scheduler.device = self.device
- self._model = self._model_manager.load_model(self.model_info.model_id, {}).to(
- self.device
- )
+ self._inference_pipeline = get_pipeline(self.model_info.model_id, self.device)
self.ready_event.set()
activate_daemon = threading.Thread(
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/__init__.py
similarity index 63%
copy from iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
copy to iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/__init__.py
index 2a1e720805f..617c4e67380 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/__init__.py
@@ -15,3 +15,15 @@
# specific language governing permissions and limitations
# under the License.
#
+
+from iotdb.ainode.core.inference.pipeline.sundial_pipeline import SundialPipeline
+from iotdb.ainode.core.inference.pipeline.timerxl_pipeline import TimerxlPipeline
+
+
+def get_pipeline(model_id, device):
+ if model_id == "timerxl":
+ return TimerxlPipeline(model_id, device=device)
+ elif model_id == "sundial":
+ return SundialPipeline(model_id, device=device)
+ else:
+ raise ValueError(f"Unsupported model_id: {model_id} with pipeline")
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/basic_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/basic_pipeline.py
new file mode 100644
index 00000000000..c413a92e82d
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/basic_pipeline.py
@@ -0,0 +1,111 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from abc import ABC
+
+import torch
+
+from iotdb.ainode.core.exception import InferenceModelInternalError
+from iotdb.ainode.core.manager.model_manager import get_model_manager
+
+
+class BasicPipeline(ABC):
+ def __init__(self, model_id, **infer_kwargs):
+ self.model_id = model_id
+ self.device = infer_kwargs.get("device", "cpu")
+ # self.model = get_model_manager().load_model(model_id).to(self.device)
+ self.model = get_model_manager().load_model(
+ model_id, device_map=str(self.device)
+ )
+
+ def _preprocess(self, inputs):
+ """
+ Preprocess the input before inference, including shape validation and value transformation.
+ """
+ # TODO: Integrate with the data processing pipeline operators
+ pass
+
+ def infer(self, inputs):
+ pass
+
+ def _post_decode(self):
+ """
+ Post-process the outputs after each decode step.
+ """
+ pass
+
+ def _postprocess(self, output: torch.Tensor):
+ """
+ Post-process the outputs after the entire inference task.
+ """
+ pass
+
+
+class ForecastPipeline(BasicPipeline):
+ def __init__(self, model_id, **infer_kwargs):
+ super().__init__(model_id, infer_kwargs=infer_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
+
+ def forecast(self, inputs, **infer_kwargs):
+ pass
+
+ def _post_decode(self):
+ pass
+
+ def _postprocess(self, output: torch.Tensor):
+ pass
+
+
+class ClassificationPipeline(BasicPipeline):
+ def __init__(self, model_id, **infer_kwargs):
+ super().__init__(model_id, infer_kwargs=infer_kwargs)
+
+ def _preprocess(self, inputs):
+ pass
+
+ def classify(self, inputs, **kwargs):
+ pass
+
+ def _post_decode(self):
+ pass
+
+ def _postprocess(self, output: torch.Tensor):
+ pass
+
+
+class ChatPipeline(BasicPipeline):
+ def __init__(self, model_id, **infer_kwargs):
+ super().__init__(model_id, infer_kwargs=infer_kwargs)
+
+ def _preprocess(self, inputs):
+ pass
+
+ def chat(self, inputs, **kwargs):
+ pass
+
+ def _post_decode(self):
+ pass
+
+ def _postprocess(self, output: torch.Tensor):
+ pass
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/sktime_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/sktime_pipeline.py
new file mode 100644
index 00000000000..004222db7ca
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/sktime_pipeline.py
@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import numpy as np
+import pandas as pd
+import torch
+
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import BasicPipeline
+
+
+class SktimePipeline(BasicPipeline):
+ def __init__(self, model_id, **infer_kwargs):
+ super().__init__(model_id, infer_kwargs=infer_kwargs)
+
+ def _preprocess(self, inputs):
+ return super()._preprocess(inputs)
+
+ def infer(self, inputs, **infer_kwargs):
+ 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:
+ # Batch processing: convert each row to Series
+ outputs = []
+ for i in range(input_ids.shape[0]):
+ series = pd.Series(input_ids[i].cpu().numpy() if isinstance(input_ids, torch.Tensor) else input_ids[i])
+ output = self.model.generate(series)
+ 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())
+ else:
+ series = pd.Series(input_ids.squeeze())
+ output = self.model.generate(series)
+ # Add batch dimension if needed
+ if len(output.shape) == 1:
+ output = output[np.newaxis, :]
+
+ return self._postprocess(output)
+
+ def _postprocess(self, output):
+ if isinstance(output, np.ndarray):
+ return torch.from_numpy(output).float()
+ return output
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/sundial_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/sundial_pipeline.py
new file mode 100644
index 00000000000..8d0909954bf
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/sundial_pipeline.py
@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import pandas as pd
+import torch
+
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
+from iotdb.ainode.core.util.serde import convert_to_binary
+
+
+class SundialPipeline(ForecastPipeline):
+ def __init__(self, model_id, **infer_kwargs):
+ super().__init__(model_id, infer_kwargs=infer_kwargs)
+
+ def _preprocess(self, inputs):
+ return super()._preprocess(inputs)
+
+ def infer(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)
+
+ input_ids = self._preprocess(inputs)
+ output = self.model.generate(
+ input_ids,
+ max_new_tokens=predict_length,
+ num_samples=num_samples,
+ revin=revin,
+ )
+ return self._postprocess(output)
+
+ def _postprocess(self, output: torch.Tensor):
+ return output.mean(dim=1)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/timerxl_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/timerxl_pipeline.py
new file mode 100644
index 00000000000..cf6d35c805c
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/timerxl_pipeline.py
@@ -0,0 +1,44 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import pandas as pd
+import torch
+
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
+from iotdb.ainode.core.util.serde import convert_to_binary
+
+
+class TimerxlPipeline(ForecastPipeline):
+ def __init__(self, model_id, **infer_kwargs):
+ super().__init__(model_id, infer_kwargs=infer_kwargs)
+
+ def _preprocess(self, inputs):
+ return super()._preprocess(inputs)
+
+ def infer(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
+ )
+ return self._postprocess(output)
+
+ def _postprocess(self, output: torch.Tensor):
+ return output
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
index 54580402ec2..5af5bb95102 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py
@@ -41,9 +41,6 @@ from iotdb.ainode.core.inference.pool_scheduler.basic_pool_scheduler import (
)
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.manager.model_manager import ModelManager
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
-from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
-from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
from iotdb.ainode.core.util.atmoic_int import AtomicInt
from iotdb.ainode.core.util.batch_executor import BatchExecutor
from iotdb.ainode.core.util.decorator import synchronized
@@ -78,7 +75,7 @@ class PoolController:
# =============== Pool Management ===============
@synchronized(threading.Lock())
- def first_req_init(self, model_id: str):
+ def first_req_init(self, model_id: str, device):
"""
Initialize the pools when the first request for the given model_id arrives.
"""
@@ -110,17 +107,12 @@ class PoolController:
device = torch.device(device_str)
device_id = device.index
- if model_id == "sundial":
- config = SundialConfig()
- elif model_id == "timer_xl":
- config = TimerConfig()
first_queue = mp.Queue()
ready_event = mp.Event()
first_pool = InferenceRequestPool(
pool_id=0,
model_id=model_id,
device=device_str,
- config=config,
request_queue=first_queue,
result_queue=self._result_queue,
ready_event=ready_event,
@@ -255,29 +247,19 @@ class PoolController:
"""
def _expand_pool_on_device(*_):
- result_queue = mp.Queue()
+ request_queue = mp.Queue()
pool_id = self._new_pool_id.get_and_increment()
model_info = self._model_manager.get_model_info(model_id)
- model_type = model_info.model_type
- if model_type == BuiltInModelType.SUNDIAL.value:
- config = SundialConfig()
- elif model_type == BuiltInModelType.TIMER_XL.value:
- config = TimerConfig()
- else:
- raise InferenceModelInternalError(
- f"Unsupported model type {model_type} for loading model {model_id}"
- )
pool = InferenceRequestPool(
pool_id=pool_id,
model_info=model_info,
device=device_id,
- config=config,
- request_queue=result_queue,
+ request_queue=request_queue,
result_queue=self._result_queue,
ready_event=mp.Event(),
)
pool.start()
- self._register_pool(model_id, device_id, pool_id, pool, result_queue)
+ self._register_pool(model_id, device_id, pool_id, pool, request_queue)
if not pool.ready_event.wait(timeout=300):
logger.error(
f"[Inference][Device-{device_id}][Pool-{pool_id}] Pool failed to be ready in time"
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
index 6a2bd2b619a..6ad55f742a4 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py
@@ -28,7 +28,7 @@ from iotdb.ainode.core.inference.pool_scheduler.abstract_pool_scheduler import (
ScaleActionType,
)
from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.manager.model_manager import ModelManager
+from iotdb.ainode.core.manager.model_manager import get_model_manager, ModelManager
from iotdb.ainode.core.manager.utils import (
INFERENCE_EXTRA_MEMORY_RATIO,
INFERENCE_MEMORY_USAGE_RATIO,
@@ -36,12 +36,11 @@ from iotdb.ainode.core.manager.utils import (
estimate_pool_size,
evaluate_system_resources,
)
-from iotdb.ainode.core.model.model_info import BUILT_IN_LTSM_MAP, ModelInfo
+from iotdb.ainode.core.model.model_info import ModelInfo
from iotdb.ainode.core.util.gpu_mapping import convert_device_id_to_torch_device
logger = Logger()
-
def _estimate_shared_pool_size_by_total_mem(
device: torch.device,
existing_model_infos: List[ModelInfo],
@@ -63,7 +62,7 @@ def _estimate_shared_pool_size_by_total_mem(
mem_usages: Dict[str, float] = {}
for model_info in all_models:
mem_usages[model_info.model_id] = (
- MODEL_MEM_USAGE_MAP[model_info.model_type] * INFERENCE_EXTRA_MEMORY_RATIO
+ MODEL_MEM_USAGE_MAP[model_info.model_id] * INFERENCE_EXTRA_MEMORY_RATIO
)
# Evaluate system resources and get TOTAL memory
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py
deleted file mode 100644
index 2300169a6ee..00000000000
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/abstract_inference_pipeline.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-from abc import ABC, abstractmethod
-
-import torch
-
-
-class AbstractInferencePipeline(ABC):
- """
- Abstract assistance strategy class for model inference.
- This class shall define the interface process for specific model.
- """
-
- def __init__(self, model_config, **infer_kwargs):
- self.model_config = model_config
- self.infer_kwargs = infer_kwargs
-
- @abstractmethod
- def preprocess_inputs(self, inputs: torch.Tensor):
- """
- Preprocess the inputs before inference, including shape validation and value transformation.
-
- Args:
- inputs (torch.Tensor): The input tensor to be preprocessed.
-
- Returns:
- torch.Tensor: The preprocessed input tensor.
- """
- # TODO: Integrate with the data processing pipeline operators
- pass
-
- @abstractmethod
- def post_decode(self):
- """
- Post-process the outputs after each decode step.
- """
- pass
-
- @abstractmethod
- def post_inference(self):
- """
- Post-process the outputs after the entire inference task.
- """
- pass
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py
deleted file mode 100644
index 17c88e32fb5..00000000000
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timer_sundial_inference_pipeline.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import torch
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
- AbstractInferencePipeline,
-)
-from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
-
-
-class TimerSundialInferencePipeline(AbstractInferencePipeline):
- """
- Strategy for Timer-Sundial model inference.
- """
-
- def __init__(self, model_config: SundialConfig, **infer_kwargs):
- super().__init__(model_config, infer_kwargs=infer_kwargs)
-
- def preprocess_inputs(self, inputs: torch.Tensor):
- super().preprocess_inputs(inputs)
- if len(inputs.shape) != 2:
- raise InferenceModelInternalError(
- f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
- )
- # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
- return inputs
-
- def post_decode(self):
- # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
- pass
-
- def post_inference(self):
- # TODO: Disassemble and adapt with Sundial's ts_generation_mixin.py
- pass
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py b/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py
deleted file mode 100644
index dc1dd304f68..00000000000
--- a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/timerxl_inference_pipeline.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-#
-
-import torch
-
-from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
- AbstractInferencePipeline,
-)
-from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
-
-
-class TimerXLInferencePipeline(AbstractInferencePipeline):
- """
- Strategy for Timer-XL model inference.
- """
-
- def __init__(self, model_config: TimerConfig, **infer_kwargs):
- super().__init__(model_config, infer_kwargs=infer_kwargs)
-
- def preprocess_inputs(self, inputs: torch.Tensor):
- super().preprocess_inputs(inputs)
- if len(inputs.shape) != 2:
- raise InferenceModelInternalError(
- f"[Inference] Input shape must be: [batch_size, seq_len], but receives {inputs.shape}"
- )
- # Considering that we are currently using the generate function interface, it seems that no pre-processing is required
- return inputs
-
- def post_decode(self):
- # Considering that we are currently using the generate function interface, it seems that no post-processing is required
- pass
-
- def post_inference(self):
- # Considering that we are currently using the generate function interface, it seems that no post-processing is required
- pass
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
index a67d576b0ec..6f022f0ddd1 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/inference_manager.py
@@ -37,21 +37,11 @@ from iotdb.ainode.core.inference.inference_request import (
InferenceRequest,
InferenceRequestProxy,
)
+from iotdb.ainode.core.inference.pipeline import get_pipeline
from iotdb.ainode.core.inference.pool_controller import PoolController
-from iotdb.ainode.core.inference.strategy.timer_sundial_inference_pipeline import (
- TimerSundialInferencePipeline,
-)
-from iotdb.ainode.core.inference.strategy.timerxl_inference_pipeline import (
- TimerXLInferencePipeline,
-)
from iotdb.ainode.core.inference.utils import generate_req_id
from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.manager.model_manager import ModelManager
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
-from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
-from iotdb.ainode.core.model.sundial.modeling_sundial import SundialForPrediction
-from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
-from iotdb.ainode.core.model.timerxl.modeling_timer import TimerForPrediction
+from iotdb.ainode.core.manager.model_manager import get_model_manager
from iotdb.ainode.core.rpc.status import get_status
from iotdb.ainode.core.util.gpu_mapping import get_available_devices
from iotdb.ainode.core.util.serde import convert_to_binary
@@ -71,90 +61,13 @@ from iotdb.tsfile.utils.tsblock_serde import deserialize
logger = Logger()
-class InferenceStrategy(ABC):
- def __init__(self, model):
- self.model = model
-
- @abstractmethod
- def infer(self, full_data, **kwargs):
- pass
-
-
-# [IoTDB] full data deserialized from iotdb is composed of [timestampList, valueList, length],
-# we only get valueList currently.
-class TimerXLStrategy(InferenceStrategy):
- def infer(self, full_data, predict_length=96, **_):
- data = full_data[1][0]
- if data.dtype.byteorder not in ("=", "|"):
- np_data = data.byteswap()
- data = np_data.view(np_data.dtype.newbyteorder())
- seqs = torch.tensor(data).unsqueeze(0).float()
- # TODO: unify model inference input
- output = self.model.generate(seqs, max_new_tokens=predict_length, revin=True)
- df = pd.DataFrame(output[0])
- return convert_to_binary(df)
-
-
-class SundialStrategy(InferenceStrategy):
- def infer(self, full_data, predict_length=96, **_):
- data = full_data[1][0]
- if data.dtype.byteorder not in ("=", "|"):
- np_data = data.byteswap()
- data = np_data.view(np_data.dtype.newbyteorder())
- seqs = torch.tensor(data).unsqueeze(0).float()
- # TODO: unify model inference input
- output = self.model.generate(
- seqs, max_new_tokens=predict_length, num_samples=10, revin=True
- )
- df = pd.DataFrame(output[0].mean(dim=0))
- return convert_to_binary(df)
-
-
-class BuiltInStrategy(InferenceStrategy):
- def infer(self, full_data, **_):
- data = pd.DataFrame(full_data[1]).T
- output = self.model.inference(data)
- df = pd.DataFrame(output)
- return convert_to_binary(df)
-
-
-class RegisteredStrategy(InferenceStrategy):
- def infer(self, full_data, window_interval=None, window_step=None, **_):
- _, dataset, _, length = full_data
- if window_interval is None or window_step is None:
- window_interval = length
- window_step = float("inf")
-
- if window_interval <= 0 or window_step <= 0 or window_interval > length:
- raise InvalidWindowArgumentError(window_interval, window_step, length)
-
- data = torch.tensor(dataset, dtype=torch.float32).unsqueeze(0).permute(0, 2, 1)
-
- times = int((length - window_interval) // window_step + 1)
- results = []
- try:
- for i in range(times):
- start = 0 if window_step == float("inf") else i * window_step
- end = start + window_interval
- window = data[:, start:end, :]
- out = self.model(window)
- df = pd.DataFrame(out.squeeze(0).detach().numpy())
- results.append(df)
- except Exception as e:
- msg = runtime_error_extractor(str(e)) or str(e)
- raise InferenceModelInternalError(msg)
-
- # concatenate or return first window for forecast
- return [convert_to_binary(df) for df in results]
-
-
class InferenceManager:
WAITING_INTERVAL_IN_MS = (
AINodeDescriptor().get_config().get_ain_inference_batch_interval_in_ms()
) # How often to check for requests in the result queue
def __init__(self):
- self._model_manager = ModelManager()
+ self._model_manager = get_model_manager()
self._model_mem_usage_map: Dict[str, int] = (
{}
) # store model memory usage for each model
@@ -251,15 +164,6 @@ class InferenceManager:
with self._result_wrapper_lock:
del self._result_wrapper_map[req_id]
- def _get_strategy(self, model_id, model):
- if isinstance(model, TimerForPrediction):
- return TimerXLStrategy(model)
- if isinstance(model, SundialForPrediction):
- return SundialStrategy(model)
- if self._model_manager.model_storage.is_built_in_or_fine_tuned(model_id):
- return BuiltInStrategy(model)
- return RegisteredStrategy(model)
-
def _run(
self,
req,
@@ -272,9 +176,17 @@ class InferenceManager:
model_id = req.modelId
try:
raw = data_getter(req)
+ # full data deserialized from iotdb is composed of [timestampList, valueList, None, length], we only get valueList currently.
full_data = deserializer(raw)
- inference_attrs = extract_attrs(req)
+ # TODO: TSBlock -> Tensor codes should be unified
+ data = full_data[1][0] # get valueList in ndarray
+ if data.dtype.byteorder not in ("=", "|"):
+ np_data = data.byteswap()
+ data = np_data.view(np_data.dtype.newbyteorder())
+ # the inputs should be on CPU before passing to the inference request
+ inputs = torch.tensor(data).unsqueeze(0).float().to("cpu")
+ inference_attrs = extract_attrs(req)
predict_length = int(inference_attrs.pop("predict_length", 96))
if (
predict_length
@@ -290,41 +202,20 @@ class InferenceManager:
)
if self._pool_controller.has_request_pools(model_id):
- # use request pool to accelerate inference when the model instance is already loaded.
- # TODO: TSBlock -> Tensor codes should be unified
- data = full_data[1][0]
- if data.dtype.byteorder not in ("=", "|"):
- np_data = data.byteswap()
- data = np_data.view(np_data.dtype.newbyteorder())
- # the inputs should be on CPU before passing to the inference request
- inputs = torch.tensor(data).unsqueeze(0).float().to("cpu")
- model_type = self._model_manager.get_model_info(model_id).model_type
- if model_type == BuiltInModelType.SUNDIAL.value:
- inference_pipeline = TimerSundialInferencePipeline(SundialConfig())
- elif model_type == BuiltInModelType.TIMER_XL.value:
- inference_pipeline = TimerXLInferencePipeline(TimerConfig())
- else:
- raise InferenceModelInternalError(
- f"Unsupported model_id: {model_id}"
- )
infer_req = InferenceRequest(
req_id=generate_req_id(),
model_id=model_id,
inputs=inputs,
- inference_pipeline=inference_pipeline,
max_new_tokens=predict_length,
)
outputs = self._process_request(infer_req)
outputs = convert_to_binary(pd.DataFrame(outputs[0]))
else:
- # load model
- accel = str(inference_attrs.get("acceleration", "")).lower() == "true"
- model = self._model_manager.load_model(model_id, inference_attrs, accel)
- # inference by strategy
- strategy = self._get_strategy(model_id, model)
- outputs = strategy.infer(
- full_data, predict_length=predict_length, **inference_attrs
+ inference_pipeline = get_pipeline(model_id, device="cpu")
+ outputs = inference_pipeline.infer(
+ inputs, predict_length=predict_length, **inference_attrs
)
+ outputs = convert_to_binary(pd.DataFrame(outputs[0]))
# construct response
status = get_status(TSStatusCode.SUCCESS_STATUS)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
new file mode 100644
index 00000000000..a07552922ff
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import os
+from typing import Any, List, Optional
+
+from iotdb.ainode.core.config import AINodeDescriptor
+from iotdb.ainode.core.constant import TSStatusCode
+from iotdb.ainode.core.exception import BuiltInModelDeletionError
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.model.model_loader import ModelLoader
+from iotdb.ainode.core.model.model_storage import ModelCategory, ModelInfo, ModelStorage
+from iotdb.ainode.core.rpc.status import get_status
+from iotdb.thrift.ainode.ttypes import (
+ TDeleteModelReq,
+ TRegisterModelReq,
+ TRegisterModelResp,
+ TShowModelsReq,
+ TShowModelsResp,
+)
+from iotdb.thrift.common.ttypes import TSStatus
+
+logger = Logger()
+
+
+class ModelManager:
+ def __init__(self):
+ self.models_dir = os.path.join(
+ os.getcwd(), AINodeDescriptor().get_config().get_ain_models_dir()
+ )
+ self.storage = ModelStorage(models_dir=self.models_dir)
+ self.loader = ModelLoader(storage=self.storage)
+
+ # Automatically discover all models
+ self._models = self.storage.discover_all()
+
+ def register_model(
+ self,
+ req: TRegisterModelReq,
+ ) -> TRegisterModelResp:
+ try:
+ success = self.storage.register_model(model_id=req.modelId, uri=req.uri)
+ if success:
+ return TRegisterModelResp(get_status(TSStatusCode.SUCCESS_STATUS))
+ else:
+ return TRegisterModelResp(
+ get_status(TSStatusCode.AINODE_INTERNAL_ERROR)
+ )
+ except ValueError as e:
+ return TRegisterModelResp(
+ get_status(TSStatusCode.INVALID_URI_ERROR, str(e))
+ )
+ except Exception as e:
+ return TRegisterModelResp(get_status(TSStatusCode.AINODE_INTERNAL_ERROR))
+
+ def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
+ return self.storage.show_models(req)
+
+ def delete_model(self, req: TDeleteModelReq) -> TSStatus:
+ try:
+ self.storage.delete_model(req.modelId)
+ return get_status(TSStatusCode.SUCCESS_STATUS)
+ except BuiltInModelDeletionError as e:
+ logger.warning(e)
+ return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+ except Exception as e:
+ logger.warning(e)
+ return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
+
+ def load_model(self, model_id: str, **kwargs) -> Any:
+ return self.loader.load_model(model_id=model_id, **kwargs)
+
+ def get_model_info(
+ self,
+ model_id: str,
+ category: Optional[ModelCategory] = None,
+ ) -> Optional[ModelInfo]:
+ return self.storage.get_model_info(model_id, category)
+
+ def get_model_infos(
+ self,
+ category: Optional[ModelCategory] = None,
+ model_type: Optional[str] = None,
+ ) -> List[ModelInfo]:
+ return self.storage.get_model_infos(category, model_type)
+
+ def refresh(self):
+ """Refresh the model list (re-scan the file system)"""
+ self._models = self.storage.discover_all()
+
+ def get_registered_models(self) -> List[str]:
+ return self.storage.get_registered_models()
+
+ def is_model_registered(self, model_id: str) -> bool:
+ return self.storage.is_model_registered(model_id)
+
+
+# Create a global model manager instance
+_default_manager: Optional[ModelManager] = None
+
+
+def get_model_manager() -> ModelManager:
+ global _default_manager
+ if _default_manager is None:
+ _default_manager = ModelManager()
+ return _default_manager
diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
index 0264e27331a..297a2b832d7 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py
@@ -24,8 +24,7 @@ import torch
from iotdb.ainode.core.config import AINodeDescriptor
from iotdb.ainode.core.exception import ModelNotExistError
from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.manager.model_manager import ModelManager
-from iotdb.ainode.core.model.model_info import BUILT_IN_LTSM_MAP
+from iotdb.ainode.core.manager.model_manager import get_model_manager
logger = Logger()
@@ -47,7 +46,7 @@ def measure_model_memory(device: torch.device, model_id: str) -> int:
torch.cuda.synchronize(device)
start = torch.cuda.memory_reserved(device)
- model = ModelManager().load_model(model_id, {}).to(device)
+ model = get_model_manager().load_model(model_id, {}).to(device)
torch.cuda.synchronize(device)
end = torch.cuda.memory_reserved(device)
usage = end - start
@@ -80,8 +79,8 @@ def evaluate_system_resources(device: torch.device) -> dict:
def estimate_pool_size(device: torch.device, model_id: str) -> int:
- model_info = BUILT_IN_LTSM_MAP.get(model_id, None)
- if model_info is None or model_info.model_type not in MODEL_MEM_USAGE_MAP:
+ model_info = get_model_manager.get_model_info(model_id)
+ if model_info is None or model_info.model_id not in MODEL_MEM_USAGE_MAP:
logger.error(
f"[Inference] Cannot estimate inference pool size on device: {device}, because model: {model_id} is not supported."
)
@@ -90,9 +89,7 @@ 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_id] * INFERENCE_EXTRA_MEMORY_RATIO
size = int((free_mem * INFERENCE_MEMORY_USAGE_RATIO) // mem_usage)
if size <= 0:
logger.error(
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/__init__.py
similarity index 100%
copy from iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
copy to iotdb-core/ainode/iotdb/ainode/core/model/__init__.py
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_enums.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_enums.py
new file mode 100644
index 00000000000..a6a234a1ab8
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_enums.py
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+from enum import Enum
+
+
+class ModelCategory(Enum):
+ BUILTIN = "builtin"
+ USER_DEFINED = "user_defined"
+ FINETUNE = "finetune"
+
+
+class ModelStates(Enum):
+ INACTIVE = "inactive"
+ ACTIVATING = "activating"
+ ACTIVE = "active"
+ LOADING = "loading"
+ LOADED = "loaded"
+ DROPPING = "dropping"
+ TRAINING = "training"
+ FAILED = "failed"
+
+
+class ModelFileType(Enum):
+ SAFETENSORS = "safetensors"
+ PYTORCH = "pytorch"
+ UNKNOWN = "unknown"
+
+
+class UriType(Enum):
+ REPO = "repo"
+ FILE = "file"
+
+
+# Map for inferring which HuggingFace repository to download from based on model ID
+REPO_ID_MAP = {
+ "timerxl": "thuml/timer-base-84m",
+ "sundial": "thuml/sundial-base-128m",
+ # More mappings can be added as needed
+}
+
+# Model file constants
+MODEL_CONFIG_FILE = "config.json"
+MODEL_WEIGHTS_FILE = "model.safetensors"
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
new file mode 100644
index 00000000000..f36ad582a83
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import Dict, List, Optional, Tuple
+
+from iotdb.ainode.core.model.model_enums import ModelCategory, ModelStates
+
+# Map for inferring which HuggingFace repository to download from based on model ID
+REPO_ID_MAP = {
+ "timerxl": "thuml/timer-base-84m",
+ "sundial": "thuml/sundial-base-128m",
+ # More mappings can be added as needed
+}
+
+# Model file constants
+MODEL_CONFIG_FILE = "config.json"
+MODEL_WEIGHTS_FILE = "model.safetensors"
+
+
+class ModelInfo:
+ def __init__(
+ self,
+ model_id: str,
+ model_type: str,
+ category: ModelCategory,
+ state: ModelStates,
+ path: str = "",
+ auto_map: Optional[Dict] = None,
+ _transformers_registered: bool = False,
+ ):
+ self.model_id = model_id
+ self.model_type = model_type
+ self.category = category
+ self.state = state
+ self.path = path
+ self.auto_map = auto_map # If exists, indicates it's a Transformers model
+ self._transformers_registered = _transformers_registered # Internal flag: whether registered to Transformers
+
+ def __repr__(self):
+ return (
+ f"ModelInfo(model_id={self.model_id}, model_type={self.model_type}, "
+ f"category={self.category.value}, state={self.state.value}, "
+ f"path={self.path}, has_auto_map={self.auto_map is not None})"
+ )
+
+
+BUILTIN_SKTIME_MODEL_MAP = {
+ # forecast models
+ "arima": ModelInfo(
+ model_id="arima",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ "holtwinters": ModelInfo(
+ model_id="holtwinters",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ "exponential_smoothing": ModelInfo(
+ model_id="exponential_smoothing",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ "naive_forecaster": ModelInfo(
+ model_id="naive_forecaster",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ "stl_forecaster": ModelInfo(
+ model_id="stl_forecaster",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ # anomaly detection models
+ "gaussian_hmm": ModelInfo(
+ model_id="gaussian_hmm",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ "gmm_hmm": ModelInfo(
+ model_id="gmm_hmm",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+ "stray": ModelInfo(
+ model_id="stray",
+ model_type="sktime",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.ACTIVE,
+ ),
+}
+
+# Built-in huggingface transformers models, their weights are not included in AINode by default
+BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
+ "timerxl": ModelInfo(
+ model_id="timerxl",
+ model_type="timer",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.INACTIVE,
+ ),
+ "sundial": ModelInfo(
+ model_id="sundial",
+ model_type="sundial",
+ category=ModelCategory.BUILTIN,
+ state=ModelStates.INACTIVE,
+ ),
+}
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py
new file mode 100644
index 00000000000..9b70b5f8040
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import os
+from typing import Any
+
+import torch
+from transformers import (
+ AutoConfig,
+ AutoModelForCausalLM,
+ AutoModelForNextSentencePrediction,
+ AutoModelForSeq2SeqLM,
+ AutoModelForSequenceClassification,
+ AutoModelForTimeSeriesPrediction,
+ AutoModelForTokenClassification,
+)
+
+from iotdb.ainode.core.exception import ModelNotExistError
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.model.model_enums import ModelCategory
+from iotdb.ainode.core.model.model_info import ModelInfo
+from iotdb.ainode.core.model.model_storage import ModelStorage
+from iotdb.ainode.core.model.sktime.modeling_sktime import create_sktime_model
+
+logger = Logger()
+
+
+class ModelLoader:
+ """Model loader - unified interface for loading different types of models"""
+
+ def __init__(self, storage: ModelStorage):
+ self.storage = storage
+
+ def load_model(self, model_id: str, **kwargs) -> Any:
+ # Lazy registration: if it's a Transformers model and not registered, register it first
+ model_info = self.storage.ensure_transformers_registered(model_id)
+ if not model_info:
+ logger.error(
+ f"Model {model_id} failed to register to Transformers, cannot load."
+ )
+ return None
+
+ if model_info.auto_map is not None:
+ model = self.load_model_from_transformers(model_info, **kwargs)
+ else:
+ if model_info.model_type == "sktime":
+ model = create_sktime_model(model_id)
+ else:
+ model = self.load_model_from_pt(model_info, **kwargs)
+
+ logger.info(f"Model {model_id} loaded to device {model.device} successfully.")
+ return model
+
+ def load_model_from_transformers(self, model_info: ModelInfo, **kwargs):
+ model_config, load_class = None, None
+ device_map = kwargs.get("device_map", "cpu")
+ trust_remote_code = kwargs.get("trust_remote_code", True)
+ train_from_scratch = kwargs.get("train_from_scratch", False)
+
+ if model_info.category == ModelCategory.BUILTIN:
+ if model_info.model_id == "timerxl":
+ from iotdb.ainode.core.model.timerxl.configuration_timer import (
+ TimerConfig,
+ )
+
+ model_config = TimerConfig()
+ from iotdb.ainode.core.model.timerxl.modeling_timer import (
+ TimerForPrediction,
+ )
+
+ load_class = TimerForPrediction
+ elif model_info.model_id == "sundial":
+ from iotdb.ainode.core.model.sundial.configuration_sundial import (
+ SundialConfig,
+ )
+
+ model_config = SundialConfig()
+ from iotdb.ainode.core.model.sundial.modeling_sundial import (
+ SundialForPrediction,
+ )
+
+ load_class = SundialForPrediction
+ else:
+ logger.error(
+ f"Unsupported built-in Transformers model {model_info.model_id}."
+ )
+ else:
+ model_config = AutoConfig.from_pretrained(model_info.path)
+ if (
+ type(model_config)
+ in AutoModelForTimeSeriesPrediction._model_mapping.keys()
+ ):
+ load_class = AutoModelForTimeSeriesPrediction
+ elif (
+ type(model_config)
+ in AutoModelForNextSentencePrediction._model_mapping.keys()
+ ):
+ load_class = AutoModelForNextSentencePrediction
+ elif type(model_config) in AutoModelForSeq2SeqLM._model_mapping.keys():
+ load_class = AutoModelForSeq2SeqLM
+ elif (
+ type(model_config)
+ in AutoModelForSequenceClassification._model_mapping.keys()
+ ):
+ load_class = AutoModelForSequenceClassification
+ elif (
+ type(model_config)
+ in AutoModelForTokenClassification._model_mapping.keys()
+ ):
+ load_class = AutoModelForTokenClassification
+ else:
+ load_class = AutoModelForCausalLM
+
+ if train_from_scratch:
+ model = load_class.from_config(
+ model_config, trust_remote_code=trust_remote_code, device_map=device_map
+ )
+ else:
+ model = load_class.from_pretrained(
+ model_info.path,
+ trust_remote_code=trust_remote_code,
+ device_map=device_map,
+ )
+
+ return model
+
+ def load_model_from_pt(self, model_info: ModelInfo, **kwargs):
+ device_map = kwargs.get("device_map", "cpu")
+ acceleration = kwargs.get("acceleration", False)
+ model_path = os.path.join(model_info.path, "model.pt")
+ if not os.path.exists(model_path):
+ logger.error(f"Model file not found at {model_path}.")
+ raise ModelNotExistError(model_path)
+ model = torch.jit.load(model_path)
+ if (
+ isinstance(model, torch._dynamo.eval_frame.OptimizedModule)
+ or not acceleration
+ ):
+ return model
+ try:
+ model = torch.compile(model)
+ except Exception as e:
+ logger.warning(f"acceleration failed, fallback to normal mode: {str(e)}")
+ return model.to(device_map)
+
+ def load_model_for_efficient_inference(self):
+ # TODO: An efficient model loading method for inference based on model_arguments
+ pass
+
+ def load_model_for_powerful_finetune(self):
+ # TODO: An powerful model loading method for finetune based on model_arguments
+ pass
+
+ def unload_model(self):
+ pass
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
new file mode 100644
index 00000000000..5ba67e158a5
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
@@ -0,0 +1,576 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import concurrent.futures
+import os
+import shutil
+from typing import List, Optional
+
+from huggingface_hub import hf_hub_download, snapshot_download
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from iotdb.ainode.core.constant import TSStatusCode
+from iotdb.ainode.core.exception import BuiltInModelDeletionError
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.model.model_enums import REPO_ID_MAP, ModelCategory, ModelStates
+from iotdb.ainode.core.model.model_info import (
+ BUILTIN_HF_TRANSFORMERS_MODEL_MAP,
+ BUILTIN_SKTIME_MODEL_MAP,
+ ModelInfo,
+)
+from iotdb.ainode.core.model.utils import *
+from iotdb.ainode.core.util.lock import ModelLockPool
+from iotdb.thrift.ainode.ttypes import TShowModelsReq, TShowModelsResp
+from iotdb.thrift.common.ttypes import TSStatus
+
+logger = Logger()
+
+
+class ModelStorage:
+ """Model storage class - unified management of model discovery and registration"""
+
+ def __init__(self, models_dir: str):
+ self.models_dir = Path(models_dir)
+ # Unified storage: category -> {model_id -> ModelInfo}
+ self._models: Dict[str, Dict[str, ModelInfo]] = {
+ ModelCategory.BUILTIN.value: {},
+ ModelCategory.USER_DEFINED.value: {},
+ ModelCategory.FINETUNE.value: {},
+ }
+ # Async download executor (using single-threaded executor because hf download interface is unstable with concurrent downloads)
+ self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+ # Thread lock pool for protecting concurrent access to model information
+ self._lock_pool = ModelLockPool()
+ self._initialize_directories()
+
+ def _initialize_directories(self):
+ """Initialize directory structure and ensure __init__.py files exist"""
+ self.models_dir.mkdir(parents=True, exist_ok=True)
+ ensure_init_file(self.models_dir)
+
+ for category in ModelCategory:
+ category_path = self.models_dir / category.value
+ category_path.mkdir(parents=True, exist_ok=True)
+ ensure_init_file(category_path)
+
+ # ==================== Discovery Methods ====================
+
+ def discover_all(self) -> Dict[str, Dict[str, ModelInfo]]:
+ """Scan file system to discover all models"""
+ self._discover_category(ModelCategory.BUILTIN)
+ self._discover_category(ModelCategory.USER_DEFINED)
+ self._discover_category(ModelCategory.FINETUNE)
+ return self._models
+
+ def _discover_category(self, category: ModelCategory):
+ """Discover all models in a category directory"""
+ category_path = self.models_dir / category.value
+ if not category_path.exists():
+ return
+
+ if category == ModelCategory.BUILTIN:
+ self._discover_builtin_models(category_path)
+ else:
+ # For finetune and user_defined, scan directories
+ for item in category_path.iterdir():
+ if item.is_dir() and not item.name.startswith("__"):
+ relative_path = item.relative_to(category_path)
+ model_id = str(relative_path).replace("/", "_").replace("\\", "_")
+ self._process_model_directory(item, model_id, category)
+
+ def _discover_builtin_models(self, category_path: Path):
+ # Register SKTIME models directly from map
+ for model_id in BUILTIN_SKTIME_MODEL_MAP.keys():
+ with self._lock_pool.get_lock(model_id).write_lock():
+ self._models[ModelCategory.BUILTIN.value][model_id] = (
+ BUILTIN_SKTIME_MODEL_MAP[model_id]
+ )
+
+ # Process HuggingFace Transformers models
+ for model_id in BUILTIN_HF_TRANSFORMERS_MODEL_MAP.keys():
+ model_dir = category_path / model_id
+ model_dir.mkdir(parents=True, exist_ok=True)
+ self._process_model_directory(model_dir, model_id, ModelCategory.BUILTIN)
+
+ def _process_model_directory(
+ self, model_dir: Path, model_id: str, category: ModelCategory
+ ):
+ """Handling the discovery logic for a single model directory."""
+ ensure_init_file(model_dir)
+
+ config_path = model_dir / MODEL_CONFIG_FILE
+ weights_path = model_dir / MODEL_WEIGHTS_FILE
+ needs_download = not config_path.exists() or not weights_path.exists()
+
+ if needs_download:
+ with self._lock_pool.get_lock(model_id).write_lock():
+ model_info = ModelInfo(
+ model_id=model_id,
+ model_type="", # Read from config.json after download
+ category=category,
+ state=ModelStates.ACTIVATING,
+ path=str(model_dir),
+ auto_map=None,
+ _transformers_registered=False,
+ )
+ self._models[category.value][model_id] = model_info
+
+ future = self._executor.submit(
+ self._download_model_if_necessary, str(model_dir), model_id
+ )
+ future.add_done_callback(
+ lambda f, mid=model_id, cat=category: self._callback_model_download_result(
+ f, mid, cat
+ )
+ )
+ else:
+ config = load_model_config(config_path)
+ model_type = config.get("model_type", "")
+ auto_map = config.get("auto_map")
+
+ with self._lock_pool.get_lock(model_id).write_lock():
+ model_info = ModelInfo(
+ model_id=model_id,
+ model_type=model_type,
+ category=category,
+ state=ModelStates.ACTIVE,
+ path=str(model_dir),
+ auto_map=auto_map,
+ _transformers_registered=False, # Lazy registration
+ )
+ self._models[category.value][model_id] = model_info
+
+ def _download_model_if_necessary(self, model_dir: str, model_id: str) -> bool:
+ """Returns: True if the model is existed or downloaded successfully, False otherwise."""
+ if model_id in REPO_ID_MAP:
+ repo_id = REPO_ID_MAP[model_id]
+ else:
+ logger.error(f"Model {model_id} not found in REPO_ID_MAP")
+ return False
+
+ weights_path = os.path.join(model_dir, MODEL_WEIGHTS_FILE)
+ config_path = os.path.join(model_dir, MODEL_CONFIG_FILE)
+
+ if not os.path.exists(weights_path):
+ try:
+ hf_hub_download(
+ repo_id=repo_id,
+ filename=MODEL_WEIGHTS_FILE,
+ local_dir=model_dir,
+ )
+ except Exception as e:
+ logger.error(f"Failed to download model weights from HuggingFace: {e}")
+ return False
+
+ if not os.path.exists(config_path):
+ try:
+ hf_hub_download(
+ repo_id=repo_id,
+ filename=MODEL_CONFIG_FILE,
+ local_dir=model_dir,
+ )
+ except Exception as e:
+ logger.error(f"Failed to download model config from HuggingFace: {e}")
+ return False
+
+ return True
+
+ def _callback_model_download_result(
+ self, future, model_id: str, category: ModelCategory
+ ):
+ """Callback function for handling model download results"""
+ with self._lock_pool.get_lock(model_id).write_lock():
+ try:
+ if future.result():
+ if model_id in self._models[category.value]:
+ model_info = self._models[category.value][model_id]
+ model_info.state = ModelStates.ACTIVE
+ config_path = os.path.join(model_info.path, MODEL_CONFIG_FILE)
+ if os.path.exists(config_path):
+ with open(config_path, "r", encoding="utf-8") as f:
+ config = json.load(f)
+ model_info.model_type = config.get("model_type", "")
+ model_info.auto_map = config.get("auto_map")
+ logger.info(
+ f"Model {model_id} downloaded successfully and is ready to use."
+ )
+ else:
+ if model_id in self._models[category.value]:
+ self._models[category.value][
+ model_id
+ ].state = ModelStates.INACTIVE
+ logger.warning(f"Failed to download model {model_id}.")
+ except Exception as e:
+ logger.error(f"Error in download callback for model {model_id}: {e}")
+ if model_id in self._models[category.value]:
+ self._models[category.value][model_id].state = ModelStates.INACTIVE
+
+ # ==================== Registration Methods ====================
+
+ def register_model(self, model_id: str, uri: str) -> bool:
+ """
+ Supported URI formats:
+ - repo://<huggingface_repo_id>
+ - file://<local_path>
+ """
+ uri_type = parse_uri_type(uri)
+ parsed_uri = get_parsed_uri(uri)
+
+ model_dir = Path(self.models_dir) / "user_defined" / model_id
+ model_dir.mkdir(parents=True, exist_ok=True)
+ ensure_init_file(model_dir)
+
+ if uri_type == UriType.REPO:
+ self._fetch_model_from_hf_repo(parsed_uri, str(model_dir))
+ else:
+ self._fetch_model_from_local(os.path.expanduser(parsed_uri), str(model_dir))
+
+ config_path, _ = validate_model_files(model_dir)
+ config = load_model_config(config_path)
+ model_type = config.get("model_type", "")
+ auto_map = config.get("auto_map")
+
+ with self._lock_pool.get_lock(model_id).write_lock():
+ model_info = ModelInfo(
+ model_id=model_id,
+ model_type=model_type,
+ category=ModelCategory.USER_DEFINED,
+ state=ModelStates.ACTIVE,
+ path=str(model_dir),
+ auto_map=auto_map,
+ _transformers_registered=False, # Register later
+ )
+ self._models[ModelCategory.USER_DEFINED.value][model_id] = model_info
+
+ if auto_map:
+ # Transformers model: immediately register to Transformers auto-loading mechanism
+ success = self._register_transformers_model(model_info)
+ if success:
+ with self._lock_pool.get_lock(model_id).write_lock():
+ model_info._transformers_registered = True
+ else:
+ with self._lock_pool.get_lock(model_id).write_lock():
+ model_info.state = ModelStates.INACTIVE
+ logger.error(f"Failed to register Transformers model {model_id}")
+ return False
+ else:
+ # Other type models: only log
+ self._register_other_model(model_info)
+
+ logger.info(f"Successfully registered model {model_id} from URI: {uri}")
+ return True
+
+ def _fetch_model_from_hf_repo(self, repo_id: str, storage_path: str):
+ logger.info(
+ f"Downloading model from HuggingFace repository: {repo_id} -> {storage_path}"
+ )
+
+ # Use snapshot_download to download entire repository (including config.json and model.safetensors)
+ try:
+ snapshot_download(
+ repo_id=repo_id,
+ local_dir=storage_path,
+ local_dir_use_symlinks=False,
+ )
+ except Exception as e:
+ logger.error(f"Failed to download model from HuggingFace: {e}")
+ raise
+
+ def _fetch_model_from_local(self, source_path: str, storage_path: str):
+ logger.info(f"Copying model from local path: {source_path} -> {storage_path}")
+
+ source_dir = Path(source_path)
+ if not source_dir.is_dir():
+ raise ValueError(
+ f"Source path does not exist or is not a directory: {source_path}"
+ )
+
+ source_config = source_dir / MODEL_CONFIG_FILE
+ source_weights = source_dir / MODEL_WEIGHTS_FILE
+ if not source_config.exists():
+ raise ValueError(
+ f"Config file missing in source directory: {source_config}"
+ )
+ if not source_weights.exists():
+ raise ValueError(
+ f"Weights file missing in source directory: {source_weights}"
+ )
+
+ # Copy all files
+ storage_dir = Path(storage_path)
+ for file in source_dir.iterdir():
+ if file.is_file():
+ shutil.copy2(file, storage_dir / file.name)
+
+ def _register_transformers_model(self, model_info: ModelInfo) -> bool:
+ """
+ Register Transformers model to auto-loading mechanism (internal method)
+ """
+ auto_map = model_info.auto_map
+ if not auto_map:
+ return False
+
+ auto_config_path = auto_map.get("AutoConfig")
+ auto_model_path = auto_map.get("AutoModelForCausalLM")
+
+ try:
+ module_parent = str(Path(model_info.path).parent.absolute())
+ with temporary_sys_path(module_parent):
+ config_class = import_class_from_path(
+ model_info.model_id, auto_config_path
+ )
+ AutoConfig.register(model_info.model_type, config_class)
+ logger.info(
+ f"Registered AutoConfig: {model_info.model_type} -> {auto_config_path}"
+ )
+
+ model_class = import_class_from_path(
+ model_info.model_id, auto_model_path
+ )
+ AutoModelForCausalLM.register(config_class, model_class)
+ logger.info(
+ f"Registered AutoModelForCausalLM: {config_class.__name__} -> {auto_model_path}"
+ )
+
+ return True
+ except Exception as e:
+ logger.warning(
+ f"Failed to register Transformers model {model_info.model_id}: {e}. Model may still work via auto_map, but ensure module path is correct."
+ )
+ return False
+
+ def _register_other_model(self, model_info: ModelInfo):
+ """Register other type models (non-Transformers models)"""
+ logger.info(
+ f"Registered other type model: {model_info.model_id} ({model_info.model_type})"
+ )
+
+ def ensure_transformers_registered(self, model_id: str) -> "ModelInfo":
+ """
+ Ensure Transformers model is registered (called for lazy registration)
+ This method uses locks to ensure thread safety. All check logic is within lock protection.
+ Returns:
+ str: If None, registration failed, otherwise returns model path
+ """
+ # Use lock to protect entire check-execute process
+ with self._lock_pool.get_lock(model_id).write_lock():
+ # Directly access _models dictionary (avoid calling get_model_info which may cause deadlock)
+ model_info = None
+ for category_dict in self._models.values():
+ if model_id in category_dict:
+ model_info = category_dict[model_id]
+ break
+
+ if not model_info:
+ logger.warning(f"Model {model_id} does not exist, cannot register")
+ return None
+
+ # If already registered, return directly
+ if model_info._transformers_registered:
+ return model_info
+
+ # If no auto_map, not a Transformers model, mark as registered (avoid duplicate checks)
+ if (
+ not model_info.auto_map
+ or model_id in BUILTIN_HF_TRANSFORMERS_MODEL_MAP.keys()
+ ):
+ model_info._transformers_registered = True
+ return model_info
+
+ # Execute registration (under lock protection)
+ try:
+ success = self._register_transformers_model(model_info)
+ if success:
+ model_info._transformers_registered = True
+ logger.info(
+ f"Model {model_id} successfully registered to Transformers"
+ )
+ return model_info
+ else:
+ model_info.state = ModelStates.INACTIVE
+ logger.error(f"Model {model_id} failed to register to Transformers")
+ return None
+
+ except Exception as e:
+ # Ensure state consistency in exception cases
+ model_info.state = ModelStates.INACTIVE
+ model_info._transformers_registered = False
+ logger.error(
+ f"Exception occurred while registering model {model_id} to Transformers: {e}"
+ )
+ return None
+
+ # ==================== Show and Delete Models ====================
+
+ def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
+ resp_status = TSStatus(
+ code=TSStatusCode.SUCCESS_STATUS.value,
+ message="Show models successfully",
+ )
+
+ # Use global lock to protect entire dictionary structure
+ with self._lock_pool.get_lock("").read_lock():
+ if req.modelId:
+ # Find specified model
+ model_info = None
+ for category_dict in self._models.values():
+ if req.modelId in category_dict:
+ model_info = category_dict[req.modelId]
+ break
+
+ if model_info:
+ return TShowModelsResp(
+ status=resp_status,
+ modelIdList=[req.modelId],
+ modelTypeMap={req.modelId: model_info.model_type},
+ categoryMap={req.modelId: model_info.category.value},
+ stateMap={req.modelId: model_info.state.value},
+ )
+ else:
+ return TShowModelsResp(
+ status=resp_status,
+ modelIdList=[],
+ modelTypeMap={},
+ categoryMap={},
+ stateMap={},
+ )
+ else:
+ # Return all models
+ model_id_list = []
+ model_type_map = {}
+ category_map = {}
+ state_map = {}
+
+ for category_dict in self._models.values():
+ for model_id, model_info in category_dict.items():
+ model_id_list.append(model_id)
+ model_type_map[model_id] = model_info.model_type
+ category_map[model_id] = model_info.category.value
+ state_map[model_id] = model_info.state.value
+
+ return TShowModelsResp(
+ status=resp_status,
+ modelIdList=model_id_list,
+ modelTypeMap=model_type_map,
+ categoryMap=category_map,
+ stateMap=state_map,
+ )
+
+ def delete_model(self, model_id: str) -> None:
+ # Use write lock to protect entire deletion process
+ with self._lock_pool.get_lock(model_id).write_lock():
+ model_info = None
+ category_value = None
+ for cat_value, category_dict in self._models.items():
+ if model_id in category_dict:
+ model_info = category_dict[model_id]
+ category_value = cat_value
+ break
+
+ if not model_info:
+ logger.warning(f"Model {model_id} does not exist, cannot delete")
+ return
+
+ if model_info.category == ModelCategory.BUILTIN:
+ raise BuiltInModelDeletionError(model_id)
+
+ model_path = Path(model_info.path)
+ if model_path.exists():
+ try:
+ shutil.rmtree(model_path)
+ logger.info(f"Deleted model directory: {model_path}")
+ except Exception as e:
+ logger.error(f"Failed to delete model directory {model_path}: {e}")
+ raise
+
+ if category_value and model_id in self._models[category_value]:
+ del self._models[category_value][model_id]
+ logger.info(f"Model {model_id} has been removed from storage")
+
+ return
+
+ # ==================== Query Methods ====================
+
+ def get_model_info(
+ self, model_id: str, category: Optional[ModelCategory] = None
+ ) -> Optional[ModelInfo]:
+ """
+ Get single model information
+
+ If category is specified, use model_id's lock
+ If category is not specified, need to traverse all dictionaries, use global lock
+ """
+ if category:
+ # Category specified, only need to access specific dictionary, use model_id's lock
+ with self._lock_pool.get_lock(model_id).read_lock():
+ return self._models[category.value].get(model_id)
+ else:
+ # Category not specified, need to traverse all dictionaries, use global lock
+ with self._lock_pool.get_lock("").read_lock():
+ for category_dict in self._models.values():
+ if model_id in category_dict:
+ return category_dict[model_id]
+ return None
+
+ def get_model_infos(
+ self, category: Optional[ModelCategory] = None, model_type: Optional[str] = None
+ ) -> List[ModelInfo]:
+ """
+ Get model information list
+
+ Note: Since we need to traverse all models, use a global lock to protect the entire dictionary structure
+ For single model access, using model_id-based lock would be more efficient
+ """
+ matching_models = []
+
+ # For traversal operations, we need to protect the entire dictionary structure
+ # Use a special lock (using empty string as key) to protect the entire dictionary
+ with self._lock_pool.get_lock("").read_lock():
+ if category and model_type:
+ for model_info in self._models[category.value].values():
+ if model_info.model_type == model_type:
+ matching_models.append(model_info)
+ return matching_models
+ elif category:
+ return list(self._models[category.value].values())
+ elif model_type:
+ for category_dict in self._models.values():
+ for model_info in category_dict.values():
+ if model_info.model_type == model_type:
+ matching_models.append(model_info)
+ return matching_models
+ else:
+ for category_dict in self._models.values():
+ matching_models.extend(category_dict.values())
+ return matching_models
+
+ def is_model_registered(self, model_id: str) -> bool:
+ """Check if model is registered (search in _models)"""
+ with self._lock_pool.get_lock("").read_lock():
+ for category_dict in self._models.values():
+ if model_id in category_dict:
+ return True
+ return False
+
+ def get_registered_models(self) -> List[str]:
+ """Get list of all registered model IDs"""
+ with self._lock_pool.get_lock("").read_lock():
+ model_ids = []
+ for category_dict in self._models.values():
+ model_ids.extend(category_dict.keys())
+ return model_ids
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/__init__.py
similarity index 100%
copy from iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
copy to iotdb-core/ainode/iotdb/ainode/core/model/sktime/__init__.py
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/arima/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/arima/config.json
new file mode 100644
index 00000000000..dcdc1335290
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/arima/config.json
@@ -0,0 +1,22 @@
+{
+ "model_type": "sktime",
+ "name": "ARIMA",
+ "predict_length": 1,
+ "order": [1, 0, 0],
+ "seasonal_order": [0, 0, 0, 0],
+ "method": "lbfgs",
+ "maxiter": 1,
+ "suppress_warnings": true,
+ "out_of_sample_size": 0,
+ "scoring": "mse",
+ "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
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py
new file mode 100644
index 00000000000..bd780da3a73
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py
@@ -0,0 +1,379 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+"""
+Sktime model configuration module - simplified version
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Union
+
+from iotdb.ainode.core.exception import (
+ BuiltInModelNotSupportError,
+ ListRangeException,
+ NumericalRangeException,
+ StringRangeException,
+ WrongAttributeTypeError,
+)
+from iotdb.ainode.core.log import Logger
+
+logger = Logger()
+
+
+@dataclass
+class AttributeConfig:
+ """Base class for attribute configuration"""
+
+ name: str
+ default: Any
+ type: str # 'int', 'float', 'str', 'bool', 'list', 'tuple'
+ low: Union[int, float, None] = None
+ high: Union[int, float, None] = None
+ choices: List[str] = field(default_factory=list)
+ value_type: type = None # Element type for list and tuple
+
+ def validate_value(self, value):
+ """Validate if the value meets the requirements"""
+ if self.type == "int":
+ if not isinstance(value, int):
+ raise WrongAttributeTypeError(self.name, "int")
+ if self.low is not None and self.high is not None:
+ if not (self.low <= value <= self.high):
+ raise NumericalRangeException(self.name, value, self.low, self.high)
+ elif self.type == "float":
+ if not isinstance(value, (int, float)):
+ raise WrongAttributeTypeError(self.name, "float")
+ value = float(value)
+ if self.low is not None and self.high is not None:
+ if not (self.low <= value <= self.high):
+ raise NumericalRangeException(self.name, value, self.low, self.high)
+ elif self.type == "str":
+ if not isinstance(value, str):
+ raise WrongAttributeTypeError(self.name, "str")
+ if self.choices and value not in self.choices:
+ raise StringRangeException(self.name, value, self.choices)
+ elif self.type == "bool":
+ if not isinstance(value, bool):
+ raise WrongAttributeTypeError(self.name, "bool")
+ elif self.type == "list":
+ if not isinstance(value, list):
+ raise WrongAttributeTypeError(self.name, "list")
+ for item in value:
+ if not isinstance(item, self.value_type):
+ raise WrongAttributeTypeError(self.name, self.value_type)
+ elif self.type == "tuple":
+ if not isinstance(value, tuple):
+ raise WrongAttributeTypeError(self.name, "tuple")
+ for item in value:
+ if not isinstance(item, self.value_type):
+ raise WrongAttributeTypeError(self.name, self.value_type)
+ return True
+
+ def parse(self, string_value: str):
+ """Parse string value to corresponding type"""
+ if self.type == "int":
+ try:
+ return int(string_value)
+ except:
+ raise WrongAttributeTypeError(self.name, "int")
+ elif self.type == "float":
+ try:
+ return float(string_value)
+ except:
+ raise WrongAttributeTypeError(self.name, "float")
+ elif self.type == "str":
+ return string_value
+ elif self.type == "bool":
+ if string_value.lower() == "true":
+ return True
+ elif string_value.lower() == "false":
+ return False
+ else:
+ raise WrongAttributeTypeError(self.name, "bool")
+ elif self.type == "list":
+ try:
+ list_value = eval(string_value)
+ except:
+ raise WrongAttributeTypeError(self.name, "list")
+ if not isinstance(list_value, list):
+ raise WrongAttributeTypeError(self.name, "list")
+ for i in range(len(list_value)):
+ try:
+ list_value[i] = self.value_type(list_value[i])
+ except:
+ raise ListRangeException(
+ self.name, list_value, str(self.value_type)
+ )
+ return list_value
+ elif self.type == "tuple":
+ try:
+ tuple_value = eval(string_value)
+ except:
+ raise WrongAttributeTypeError(self.name, "tuple")
+ if not isinstance(tuple_value, tuple):
+ raise WrongAttributeTypeError(self.name, "tuple")
+ list_value = list(tuple_value)
+ for i in range(len(list_value)):
+ try:
+ list_value[i] = self.value_type(list_value[i])
+ except:
+ raise ListRangeException(
+ self.name, list_value, str(self.value_type)
+ )
+ return tuple(list_value)
+
+
+# Model configuration definitions - using concise dictionary format
+MODEL_CONFIGS = {
+ "NAIVE_FORECASTER": {
+ "predict_length": AttributeConfig("predict_length", 1, "int", 1, 5000),
+ "pipeline": AttributeConfig(
+ "pipeline", "last", "str", choices=["last", "mean"]
+ ),
+ "sp": AttributeConfig("sp", 1, "int", 1, 5000),
+ },
+ "EXPONENTIAL_SMOOTHING": {
+ "predict_length": AttributeConfig("predict_length", 1, "int", 1, 5000),
+ "damped_trend": AttributeConfig("damped_trend", False, "bool"),
+ "initialization_method": AttributeConfig(
+ "initialization_method",
+ "estimated",
+ "str",
+ choices=["estimated", "heuristic", "legacy-heuristic", "known"],
+ ),
+ "optimized": AttributeConfig("optimized", True, "bool"),
+ "remove_bias": AttributeConfig("remove_bias", False, "bool"),
+ "use_brute": AttributeConfig("use_brute", False, "bool"),
+ },
+ "ARIMA": {
+ "predict_length": AttributeConfig("predict_length", 1, "int", 1, 5000),
+ "order": AttributeConfig("order", (1, 0, 0), "tuple", value_type=int),
+ "seasonal_order": AttributeConfig(
+ "seasonal_order", (0, 0, 0, 0), "tuple", value_type=int
+ ),
+ "method": AttributeConfig(
+ "method",
+ "lbfgs",
+ "str",
+ choices=["lbfgs", "bfgs", "newton", "nm", "cg", "ncg", "powell"],
+ ),
+ "maxiter": AttributeConfig("maxiter", 1, "int", 1, 5000),
+ "suppress_warnings": AttributeConfig("suppress_warnings", True, "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"],
+ ),
+ "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),
+ "sp": AttributeConfig("sp", 2, "int", 1, 5000),
+ "seasonal": AttributeConfig("seasonal", 7, "int", 1, 5000),
+ "seasonal_deg": AttributeConfig("seasonal_deg", 1, "int", 0, 5000),
+ "trend_deg": AttributeConfig("trend_deg", 1, "int", 0, 5000),
+ "low_pass_deg": AttributeConfig("low_pass_deg", 1, "int", 0, 5000),
+ "seasonal_jump": AttributeConfig("seasonal_jump", 1, "int", 0, 5000),
+ "trend_jump": AttributeConfig("trend_jump", 1, "int", 0, 5000),
+ "low_pass_jump": AttributeConfig("low_pass_jump", 1, "int", 0, 5000),
+ },
+ "GAUSSIAN_HMM": {
+ "n_components": AttributeConfig("n_components", 1, "int", 1, 5000),
+ "covariance_type": AttributeConfig(
+ "covariance_type",
+ "diag",
+ "str",
+ choices=["spherical", "diag", "full", "tied"],
+ ),
+ "min_covar": AttributeConfig("min_covar", 1e-3, "float", -1e10, 1e10),
+ "startprob_prior": AttributeConfig(
+ "startprob_prior", 1.0, "float", -1e10, 1e10
+ ),
+ "transmat_prior": AttributeConfig("transmat_prior", 1.0, "float", -1e10, 1e10),
+ "means_prior": AttributeConfig("means_prior", 0.0, "float", -1e10, 1e10),
+ "means_weight": AttributeConfig("means_weight", 0.0, "float", -1e10, 1e10),
+ "covars_prior": AttributeConfig("covars_prior", 1e-2, "float", -1e10, 1e10),
+ "covars_weight": AttributeConfig("covars_weight", 1.0, "float", -1e10, 1e10),
+ "algorithm": AttributeConfig(
+ "algorithm", "viterbi", "str", choices=["viterbi", "map"]
+ ),
+ "n_iter": AttributeConfig("n_iter", 10, "int", 1, 5000),
+ "tol": AttributeConfig("tol", 1e-2, "float", -1e10, 1e10),
+ "params": AttributeConfig("params", "stmc", "str", choices=["stmc", "stm"]),
+ "init_params": AttributeConfig(
+ "init_params", "stmc", "str", choices=["stmc", "stm"]
+ ),
+ "implementation": AttributeConfig(
+ "implementation", "log", "str", choices=["log", "scaling"]
+ ),
+ },
+ "GMM_HMM": {
+ "n_components": AttributeConfig("n_components", 1, "int", 1, 5000),
+ "n_mix": AttributeConfig("n_mix", 1, "int", 1, 5000),
+ "min_covar": AttributeConfig("min_covar", 1e-3, "float", -1e10, 1e10),
+ "startprob_prior": AttributeConfig(
+ "startprob_prior", 1.0, "float", -1e10, 1e10
+ ),
+ "transmat_prior": AttributeConfig("transmat_prior", 1.0, "float", -1e10, 1e10),
+ "weights_prior": AttributeConfig("weights_prior", 1.0, "float", -1e10, 1e10),
+ "means_prior": AttributeConfig("means_prior", 0.0, "float", -1e10, 1e10),
+ "means_weight": AttributeConfig("means_weight", 0.0, "float", -1e10, 1e10),
+ "algorithm": AttributeConfig(
+ "algorithm", "viterbi", "str", choices=["viterbi", "map"]
+ ),
+ "covariance_type": AttributeConfig(
+ "covariance_type",
+ "diag",
+ "str",
+ choices=["sperical", "diag", "full", "tied"],
+ ),
+ "n_iter": AttributeConfig("n_iter", 10, "int", 1, 5000),
+ "tol": AttributeConfig("tol", 1e-2, "float", -1e10, 1e10),
+ "init_params": AttributeConfig(
+ "init_params",
+ "stmcw",
+ "str",
+ choices=[
+ "s",
+ "t",
+ "m",
+ "c",
+ "w",
+ "st",
+ "sm",
+ "sc",
+ "sw",
+ "tm",
+ "tc",
+ "tw",
+ "mc",
+ "mw",
+ "cw",
+ "stm",
+ "stc",
+ "stw",
+ "smc",
+ "smw",
+ "scw",
+ "tmc",
+ "tmw",
+ "tcw",
+ "mcw",
+ "stmc",
+ "stmw",
+ "stcw",
+ "smcw",
+ "tmcw",
+ "stmcw",
+ ],
+ ),
+ "params": AttributeConfig(
+ "params",
+ "stmcw",
+ "str",
+ choices=[
+ "s",
+ "t",
+ "m",
+ "c",
+ "w",
+ "st",
+ "sm",
+ "sc",
+ "sw",
+ "tm",
+ "tc",
+ "tw",
+ "mc",
+ "mw",
+ "cw",
+ "stm",
+ "stc",
+ "stw",
+ "smc",
+ "smw",
+ "scw",
+ "tmc",
+ "tmw",
+ "tcw",
+ "mcw",
+ "stmc",
+ "stmw",
+ "stcw",
+ "smcw",
+ "tmcw",
+ "stmcw",
+ ],
+ ),
+ "implementation": AttributeConfig(
+ "implementation", "log", "str", choices=["log", "scaling"]
+ ),
+ },
+ "STRAY": {
+ "alpha": AttributeConfig("alpha", 0.01, "float", -1e10, 1e10),
+ "k": AttributeConfig("k", 10, "int", 1, 5000),
+ "knn_algorithm": AttributeConfig(
+ "knn_algorithm",
+ "brute",
+ "str",
+ choices=["brute", "kd_tree", "ball_tree", "auto"],
+ ),
+ "p": AttributeConfig("p", 0.5, "float", -1e10, 1e10),
+ "size_threshold": AttributeConfig("size_threshold", 50, "int", 1, 5000),
+ "outlier_tail": AttributeConfig(
+ "outlier_tail", "max", "str", choices=["min", "max"]
+ ),
+ },
+}
+
+
+def get_attributes(model_id: str) -> Dict[str, AttributeConfig]:
+ """Get attribute configuration for Sktime model"""
+ model_id = "EXPONENTIAL_SMOOTHING" if model_id == "HOLTWINTERS" else model_id
+ if model_id not in MODEL_CONFIGS:
+ raise BuiltInModelNotSupportError(model_id)
+ return MODEL_CONFIGS[model_id]
+
+
+def update_attribute(
+ input_attributes: Dict[str, str], attribute_map: Dict[str, AttributeConfig]
+) -> Dict[str, Any]:
+ """Update Sktime model attributes using input attributes"""
+ attributes = {}
+ for name, config in attribute_map.items():
+ if name in input_attributes:
+ value = config.parse(input_attributes[name])
+ config.validate_value(value)
+ attributes[name] = value
+ else:
+ attributes[name] = config.default
+ return attributes
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/exponential_smoothing/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/exponential_smoothing/config.json
new file mode 100644
index 00000000000..d6002fb26e8
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/exponential_smoothing/config.json
@@ -0,0 +1,11 @@
+{
+ "model_type": "sktime",
+ "name": "ExponentialSmoothing",
+ "predict_length": 1,
+ "damped_trend": false,
+ "initialization_method": "estimated",
+ "optimized": true,
+ "remove_bias": false,
+ "use_brute": false
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/gaussian_hmm/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/gaussian_hmm/config.json
new file mode 100644
index 00000000000..3392e1c0b57
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/gaussian_hmm/config.json
@@ -0,0 +1,20 @@
+{
+ "model_type": "sktime",
+ "name": "GaussianHMM",
+ "n_components": 1,
+ "covariance_type": "diag",
+ "min_covar": 0.001,
+ "startprob_prior": 1.0,
+ "transmat_prior": 1.0,
+ "means_prior": 0.0,
+ "means_weight": 0.0,
+ "covars_prior": 0.01,
+ "covars_weight": 1.0,
+ "algorithm": "viterbi",
+ "n_iter": 10,
+ "tol": 0.01,
+ "params": "stmc",
+ "init_params": "stmc",
+ "implementation": "log"
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/gmm_hmm/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/gmm_hmm/config.json
new file mode 100644
index 00000000000..235f8ae642d
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/gmm_hmm/config.json
@@ -0,0 +1,20 @@
+{
+ "model_type": "sktime",
+ "name": "GMMHMM",
+ "n_components": 1,
+ "n_mix": 1,
+ "min_covar": 0.001,
+ "startprob_prior": 1.0,
+ "transmat_prior": 1.0,
+ "weights_prior": 1.0,
+ "means_prior": 0.0,
+ "means_weight": 0.0,
+ "algorithm": "viterbi",
+ "covariance_type": "diag",
+ "n_iter": 10,
+ "tol": 0.01,
+ "init_params": "stmcw",
+ "params": "stmcw",
+ "implementation": "log"
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py
new file mode 100644
index 00000000000..f272e3dda35
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/modeling_sktime.py
@@ -0,0 +1,178 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+"""
+Sktime model implementation module - simplified version
+"""
+
+from abc import abstractmethod
+from typing import Any, Dict
+
+import numpy as np
+from sklearn.preprocessing import MinMaxScaler
+from sktime.detection.hmm_learn import GMMHMM, GaussianHMM
+from sktime.detection.stray import STRAY
+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
+
+from iotdb.ainode.core.exception import (
+ BuiltInModelNotSupportError,
+ InferenceModelInternalError,
+)
+from iotdb.ainode.core.log import Logger
+
+from .configuration_sktime import get_attributes, update_attribute
+
+logger = Logger()
+
+
+class SktimeModel:
+ """Base class for Sktime models"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ self._attributes = attributes
+ self._model = None
+
+ @abstractmethod
+ def generate(self, data):
+ """Execute generation/inference"""
+ raise NotImplementedError
+
+
+class ForecastingModel(SktimeModel):
+ """Base class for forecasting models"""
+
+ def generate(self, data):
+ """Execute forecasting"""
+ try:
+ predict_length = self._attributes["predict_length"]
+ self._model.fit(data)
+ output = self._model.predict(fh=range(predict_length))
+ return np.array(output, dtype=np.float64)
+ except Exception as e:
+ raise InferenceModelInternalError(str(e))
+
+
+class DetectionModel(SktimeModel):
+ """Base class for detection models"""
+
+ def generate(self, data):
+ """Execute detection"""
+ try:
+ self._model.fit(data)
+ output = self._model.predict(data)
+ return np.array(output, dtype=np.int32)
+ except Exception as e:
+ raise InferenceModelInternalError(str(e))
+
+
+class ArimaModel(ForecastingModel):
+ """ARIMA model"""
+
+ 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"}
+ )
+
+
+class ExponentialSmoothingModel(ForecastingModel):
+ """Exponential smoothing model"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ super().__init__(attributes)
+ self._model = ExponentialSmoothing(
+ **{k: v for k, v in attributes.items() if k != "predict_length"}
+ )
+
+
+class NaiveForecasterModel(ForecastingModel):
+ """Naive forecaster model"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ super().__init__(attributes)
+ self._model = NaiveForecaster(
+ **{k: v for k, v in attributes.items() if k != "predict_length"}
+ )
+
+
+class STLForecasterModel(ForecastingModel):
+ """STL forecaster model"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ super().__init__(attributes)
+ self._model = STLForecaster(
+ **{k: v for k, v in attributes.items() if k != "predict_length"}
+ )
+
+
+class GMMHMMModel(DetectionModel):
+ """GMM HMM model"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ super().__init__(attributes)
+ self._model = GMMHMM(**attributes)
+
+
+class GaussianHmmModel(DetectionModel):
+ """Gaussian HMM model"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ super().__init__(attributes)
+ self._model = GaussianHMM(**attributes)
+
+
+class STRAYModel(DetectionModel):
+ """STRAY anomaly detection model"""
+
+ def __init__(self, attributes: Dict[str, Any]):
+ super().__init__(attributes)
+ self._model = STRAY(**attributes)
+
+ def generate(self, data):
+ """STRAY requires special handling: normalize first"""
+ try:
+ data = MinMaxScaler().fit_transform(data)
+ output = self._model.fit_transform(data)
+ return np.array(output, dtype=np.int32)
+ except Exception as e:
+ raise InferenceModelInternalError(str(e))
+
+
+# Model factory mapping
+_MODEL_FACTORY = {
+ "ARIMA": ArimaModel,
+ "EXPONENTIAL_SMOOTHING": ExponentialSmoothingModel,
+ "HOLTWINTERS": ExponentialSmoothingModel, # Use the same model class
+ "NAIVE_FORECASTER": NaiveForecasterModel,
+ "STL_FORECASTER": STLForecasterModel,
+ "GMM_HMM": GMMHMMModel,
+ "GAUSSIAN_HMM": GaussianHmmModel,
+ "STRAY": STRAYModel,
+}
+
+
+def create_sktime_model(model_id: str, **kwargs) -> SktimeModel:
+ """Create a Sktime model instance"""
+ attributes = update_attribute({**kwargs}, get_attributes(model_id.upper()))
+ model_class = _MODEL_FACTORY.get(model_id.upper())
+ if model_class is None:
+ raise BuiltInModelNotSupportError(model_id)
+ return model_class(attributes)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/naive_forecaster/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/naive_forecaster/config.json
new file mode 100644
index 00000000000..20d8c1ed32b
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/naive_forecaster/config.json
@@ -0,0 +1,8 @@
+{
+ "model_type": "sktime",
+ "name": "NaiveForecaster",
+ "predict_length": 1,
+ "strategy": "last",
+ "sp": 1
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/stl_forecaster/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/stl_forecaster/config.json
new file mode 100644
index 00000000000..1005f9d944e
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/stl_forecaster/config.json
@@ -0,0 +1,14 @@
+{
+ "model_type": "sktime",
+ "name": "STLForecaster",
+ "predict_length": 1,
+ "sp": 2,
+ "seasonal": 7,
+ "seasonal_deg": 1,
+ "trend_deg": 1,
+ "low_pass_deg": 1,
+ "seasonal_jump": 1,
+ "trend_jump": 1,
+ "low_pass_jump": 1
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sktime/stray/config.json b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/stray/config.json
new file mode 100644
index 00000000000..64c64aa9e05
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sktime/stray/config.json
@@ -0,0 +1,11 @@
+{
+ "model_type": "sktime",
+ "name": "STRAY",
+ "alpha": 0.01,
+ "k": 10,
+ "knn_algorithm": "brute",
+ "p": 0.5,
+ "size_threshold": 50,
+ "outlier_tail": "max"
+}
+
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/__init__.py
similarity index 100%
copy from iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
copy to iotdb-core/ainode/iotdb/ainode/core/model/sundial/__init__.py
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/configuration_sundial.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/configuration_sundial.py
new file mode 100644
index 00000000000..21eefef2933
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/configuration_sundial.py
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+from transformers import PretrainedConfig
+
+
+class SundialConfig(PretrainedConfig):
+ model_type = "sundial"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ input_token_len: int = 16,
+ hidden_size: int = 768,
+ intermediate_size: int = 3072,
+ output_token_lens: List[int] = [720],
+ num_hidden_layers: int = 12,
+ num_attention_heads: int = 12,
+ hidden_act: str = "silu",
+ use_cache: bool = True,
+ rope_theta: int = 10000,
+ dropout_rate: float = 0.1,
+ initializer_range: float = 0.02,
+ max_position_embeddings: int = 10000,
+ flow_loss_depth: int = 3,
+ num_sampling_steps: int = 50,
+ diffusion_batch_mul: int = 4,
+ **kwargs,
+ ):
+ self.input_token_len = input_token_len
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.hidden_act = hidden_act
+ self.output_token_lens = output_token_lens
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.dropout_rate = dropout_rate
+ self.initializer_range = initializer_range
+ self.max_position_embeddings = max_position_embeddings
+ self.flow_loss_depth = flow_loss_depth
+ self.num_sampling_steps = num_sampling_steps
+ self.diffusion_batch_mul = diffusion_batch_mul
+
+ super().__init__(
+ **kwargs,
+ )
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py
new file mode 100644
index 00000000000..b3fe95dbe2d
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/flow_loss.py
@@ -0,0 +1,255 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class FlowLoss(nn.Module):
+ """Flow Loss"""
+
+ def __init__(self, target_channels, z_channels, depth, width, num_sampling_steps):
+ super(FlowLoss, self).__init__()
+ self.in_channels = target_channels
+ self.net = SimpleMLPAdaLN(
+ in_channels=target_channels,
+ model_channels=width,
+ out_channels=target_channels,
+ z_channels=z_channels,
+ num_res_blocks=depth,
+ )
+ self.num_sampling_steps = num_sampling_steps
+
+ def forward(self, target, z, mask=None, mask_y=None):
+ noise = torch.randn_like(target)
+ t = torch.rand(target.shape[0], device=target.device)
+
+ noised_target = t[:, None] * target + (1 - t[:, None]) * noise
+
+ predict_v = self.net(noised_target, t * 1000, z)
+
+ weights = 1.0 / torch.arange(
+ 1, self.in_channels + 1, dtype=torch.float32, device=target.device
+ )
+ if mask_y is not None:
+ loss = (mask_y * weights * (predict_v - target) ** 2).sum(dim=-1)
+ else:
+ loss = (weights * (predict_v - target) ** 2).sum(dim=-1)
+
+ if mask is not None:
+ loss = (loss * mask).sum() / mask.sum()
+ return loss.mean()
+
+ def sample(self, z, num_samples=1):
+ z = z.repeat(num_samples, 1)
+ noise = torch.randn(z.shape[0], self.in_channels).to(z.device)
+ x = noise
+ dt = 1.0 / self.num_sampling_steps
+ for i in range(self.num_sampling_steps):
+ t = (torch.ones((x.shape[0])) * i / self.num_sampling_steps).to(x.device)
+ pred = self.net(x, t * 1000, z)
+ x = x + (pred - noise) * dt
+ x = x.reshape(num_samples, -1, self.in_channels).transpose(0, 1)
+ return x
+
+
+def modulate(x, shift, scale):
+ return x * (1 + scale) + shift
+
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period)
+ * torch.arange(start=0, end=half, dtype=torch.float32)
+ / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat(
+ [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+ )
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class ResBlock(nn.Module):
+ """
+ A residual block that can optionally change the number of channels.
+ :param channels: the number of input channels.
+ """
+
+ def __init__(self, channels):
+ super().__init__()
+ self.channels = channels
+
+ self.in_ln = nn.LayerNorm(channels, eps=1e-6)
+ self.mlp = nn.Sequential(
+ nn.Linear(channels, channels, bias=True),
+ nn.SiLU(),
+ nn.Linear(channels, channels, bias=True),
+ )
+
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
+ )
+
+ def forward(self, x, y):
+ shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
+ h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
+ h = self.mlp(h)
+ return x + gate_mlp * h
+
+
+class FinalLayer(nn.Module):
+ """
+ The final layer adopted from DiT.
+ """
+
+ def __init__(self, model_channels, out_channels):
+ super().__init__()
+ self.norm_final = nn.LayerNorm(
+ model_channels, elementwise_affine=False, eps=1e-6
+ )
+ self.linear = nn.Linear(model_channels, out_channels, bias=True)
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True)
+ )
+
+ def forward(self, x, c):
+ shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+ x = modulate(self.norm_final(x), shift, scale)
+ x = self.linear(x)
+ return x
+
+
+class SimpleMLPAdaLN(nn.Module):
+ """
+ The MLP for Diffusion Loss.
+ :param in_channels: channels in the input Tensor.
+ :param model_channels: base channel count for the model.
+ :param out_channels: channels in the output Tensor.
+ :param z_channels: channels in the condition.
+ :param num_res_blocks: number of residual blocks per downsample.
+ """
+
+ def __init__(
+ self,
+ in_channels,
+ model_channels,
+ out_channels,
+ z_channels,
+ num_res_blocks,
+ ):
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.out_channels = out_channels
+ self.num_res_blocks = num_res_blocks
+
+ self.time_embed = TimestepEmbedder(model_channels)
+ self.cond_embed = nn.Linear(z_channels, model_channels)
+
+ self.input_proj = nn.Linear(in_channels, model_channels)
+
+ res_blocks = []
+ for i in range(num_res_blocks):
+ res_blocks.append(
+ ResBlock(
+ model_channels,
+ )
+ )
+
+ self.res_blocks = nn.ModuleList(res_blocks)
+ self.final_layer = FinalLayer(model_channels, out_channels)
+
+ self.initialize_weights()
+
+ def initialize_weights(self):
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+
+ self.apply(_basic_init)
+
+ # Initialize timestep embedding MLP
+ nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
+
+ # Zero-out adaLN modulation layers
+ for block in self.res_blocks:
+ nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+ # Zero-out output layers
+ nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+ nn.init.constant_(self.final_layer.linear.weight, 0)
+ nn.init.constant_(self.final_layer.linear.bias, 0)
+
+ def forward(self, x, t, c):
+ """
+ Apply the model to an input batch.
+ :param x: an [N x C] Tensor of inputs.
+ :param t: a 1-D batch of timesteps.
+ :param c: conditioning from AR transformer.
+ :return: an [N x C] Tensor of outputs.
+ """
+ x = self.input_proj(x)
+ t = self.time_embed(t)
+ c = self.cond_embed(c)
+ y = t + c
+
+ for block in self.res_blocks:
+ x = block(x, y)
+
+ return self.final_layer(x, y)
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py
new file mode 100644
index 00000000000..544193e4d9c
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/modeling_sundial.py
@@ -0,0 +1,651 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import os
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file as load_safetensors
+from torch import nn
+from transformers import Cache, DynamicCache, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.modeling_outputs import (
+ MoeCausalLMOutputWithPast,
+ MoeModelOutputWithPast,
+)
+
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.model.sundial.configuration_sundial import SundialConfig
+from iotdb.ainode.core.model.sundial.flow_loss import FlowLoss
+from iotdb.ainode.core.model.sundial.ts_generation_mixin import TSGenerationMixin
+
+logger = Logger()
+
+
+def rotate_half(x):
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class SundialPatchEmbedding(nn.Module):
+ def __init__(self, config: SundialConfig):
+ super().__init__()
+ self.dropout = nn.Dropout(config.dropout_rate)
+ self.hidden_layer = nn.Linear(
+ config.input_token_len * 2, config.intermediate_size
+ )
+ self.act = ACT2FN[config.hidden_act]
+ self.output_layer = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.residual_layer = nn.Linear(config.input_token_len * 2, config.hidden_size)
+ self.input_token_len = config.input_token_len
+
+ def forward(self, x):
+ mask = torch.ones_like(x, dtype=torch.float32)
+ input_length = x.shape[-1]
+ padding_length = (
+ self.input_token_len - (input_length % self.input_token_len)
+ ) % self.input_token_len
+ x = F.pad(x, (padding_length, 0))
+ mask = F.pad(mask, (padding_length, 0))
+ x = x.unfold(dimension=-1, size=self.input_token_len, step=self.input_token_len)
+ mask = mask.unfold(
+ dimension=-1, size=self.input_token_len, step=self.input_token_len
+ )
+
+ x = torch.cat([x, mask], dim=-1)
+ hid = self.act(self.hidden_layer(x))
+ out = self.dropout(self.output_layer(hid))
+ res = self.residual_layer(x)
+ out = out + res
+ return out
+
+
+class SundialRotaryEmbedding(torch.nn.Module):
+ def __init__(self, dim, max_position_embeddings=10000, base=10000, device=None):
+ super().__init__()
+ self.dim = dim
+ self.max_position_embeddings = max_position_embeddings
+ self.base = base
+ inv_freq = 1.0 / (
+ self.base
+ ** (
+ torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
+ / self.dim
+ )
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ # Build here to make `torch.jit.trace` work.
+ self._set_cos_sin_cache(
+ seq_len=max_position_embeddings,
+ device=self.inv_freq.device,
+ dtype=torch.get_default_dtype(),
+ )
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(
+ self.max_seq_len_cached, device=device, dtype=torch.int64
+ ).type_as(self.inv_freq)
+
+ freqs = torch.outer(t, self.inv_freq)
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+ def forward(self, x, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ if seq_len > self.max_seq_len_cached:
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+ return (
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
+ )
+
+
+class SundialAttention(nn.Module):
+ def __init__(self, config: SundialConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.hidden_size // self.num_heads
+ self.attention_dropout = config.dropout_rate
+ self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+ self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+ self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+ self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+ self.rotary_emb = SundialRotaryEmbedding(
+ self.head_dim, max_position_embeddings=config.max_position_embeddings
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.view(
+ bsz, q_len, self.num_heads, self.head_dim
+ ).transpose(1, 2)
+ key_states = key_states.view(
+ bsz, q_len, self.num_heads, self.head_dim
+ ).transpose(1, 2)
+ value_states = value_states.view(
+ bsz, q_len, self.num_heads, self.head_dim
+ ).transpose(1, 2)
+
+ kv_seq_len = key_states.shape[-2]
+ if past_key_value is not None:
+ kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+ query_states, key_states = apply_rotary_pos_emb(
+ query_states, key_states, cos, sin, position_ids
+ )
+
+ if past_key_value is not None:
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx
+ )
+
+ attn_output = F.scaled_dot_product_attention(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout_p=(self.attention_dropout if self.training else 0.0),
+ )
+
+ attn_output = attn_output.transpose(1, 2).contiguous()
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+class SundialMLP(nn.Module):
+ def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = ACT2FN[hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(
+ self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)
+ )
+
+
+class SundialDecoderLayer(nn.Module):
+ def __init__(self, config: SundialConfig, layer_idx: int):
+ super().__init__()
+ self.self_attn = SundialAttention(config, layer_idx)
+
+ self.ffn_layer = SundialMLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.intermediate_size,
+ hidden_act=config.hidden_act,
+ )
+ self.norm1 = torch.nn.LayerNorm(config.hidden_size)
+ self.norm2 = torch.nn.LayerNorm(config.hidden_size)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: Optional[bool] = False,
+ **kwargs,
+ ) -> Tuple[
+ torch.FloatTensor,
+ Optional[torch.Tensor],
+ Optional[Cache],
+ ]:
+ residual = hidden_states
+
+ hidden_states = self.norm1(hidden_states)
+
+ # Self Attention
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ output_attentions=output_attentions,
+ )
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.norm2(hidden_states)
+ hidden_states = self.ffn_layer(hidden_states)
+ hidden_states = residual + hidden_states
+
+ if not output_attentions:
+ self_attn_weights = None
+
+ return hidden_states, self_attn_weights, present_key_value
+
+
+class SundialPreTrainedModel(PreTrainedModel):
+ config_class = SundialConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["SundialDecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn_2 = True
+ _supports_sdpa = False
+ _supports_cache_class = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, torch.nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, torch.nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+
+class SundialModel(SundialPreTrainedModel):
+ def __init__(self, config: SundialConfig):
+ super().__init__(config)
+ self.embed_layer = SundialPatchEmbedding(config)
+ self.layers = nn.ModuleList(
+ [
+ SundialDecoderLayer(config, layer_idx)
+ for layer_idx in range(config.num_hidden_layers)
+ ]
+ )
+ self.norm = torch.nn.LayerNorm(config.hidden_size)
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ input_ids: torch.FloatTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[
+ Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor]]]
+ ] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
+ # input_ids is the input of time series, its shape is [batch_size, seq_len]
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
+ )
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape
+ elif inputs_embeds is not None:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ else:
+ raise ValueError(
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_layer(input_ids)
+ seq_length = inputs_embeds.shape[1]
+
+ past_key_values_length = 0
+ use_legacy_cache = False
+
+ if past_key_values is not None:
+ use_legacy_cache = not isinstance(past_key_values, Cache)
+ # Converts the legacy cache which is tuple into an equivalent Cache. Used for backward compatibility.
+ if use_legacy_cache:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ # Suppose the sequence length of each layer is the same
+ past_key_values_length = past_key_values.get_seq_length()
+
+ # When training + checkpoints, caching is usually disabled (just do not transfer)
+ if (
+ self.gradient_checkpointing
+ and self.training
+ and isinstance(past_key_values, Cache)
+ ):
+ past_key_values = None
+ past_key_values_length = 0
+
+ if position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length,
+ seq_length + past_key_values_length,
+ dtype=torch.long,
+ device=device,
+ )
+ # position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+ position_ids = position_ids.view(-1, seq_length)
+ else:
+ position_ids = position_ids.view(-1, seq_length).long()
+
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ sliding_window=None,
+ )
+
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ if isinstance(past_key_values, Cache):
+ next_decoder_cache = layer_outputs[2]
+
+ hidden_states = self.norm(hidden_states)
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = None
+ if isinstance(past_key_values, Cache):
+ next_cache = (
+ next_decoder_cache.to_legacy_cache()
+ if use_legacy_cache
+ else next_decoder_cache
+ )
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+ if v is not None
+ )
+ return MoeModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class SundialForPrediction(SundialPreTrainedModel, TSGenerationMixin):
+ def __init__(self, config: SundialConfig):
+ super().__init__(config)
+ self.config = config
+ self.model = SundialModel(self.config)
+ self.flow_loss = FlowLoss(
+ self.config.output_token_lens[-1],
+ self.config.hidden_size,
+ self.config.flow_loss_depth,
+ self.config.hidden_size,
+ self.config.num_sampling_steps,
+ )
+ self.post_init()
+
+ def set_decoder(self, decoder):
+ self.model = decoder
+
+ def get_decoder(self):
+ return self.model
+
+ def forward(
+ self,
+ input_ids: torch.FloatTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[
+ Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor]]]
+ ] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.FloatTensor] = None,
+ loss_masks: Optional[torch.FloatTensor] = None,
+ mask_y: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ max_output_length: Optional[int] = None,
+ revin: Optional[bool] = False,
+ num_samples: Optional[int] = 1,
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ if revin:
+ means = input_ids.mean(1, keepdim=True).detach()
+ stdev = input_ids.std(dim=1, keepdim=True, unbiased=False).detach()
+ stdev = torch.where(
+ stdev > 1e-2, stdev, torch.tensor(1.0, device=input_ids.device)
+ )
+ input_ids = (input_ids - means) / stdev
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+ predictions = None
+
+ loss = None
+ if labels is not None:
+ if revin:
+ labels = (labels - means) / stdev
+ output_token_len = self.config.output_token_lens[-1]
+ seq_len = hidden_states.shape[1] * self.config.input_token_len
+ labels = labels[
+ :, : seq_len - self.config.input_token_len + output_token_len
+ ]
+ shift_labels = labels.unfold(
+ dimension=-1, size=output_token_len, step=self.config.input_token_len
+ )
+
+ bsz, L, _ = shift_labels.shape
+ shift_labels = shift_labels.reshape(bsz * L, -1).repeat(
+ self.config.diffusion_batch_mul, 1
+ )
+ hidden_states = hidden_states.reshape(bsz * L, -1).repeat(
+ self.config.diffusion_batch_mul, 1
+ )
+ loss_masks = loss_masks.reshape(bsz * L).repeat(
+ self.config.diffusion_batch_mul
+ )
+ mask_y = mask_y.repeat(L * self.config.diffusion_batch_mul, 1)
+
+ loss = self.flow_loss(shift_labels, hidden_states, loss_masks, mask_y)
+ else:
+ if max_output_length is None:
+ output_token_len = self.config.output_token_lens[0]
+ max_output_length = output_token_len
+ else:
+ output_token_len = self.config.output_token_lens[0]
+ for h in self.config.output_token_lens[1:]:
+ if h > max_output_length:
+ break
+ else:
+ output_token_len = h
+
+ bsz = hidden_states.shape[0]
+ hidden_states = hidden_states[:, -1, :]
+ predictions = self.flow_loss.sample(hidden_states, num_samples)
+ if output_token_len > max_output_length:
+ predictions = predictions[:, :, :max_output_length]
+ if revin:
+ predictions = predictions * stdev + means
+ if not return_dict:
+ output = (predictions,) + outputs[1:]
+ return (loss) + output if loss is not None else output
+
+ return MoeCausalLMOutputWithPast(
+ loss=loss,
+ logits=predictions,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ revin=False,
+ num_samples=1,
+ **kwargs,
+ ):
+ # Omit tokens covered by past_key_values
+ if past_key_values is not None:
+ if isinstance(past_key_values, Cache):
+ past_length = past_key_values.get_seq_length()
+ else:
+ past_length = past_key_values[0][0].shape[2]
+
+ # Keep only the unprocessed tokens:
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+ # input)
+ if attention_mask is not None and attention_mask.shape[1] > (
+ input_ids.shape[1] // self.config.input_token_len
+ ):
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+ # input_ids based on the past_length.
+ elif past_length < (input_ids.shape[1] // self.config.input_token_len):
+ input_ids = input_ids[:, past_length * self.config.input_token_len :]
+ # 3 - Otherwise (past_length >= (input_ids.shape[1] // self.config.input_token_len)), let's assume input_ids only has unprocessed tokens.
+
+ position_ids = kwargs.get("position_ids", None)
+ if attention_mask is not None and position_ids is None:
+ # create position_ids on the fly for batch generation
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ if past_key_values:
+ position_ids = position_ids[
+ :, -(input_ids.shape[1] // self.config.input_token_len) :
+ ]
+
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+ if inputs_embeds is not None and past_key_values is None:
+ model_inputs = {"inputs_embeds": inputs_embeds}
+ else:
+ model_inputs = {"input_ids": input_ids}
+
+ model_inputs.update(
+ {
+ "position_ids": position_ids,
+ "past_key_values": past_key_values,
+ "attention_mask": attention_mask,
+ "revin": revin,
+ "num_samples": num_samples,
+ }
+ )
+ return model_inputs
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/sundial/ts_generation_mixin.py b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/ts_generation_mixin.py
new file mode 100644
index 00000000000..f09621f2cb0
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/sundial/ts_generation_mixin.py
@@ -0,0 +1,383 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import warnings
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from transformers import GenerationMixin, LogitsProcessorList, StoppingCriteriaList
+from transformers.generation import EosTokenCriteria, validate_stopping_criteria
+from transformers.generation.utils import (
+ GenerateDecoderOnlyOutput,
+ GenerateEncoderDecoderOutput,
+ GenerateNonBeamOutput,
+ GenerateOutput,
+ GenerationConfig,
+)
+from transformers.utils import ModelOutput
+
+
+class TSGenerationMixin(GenerationMixin):
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ prefix_allowed_tokens_fn: Optional[
+ Callable[[int, torch.Tensor], List[int]]
+ ] = None,
+ synced_gpus: Optional[bool] = None,
+ assistant_model: Optional["PreTrainedModel"] = None,
+ streamer: Optional["BaseStreamer"] = None,
+ negative_prompt_ids: Optional[torch.Tensor] = None,
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ revin: Optional[bool] = True,
+ num_samples: Optional[int] = 1,
+ **kwargs,
+ ) -> Union[GenerateOutput, torch.LongTensor]:
+ if len(inputs.shape) != 2:
+ raise ValueError("Input shape must be: [batch_size, seq_len]")
+ batch_size, cur_len = inputs.shape
+ if cur_len < self.config.input_token_len:
+ raise ValueError(
+ f"Input length must be at least {self.config.input_token_len}"
+ )
+ if revin:
+ means = inputs.mean(dim=-1, keepdim=True)
+ stdev = inputs.std(dim=-1, keepdim=True, unbiased=False) + 1e-5
+ inputs = (inputs - means) / stdev
+ outputs = super().generate(
+ inputs=inputs,
+ generation_config=generation_config,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+ synced_gpus=synced_gpus,
+ assistant_model=assistant_model,
+ streamer=streamer,
+ negative_prompt_ids=negative_prompt_ids,
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
+ num_samples=num_samples,
+ **kwargs,
+ )
+ if revin:
+ stdev = stdev.unsqueeze(1).repeat(1, num_samples, 1)
+ means = means.unsqueeze(1).repeat(1, num_samples, 1)
+ outputs = (outputs * stdev) + means
+ return outputs
+
+ def _sample(
+ self,
+ input_ids: torch.Tensor,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ max_length: Optional[int] = None,
+ pad_token_id: Optional[int] = None,
+ eos_token_id: Optional[Union[int, List[int]]] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_scores: Optional[bool] = None,
+ output_logits: Optional[bool] = None,
+ return_dict_in_generate: Optional[bool] = None,
+ synced_gpus: bool = False,
+ streamer: Optional["BaseStreamer"] = None,
+ **model_kwargs,
+ ) -> Union[GenerateNonBeamOutput, torch.Tensor]:
+ input_ids = input_ids.to(self.device)
+ batch_size, cur_len = input_ids.shape
+ # init values
+ logits_processor = (
+ logits_processor if logits_processor is not None else LogitsProcessorList()
+ )
+ stopping_criteria = (
+ stopping_criteria
+ if stopping_criteria is not None
+ else StoppingCriteriaList()
+ )
+ if max_length is not None:
+ warnings.warn(
+ "`max_length` is deprecated in this function, use"
+ " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
+ UserWarning,
+ )
+ stopping_criteria = validate_stopping_criteria(
+ stopping_criteria, max_length
+ )
+ pad_token_id = (
+ pad_token_id
+ if pad_token_id is not None
+ else self.generation_config.pad_token_id
+ )
+ if eos_token_id is not None:
+ stopping_criteria.append(EosTokenCriteria(eos_token_id=eos_token_id))
+ else:
+ # remove when the method is totally private
+ # need to get `eos_token_id` and add stopping criteria, so that generation does not go forever
+ eos_token_id = [
+ criteria.eos_token_id.tolist()
+ for criteria in stopping_criteria
+ if hasattr(criteria, "eos_token_id")
+ ]
+ eos_token_id = eos_token_id[0] if eos_token_id else None
+ if eos_token_id is None and self.generation_config.eos_token_id is not None:
+ eos_token_id = self.generation_config.eos_token_id
+ stopping_criteria.append(EosTokenCriteria(eos_token_id=eos_token_id))
+
+ if isinstance(eos_token_id, int):
+ eos_token_id = [eos_token_id]
+ output_scores = (
+ output_scores
+ if output_scores is not None
+ else self.generation_config.output_scores
+ )
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.generation_config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.generation_config.output_hidden_states
+ )
+ return_dict_in_generate = (
+ return_dict_in_generate
+ if return_dict_in_generate is not None
+ else self.generation_config.return_dict_in_generate
+ )
+
+ # init attention / hidden states / scores tuples
+ raw_logits = () if (return_dict_in_generate and output_logits) else None
+ scores = () if (return_dict_in_generate and output_scores) else None
+ decoder_attentions = (
+ () if (return_dict_in_generate and output_attentions) else None
+ )
+ cross_attentions = (
+ () if (return_dict_in_generate and output_attentions) else None
+ )
+ decoder_hidden_states = (
+ () if (return_dict_in_generate and output_hidden_states) else None
+ )
+
+ # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+ if return_dict_in_generate and self.config.is_encoder_decoder:
+ encoder_attentions = (
+ model_kwargs["encoder_outputs"].get("attentions")
+ if output_attentions
+ else None
+ )
+ encoder_hidden_states = (
+ model_kwargs["encoder_outputs"].get("hidden_states")
+ if output_hidden_states
+ else None
+ )
+
+ # keep track of which sequences are already finished
+ if "inputs_embeds" in model_kwargs:
+ cur_len = model_kwargs["inputs_embeds"].shape[1]
+ this_peer_finished = False
+ unfinished_sequences = torch.ones(
+ batch_size, dtype=torch.long, device=input_ids.device
+ )
+ model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
+ true_seq_len = (
+ cur_len + self.config.input_token_len - 1
+ ) // self.config.input_token_len
+ model_kwargs["attention_mask"] = model_kwargs["attention_mask"][
+ :, -true_seq_len:
+ ]
+ max_length = stopping_criteria.max_length
+ generate_results = None
+ while self._has_unfinished_sequences(
+ this_peer_finished, synced_gpus, device=input_ids.device
+ ):
+ # prepare model inputs
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+
+ input_length = input_ids.shape[1]
+
+ # forward pass to get next token
+ outputs = self(
+ **model_inputs,
+ return_dict=True,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ max_output_length=max_length - input_length,
+ )
+
+ if synced_gpus and this_peer_finished:
+ continue # don't waste resources running the code we don't need
+ next_token_logits = outputs.logits
+
+ # pre-process distribution
+ next_tokens_scores = logits_processor(input_ids, next_token_logits)
+
+ # Store scores, attentions and hidden_states when required
+ if return_dict_in_generate:
+ if output_scores:
+ scores += (next_tokens_scores,)
+ if output_logits:
+ raw_logits += (next_token_logits,)
+ if output_attentions:
+ decoder_attentions += (
+ (outputs.decoder_attentions,)
+ if self.config.is_encoder_decoder
+ else (outputs.attentions,)
+ )
+ if self.config.is_encoder_decoder:
+ cross_attentions += (outputs.cross_attentions,)
+
+ if output_hidden_states:
+ decoder_hidden_states += (
+ (outputs.decoder_hidden_states,)
+ if self.config.is_encoder_decoder
+ else (outputs.hidden_states,)
+ )
+
+ # argmax
+ # next_tokens = torch.argmax(next_tokens_scores, dim=-1)
+ next_tokens = next_tokens_scores
+
+ # finished sentences should have their next token be a padding token
+ if eos_token_id is not None:
+ if pad_token_id is None:
+ raise ValueError(
+ "If `eos_token_id` is defined, make sure that `pad_token_id` is defined."
+ )
+ next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
+ 1 - unfinished_sequences
+ )
+
+ # update generated ids, model inputs, and length for next step
+ horizon_length = next_tokens.shape[-1] // self.config.input_token_len
+
+ past_key_values = model_kwargs.get("past_key_values")
+ if past_key_values is None or generate_results is None:
+ generate_results = next_tokens
+ else:
+ generate_results = torch.cat([generate_results, next_tokens], dim=-1)
+ input_ids = torch.cat([input_ids, next_tokens.median(dim=1)[0]], dim=-1)
+
+ if streamer is not None:
+ streamer.put(next_tokens.cpu())
+ model_kwargs = self._update_model_kwargs_for_generation(
+ outputs,
+ model_kwargs,
+ horizon_length=horizon_length,
+ is_encoder_decoder=self.config.is_encoder_decoder,
+ )
+ unfinished_sequences = unfinished_sequences & ~stopping_criteria(
+ input_ids, scores
+ )
+ this_peer_finished = unfinished_sequences.max() == 0
+
+ if input_ids.shape[-1] > max_length:
+ input_ids = input_ids[:, :max_length]
+
+ if streamer is not None:
+ streamer.end()
+
+ if return_dict_in_generate:
+ if self.config.is_encoder_decoder:
+ return GenerateEncoderDecoderOutput(
+ sequences=input_ids,
+ scores=scores,
+ logits=raw_logits,
+ encoder_attentions=encoder_attentions,
+ encoder_hidden_states=encoder_hidden_states,
+ decoder_attentions=decoder_attentions,
+ cross_attentions=cross_attentions,
+ decoder_hidden_states=decoder_hidden_states,
+ past_key_values=model_kwargs.get("past_key_values"),
+ )
+ else:
+ return GenerateDecoderOnlyOutput(
+ sequences=input_ids,
+ scores=scores,
+ logits=raw_logits,
+ attentions=decoder_attentions,
+ hidden_states=decoder_hidden_states,
+ past_key_values=model_kwargs.get("past_key_values"),
+ )
+ else:
+ return generate_results[:, :, : (max_length - cur_len)]
+
+ def _update_model_kwargs_for_generation(
+ self,
+ outputs: ModelOutput,
+ model_kwargs: Dict[str, Any],
+ horizon_length: int = 1,
+ is_encoder_decoder: bool = False,
+ standardize_cache_format: bool = False,
+ ) -> Dict[str, Any]:
+ # update past_key_values
+ if "past_key_values" in outputs:
+ model_kwargs["past_key_values"] = outputs.past_key_values
+ elif "mems" in outputs:
+ model_kwargs["past_key_values"] = outputs.mems
+ elif "past_buckets_states" in outputs:
+ model_kwargs["past_key_values"] = outputs.past_buckets_states
+
+ if getattr(outputs, "state", None) is not None:
+ model_kwargs["state"] = outputs.state
+
+ # update token_type_ids with last value
+ if "token_type_ids" in model_kwargs:
+ token_type_ids = model_kwargs["token_type_ids"]
+ model_kwargs["token_type_ids"] = torch.cat(
+ [token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1
+ )
+
+ if not is_encoder_decoder:
+ # update attention mask
+ if "attention_mask" in model_kwargs:
+ attention_mask = model_kwargs["attention_mask"]
+ model_kwargs["attention_mask"] = torch.cat(
+ [
+ attention_mask,
+ attention_mask.new_ones(
+ (attention_mask.shape[0], horizon_length)
+ ),
+ ],
+ dim=-1,
+ )
+ else:
+ # update decoder attention mask
+ if "decoder_attention_mask" in model_kwargs:
+ decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+ model_kwargs["decoder_attention_mask"] = torch.cat(
+ [
+ decoder_attention_mask,
+ decoder_attention_mask.new_ones(
+ (decoder_attention_mask.shape[0], horizon_length)
+ ),
+ ],
+ dim=-1,
+ )
+
+ if (
+ "cache_position" in model_kwargs
+ and model_kwargs["cache_position"] is not None
+ ):
+ model_kwargs["cache_position"] = (
+ model_kwargs["cache_position"][-1:] + horizon_length
+ )
+
+ return model_kwargs
diff --git a/iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/__init__.py
similarity index 100%
rename from iotdb-core/ainode/iotdb/ainode/core/inference/strategy/__init__.py
rename to iotdb-core/ainode/iotdb/ainode/core/model/timerxl/__init__.py
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/configuration_timer.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/configuration_timer.py
new file mode 100644
index 00000000000..34f9de91b63
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/configuration_timer.py
@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List
+
+from transformers import PretrainedConfig
+
+
+class TimerConfig(PretrainedConfig):
+ model_type = "timer"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ input_token_len: int = 1,
+ hidden_size: int = 1024,
+ intermediate_size: int = 2048,
+ output_token_lens: List[int] = [1, 8, 32, 64],
+ num_hidden_layers: int = 8,
+ num_attention_heads: int = 8,
+ hidden_act: str = "silu",
+ use_cache: bool = True,
+ rope_theta: int = 10000,
+ attention_dropout: float = 0.0,
+ initializer_range: float = 0.02,
+ max_position_embeddings: int = 10000,
+ **kwargs,
+ ):
+ self.input_token_len = input_token_len
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.hidden_act = hidden_act
+ self.output_token_lens = output_token_lens
+ self.use_cache = use_cache
+ self.rope_theta = rope_theta
+ self.attention_dropout = attention_dropout
+ self.initializer_range = initializer_range
+ self.max_position_embeddings = max_position_embeddings
+
+ super().__init__(
+ **kwargs,
+ )
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py
new file mode 100644
index 00000000000..37bf56dfc59
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/modeling_timer.py
@@ -0,0 +1,640 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import Cache, DynamicCache, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
+from transformers.modeling_outputs import (
+ MoeCausalLMOutputWithPast,
+ MoeModelOutputWithPast,
+)
+
+from iotdb.ainode.core.log import Logger
+from iotdb.ainode.core.model.timerxl.configuration_timer import TimerConfig
+from iotdb.ainode.core.model.timerxl.ts_generation_mixin import TSGenerationMixin
+
+logger = Logger()
+
+
+def rotate_half(x):
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+class TimerPatchEmbedding(nn.Module):
+ def __init__(self, config: TimerConfig):
+ super().__init__()
+ self.input_token_len = config.input_token_len
+ self.emb = nn.Linear(config.input_token_len, config.hidden_size, bias=False)
+
+ def forward(self, hidden_state: torch.Tensor):
+ hidden_state = hidden_state.unfold(
+ dimension=-1, size=self.input_token_len, step=self.input_token_len
+ )
+ return self.emb(hidden_state)
+
+
+class TimerPointEmbedding(nn.Module):
+ def __init__(self, config: TimerConfig):
+ super().__init__()
+ self.emb_layer = nn.Linear(
+ config.input_token_len, config.hidden_size, bias=False
+ )
+ self.gate_layer = nn.Linear(
+ config.input_token_len, config.hidden_size, bias=False
+ )
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ emb = self.act_fn(self.gate_layer(x)) * self.emb_layer(x)
+ return emb
+
+
+class TimeMoeRotaryEmbedding(torch.nn.Module):
+ def __init__(self, dim, max_position_embeddings=10000, base=10000, device=None):
+ super().__init__()
+ self.dim = dim
+ self.max_position_embeddings = max_position_embeddings
+ self.base = base
+ inv_freq = 1.0 / (
+ self.base
+ ** (
+ torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
+ / self.dim
+ )
+ )
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ # Build here to make `torch.jit.trace` work.
+ self._set_cos_sin_cache(
+ seq_len=max_position_embeddings,
+ device=self.inv_freq.device,
+ dtype=torch.get_default_dtype(),
+ )
+
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(
+ self.max_seq_len_cached, device=device, dtype=torch.int64
+ ).type_as(self.inv_freq)
+
+ freqs = torch.outer(t, self.inv_freq)
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+ def forward(self, x, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ if seq_len > self.max_seq_len_cached:
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+ return (
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
+ )
+
+
+class TimerAttention(nn.Module):
+ def __init__(self, config: TimerConfig, layer_idx: Optional[int] = None):
+ super().__init__()
+ self.layer_idx = layer_idx
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.hidden_size // self.num_heads
+ self.attention_dropout = config.attention_dropout
+ self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+ self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+ self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+ self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+ self.rotary_emb = TimeMoeRotaryEmbedding(
+ self.head_dim, max_position_embeddings=config.max_position_embeddings
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: bool = False,
+ **kwargs,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
+ bsz, q_len, _ = hidden_states.size()
+
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+
+ query_states = query_states.view(
+ bsz, q_len, self.num_heads, self.head_dim
+ ).transpose(1, 2)
+ key_states = key_states.view(
+ bsz, q_len, self.num_heads, self.head_dim
+ ).transpose(1, 2)
+ value_states = value_states.view(
+ bsz, q_len, self.num_heads, self.head_dim
+ ).transpose(1, 2)
+
+ kv_seq_len = key_states.shape[-2]
+ if past_key_value is not None:
+ kv_seq_len += past_key_value.get_seq_length(self.layer_idx)
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+ query_states, key_states = apply_rotary_pos_emb(
+ query_states, key_states, cos, sin, position_ids
+ )
+
+ if past_key_value is not None:
+ key_states, value_states = past_key_value.update(
+ key_states, value_states, self.layer_idx
+ )
+
+ attn_output = F.scaled_dot_product_attention(
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout_p=self.attention_dropout,
+ )
+
+ attn_output = attn_output.transpose(1, 2).contiguous()
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+ attn_output = self.o_proj(attn_output)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights, past_key_value
+
+
+class TimerMLP(nn.Module):
+ def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+ self.act_fn = ACT2FN[hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(
+ self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)
+ )
+
+
+class TimerDecoderLayer(nn.Module):
+ def __init__(self, config: TimerConfig, layer_idx: int):
+ super().__init__()
+ self.self_attn = TimerAttention(config, layer_idx)
+
+ self.ffn_layer = TimerMLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=config.intermediate_size,
+ hidden_act=config.hidden_act,
+ )
+ self.norm1 = torch.nn.LayerNorm(config.hidden_size)
+ self.norm2 = torch.nn.LayerNorm(config.hidden_size)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_value: Optional[Cache] = None,
+ output_attentions: Optional[bool] = False,
+ **kwargs,
+ ) -> Tuple[
+ torch.FloatTensor,
+ Optional[torch.Tensor],
+ Optional[Cache],
+ ]:
+ residual = hidden_states
+
+ # Self Attention
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_value,
+ output_attentions=output_attentions,
+ )
+ hidden_states = residual + hidden_states
+ hidden_states = self.norm1(hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.ffn_layer(hidden_states)
+ hidden_states = residual + hidden_states
+ hidden_states = self.norm2(hidden_states)
+
+ if not output_attentions:
+ self_attn_weights = None
+
+ return hidden_states, self_attn_weights, present_key_value
+
+
+class TimerPreTrainedModel(PreTrainedModel):
+ config_class = TimerConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["TimeDecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn_2 = True
+ _supports_sdpa = False
+ _supports_cache_class = True
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, torch.nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, torch.nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+
+class TimerModel(TimerPreTrainedModel):
+ def __init__(self, config: TimerConfig):
+ super().__init__(config)
+ self.embed_layer = TimerPatchEmbedding(config)
+ self.layers = nn.ModuleList(
+ [
+ TimerDecoderLayer(config, layer_idx)
+ for layer_idx in range(config.num_hidden_layers)
+ ]
+ )
+ self.norm = torch.nn.LayerNorm(config.hidden_size)
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ input_ids: torch.FloatTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[
+ Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor]]]
+ ] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
+ # input_ids is the input of time series, its shape is [batch_size, seq_len]
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
+ )
+ elif input_ids is not None:
+ batch_size, seq_length = input_ids.shape
+ elif inputs_embeds is not None:
+ batch_size, seq_length, _ = inputs_embeds.shape
+ else:
+ raise ValueError(
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_layer(input_ids)
+ seq_length = inputs_embeds.shape[1]
+
+ past_key_values_length = 0
+ use_legacy_cache = False
+
+ if past_key_values is not None:
+ use_legacy_cache = not isinstance(past_key_values, Cache)
+ # Converts the legacy cache which is tuple into an equivalent Cache. Used for backward compatibility.
+ if use_legacy_cache:
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+ past_key_values_length = past_key_values.get_seq_length()
+
+ # When training + checkpoints, caching is usually disabled (just do not transfer)
+ if (
+ self.gradient_checkpointing
+ and self.training
+ and isinstance(past_key_values, Cache)
+ ):
+ past_key_values = None
+ past_key_values_length = 0
+
+ if position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length,
+ seq_length + past_key_values_length,
+ dtype=torch.long,
+ device=device,
+ )
+ # position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+ position_ids = position_ids.view(-1, seq_length)
+ else:
+ position_ids = position_ids.view(-1, seq_length).long()
+
+ # 4d mask is passed through the layers
+ attention_mask = _prepare_4d_causal_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ inputs_embeds,
+ past_key_values_length,
+ sliding_window=None,
+ )
+
+ hidden_states = inputs_embeds
+
+ # decoder layers
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+ next_decoder_cache = None
+
+ for decoder_layer in self.layers:
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ hidden_states,
+ attention_mask,
+ position_ids,
+ past_key_values,
+ output_attentions,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_value=past_key_values,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ if output_attentions:
+ all_self_attns += (layer_outputs[1],)
+
+ if isinstance(past_key_values, Cache):
+ next_decoder_cache = layer_outputs[2]
+
+ hidden_states = self.norm(hidden_states)
+ # add hidden states from the last decoder layer
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ next_cache = None
+ if isinstance(past_key_values, Cache):
+ next_cache = (
+ next_decoder_cache.to_legacy_cache()
+ if use_legacy_cache
+ else next_decoder_cache
+ )
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+ if v is not None
+ )
+ return MoeModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=next_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attns,
+ )
+
+
+class TimerForPrediction(TimerPreTrainedModel, TSGenerationMixin):
+ def __init__(self, config: TimerConfig):
+ super().__init__(config)
+ self.config = config
+ self.model = TimerModel(self.config)
+ lm_head_list = []
+ self.output_token_len_map = {}
+ for i, output_token_len in enumerate(self.config.output_token_lens):
+ lm_head_list.append(
+ nn.Linear(self.config.hidden_size, output_token_len, bias=False)
+ )
+ self.output_token_len_map[output_token_len] = i
+ self.lm_heads = nn.ModuleList(lm_head_list)
+ self.loss_function = torch.nn.MSELoss(reduction="none")
+ self.post_init()
+
+ def set_decoder(self, decoder):
+ self.model = decoder
+
+ def get_decoder(self):
+ return self.model
+
+ def forward(
+ self,
+ input_ids: torch.FloatTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[
+ Union[Cache, tuple[tuple[torch.Tensor, torch.Tensor]]]
+ ] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ labels: Optional[torch.FloatTensor] = None,
+ loss_masks: Optional[torch.FloatTensor] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ max_output_length: Optional[int] = None,
+ revin: Optional[bool] = False,
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+
+ if revin:
+ mean, std = input_ids.mean(dim=-1, keepdim=True), input_ids.std(
+ dim=-1, keepdim=True
+ )
+ input_ids = (input_ids - mean) / std
+ outputs = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+ predictions = None
+
+ loss = None
+ if labels is not None:
+ ar_loss = 0.0
+ for lm_head, output_token_len in zip(
+ self.lm_heads, self.config.output_token_lens
+ ):
+ one_predictions = lm_head(hidden_states)
+ one_loss = self.calc_ar_loss(
+ one_predictions, labels, loss_masks, output_token_len
+ )
+ ar_loss += one_loss
+ if predictions is None:
+ predictions = one_predictions
+ loss = ar_loss / len(self.config.output_token_lens)
+ else:
+ if max_output_length is None:
+ output_token_len = self.config.output_token_lens[0]
+ max_output_length = output_token_len
+ else:
+ output_token_len = self.config.output_token_lens[0]
+ for h in self.config.output_token_lens[1:]:
+ if h > max_output_length:
+ break
+ else:
+ output_token_len = h
+ lm_head = self.lm_heads[self.output_token_len_map[output_token_len]]
+ predictions = lm_head(hidden_states)[:, -1, :]
+ if output_token_len > max_output_length:
+ predictions = predictions[:, :max_output_length]
+ if revin:
+ predictions = predictions * std + mean
+ if not return_dict:
+ output = (predictions,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return MoeCausalLMOutputWithPast(
+ loss=loss,
+ logits=predictions,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def calc_ar_loss(self, predictions, labels, loss_masks, output_token_len):
+ seq_len = predictions.shape[1] * self.config.input_token_len
+ labels = labels[:, : seq_len - self.config.input_token_len + output_token_len]
+ shift_labels = labels.unfold(
+ dimension=-1, size=output_token_len, step=self.config.input_token_len
+ )
+
+ # Calculate loss with mask
+ losses = self.loss_function(predictions, shift_labels).mean(dim=-1)
+ if loss_masks is not None:
+ losses = losses * loss_masks
+ loss = losses.sum() / loss_masks.sum()
+ else:
+ loss = torch.mean(losses)
+
+ return loss
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ revin=True,
+ **kwargs,
+ ):
+ # Omit tokens covered by past_key_values
+ if past_key_values is not None:
+ if isinstance(past_key_values, Cache):
+ past_length = past_key_values.get_seq_length()
+ else:
+ past_length = past_key_values[0][0].shape[2]
+
+ # Keep only the unprocessed tokens:
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+ # input)
+ if attention_mask is not None and attention_mask.shape[1] > (
+ input_ids.shape[1] // self.config.input_token_len
+ ):
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+ # input_ids based on the past_length.
+ elif past_length < (input_ids.shape[1] // self.config.input_token_len):
+ input_ids = input_ids[:, past_length * self.config.input_token_len :]
+ # 3 - Otherwise (past_length >= (input_ids.shape[1] // self.config.input_token_len)), let's assume input_ids only has unprocessed tokens.
+
+ position_ids = kwargs.get("position_ids", None)
+ if attention_mask is not None and position_ids is None:
+ # create position_ids on the fly for batch generation
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ if past_key_values:
+ position_ids = position_ids[
+ :, -(input_ids.shape[1] // self.config.input_token_len) :
+ ]
+
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+ if inputs_embeds is not None and past_key_values is None:
+ model_inputs = {"inputs_embeds": inputs_embeds}
+ else:
+ model_inputs = {"input_ids": input_ids}
+
+ model_inputs.update(
+ {
+ "position_ids": position_ids,
+ "past_key_values": past_key_values,
+ "attention_mask": attention_mask,
+ "revin": revin,
+ }
+ )
+ return model_inputs
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/ts_generation_mixin.py b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/ts_generation_mixin.py
new file mode 100644
index 00000000000..18f711b8e1a
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/timerxl/ts_generation_mixin.py
@@ -0,0 +1,370 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import warnings
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from transformers import GenerationMixin, LogitsProcessorList, StoppingCriteriaList
+from transformers.generation import EosTokenCriteria, validate_stopping_criteria
+from transformers.generation.utils import (
+ GenerateDecoderOnlyOutput,
+ GenerateEncoderDecoderOutput,
+ GenerateNonBeamOutput,
+ GenerateOutput,
+ GenerationConfig,
+)
+from transformers.utils import ModelOutput
+
+
+class TSGenerationMixin(GenerationMixin):
+
+ @torch.no_grad()
+ def generate(
+ self,
+ inputs: Optional[torch.Tensor] = None,
+ generation_config: Optional[GenerationConfig] = None,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ prefix_allowed_tokens_fn: Optional[
+ Callable[[int, torch.Tensor], List[int]]
+ ] = None,
+ synced_gpus: Optional[bool] = None,
+ assistant_model: Optional["PreTrainedModel"] = None,
+ streamer: Optional["BaseStreamer"] = None,
+ negative_prompt_ids: Optional[torch.Tensor] = None,
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> Union[GenerateOutput, torch.LongTensor]:
+ if len(inputs.shape) == 2:
+ batch_size, cur_len = inputs.shape
+ if cur_len < self.config.input_token_len:
+ raise ValueError(
+ f"Input length must be at least {self.config.input_token_len}"
+ )
+ elif cur_len % self.config.input_token_len != 0:
+ new_len = (
+ cur_len // self.config.input_token_len
+ ) * self.config.input_token_len
+ inputs = inputs[:, -new_len:]
+ else:
+ raise ValueError("Input shape must be: [batch_size, seq_len]")
+ return super().generate(
+ inputs=inputs,
+ generation_config=generation_config,
+ logits_processor=logits_processor,
+ stopping_criteria=stopping_criteria,
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+ synced_gpus=synced_gpus,
+ assistant_model=assistant_model,
+ streamer=streamer,
+ negative_prompt_ids=negative_prompt_ids,
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
+ **kwargs,
+ )
+
+ def _sample(
+ self,
+ input_ids: torch.Tensor,
+ logits_processor: Optional[LogitsProcessorList] = None,
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
+ max_length: Optional[int] = None,
+ pad_token_id: Optional[int] = None,
+ eos_token_id: Optional[Union[int, List[int]]] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ output_scores: Optional[bool] = None,
+ output_logits: Optional[bool] = None,
+ return_dict_in_generate: Optional[bool] = None,
+ synced_gpus: bool = False,
+ streamer: Optional["BaseStreamer"] = None,
+ **model_kwargs,
+ ) -> Union[GenerateNonBeamOutput, torch.Tensor]:
+ input_ids = input_ids.to(self.device)
+ batch_size, cur_len = input_ids.shape
+ # init values
+ logits_processor = (
+ logits_processor if logits_processor is not None else LogitsProcessorList()
+ )
+ stopping_criteria = (
+ stopping_criteria
+ if stopping_criteria is not None
+ else StoppingCriteriaList()
+ )
+ if max_length is not None:
+ warnings.warn(
+ "`max_length` is deprecated in this function, use"
+ " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
+ UserWarning,
+ )
+ stopping_criteria = validate_stopping_criteria(
+ stopping_criteria, max_length
+ )
+ pad_token_id = (
+ pad_token_id
+ if pad_token_id is not None
+ else self.generation_config.pad_token_id
+ )
+ if eos_token_id is not None:
+ stopping_criteria.append(EosTokenCriteria(eos_token_id=eos_token_id))
+ else:
+ # remove when the method is totally private
+ # need to get `eos_token_id` and add stopping criteria, so that generation does not go forever
+ eos_token_id = [
+ criteria.eos_token_id.tolist()
+ for criteria in stopping_criteria
+ if hasattr(criteria, "eos_token_id")
+ ]
+ eos_token_id = eos_token_id[0] if eos_token_id else None
+ if eos_token_id is None and self.generation_config.eos_token_id is not None:
+ eos_token_id = self.generation_config.eos_token_id
+ stopping_criteria.append(EosTokenCriteria(eos_token_id=eos_token_id))
+
+ if isinstance(eos_token_id, int):
+ eos_token_id = [eos_token_id]
+ output_scores = (
+ output_scores
+ if output_scores is not None
+ else self.generation_config.output_scores
+ )
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.generation_config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.generation_config.output_hidden_states
+ )
+ return_dict_in_generate = (
+ return_dict_in_generate
+ if return_dict_in_generate is not None
+ else self.generation_config.return_dict_in_generate
+ )
+
+ # init attention / hidden states / scores tuples
+ raw_logits = () if (return_dict_in_generate and output_logits) else None
+ scores = () if (return_dict_in_generate and output_scores) else None
+ decoder_attentions = (
+ () if (return_dict_in_generate and output_attentions) else None
+ )
+ cross_attentions = (
+ () if (return_dict_in_generate and output_attentions) else None
+ )
+ decoder_hidden_states = (
+ () if (return_dict_in_generate and output_hidden_states) else None
+ )
+
+ # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+ if return_dict_in_generate and self.config.is_encoder_decoder:
+ encoder_attentions = (
+ model_kwargs["encoder_outputs"].get("attentions")
+ if output_attentions
+ else None
+ )
+ encoder_hidden_states = (
+ model_kwargs["encoder_outputs"].get("hidden_states")
+ if output_hidden_states
+ else None
+ )
+
+ # keep track of which sequences are already finished
+ if "inputs_embeds" in model_kwargs:
+ cur_len = model_kwargs["inputs_embeds"].shape[1]
+ this_peer_finished = False
+ unfinished_sequences = torch.ones(
+ batch_size, dtype=torch.long, device=input_ids.device
+ )
+ model_kwargs["cache_position"] = torch.arange(cur_len, device=input_ids.device)
+ true_seq_len = cur_len // self.config.input_token_len
+ model_kwargs["attention_mask"] = model_kwargs["attention_mask"][
+ :, -true_seq_len:
+ ]
+ max_length = stopping_criteria.max_length
+ while self._has_unfinished_sequences(
+ this_peer_finished, synced_gpus, device=input_ids.device
+ ):
+ # prepare model inputs
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+
+ input_length = input_ids.shape[1]
+
+ # forward pass to get next token
+ outputs = self(
+ **model_inputs,
+ return_dict=True,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ max_output_length=max_length - input_length,
+ )
+
+ if synced_gpus and this_peer_finished:
+ continue # don't waste resources running the code we don't need
+
+ next_token_logits = outputs.logits
+
+ # pre-process distribution
+ next_tokens_scores = logits_processor(input_ids, next_token_logits)
+
+ # Store scores, attentions and hidden_states when required
+ if return_dict_in_generate:
+ if output_scores:
+ scores += (next_tokens_scores,)
+ if output_logits:
+ raw_logits += (next_token_logits,)
+ if output_attentions:
+ decoder_attentions += (
+ (outputs.decoder_attentions,)
+ if self.config.is_encoder_decoder
+ else (outputs.attentions,)
+ )
+ if self.config.is_encoder_decoder:
+ cross_attentions += (outputs.cross_attentions,)
+
+ if output_hidden_states:
+ decoder_hidden_states += (
+ (outputs.decoder_hidden_states,)
+ if self.config.is_encoder_decoder
+ else (outputs.hidden_states,)
+ )
+
+ # argmax
+ # next_tokens = torch.argmax(next_tokens_scores, dim=-1)
+ next_tokens = next_tokens_scores
+
+ # finished sentences should have their next token be a padding token
+ if eos_token_id is not None:
+ if pad_token_id is None:
+ raise ValueError(
+ "If `eos_token_id` is defined, make sure that `pad_token_id` is defined."
+ )
+ next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
+ 1 - unfinished_sequences
+ )
+
+ # update generated ids, model inputs, and length for next step
+ horizon_length = next_tokens.shape[1] // self.config.input_token_len
+
+ input_ids = torch.cat([input_ids, next_tokens], dim=-1)
+ if streamer is not None:
+ streamer.put(next_tokens.cpu())
+ model_kwargs = self._update_model_kwargs_for_generation(
+ outputs,
+ model_kwargs,
+ horizon_length=horizon_length,
+ is_encoder_decoder=self.config.is_encoder_decoder,
+ )
+ unfinished_sequences = unfinished_sequences & ~stopping_criteria(
+ input_ids, scores
+ )
+ this_peer_finished = unfinished_sequences.max() == 0
+
+ if input_ids.shape[1] > max_length:
+ input_ids = input_ids[:, :max_length]
+
+ if streamer is not None:
+ streamer.end()
+
+ if return_dict_in_generate:
+ if self.config.is_encoder_decoder:
+ return GenerateEncoderDecoderOutput(
+ sequences=input_ids,
+ scores=scores,
+ logits=raw_logits,
+ encoder_attentions=encoder_attentions,
+ encoder_hidden_states=encoder_hidden_states,
+ decoder_attentions=decoder_attentions,
+ cross_attentions=cross_attentions,
+ decoder_hidden_states=decoder_hidden_states,
+ past_key_values=model_kwargs.get("past_key_values"),
+ )
+ else:
+ return GenerateDecoderOnlyOutput(
+ sequences=input_ids,
+ scores=scores,
+ logits=raw_logits,
+ attentions=decoder_attentions,
+ hidden_states=decoder_hidden_states,
+ past_key_values=model_kwargs.get("past_key_values"),
+ )
+ else:
+ return input_ids[:, -(max_length - cur_len) :]
+
+ def _update_model_kwargs_for_generation(
+ self,
+ outputs: ModelOutput,
+ model_kwargs: Dict[str, Any],
+ horizon_length: int = 1,
+ is_encoder_decoder: bool = False,
+ standardize_cache_format: bool = False,
+ ) -> Dict[str, Any]:
+ # update past_key_values
+ if "past_key_values" in outputs:
+ model_kwargs["past_key_values"] = outputs.past_key_values
+ elif "mems" in outputs:
+ model_kwargs["past_key_values"] = outputs.mems
+ elif "past_buckets_states" in outputs:
+ model_kwargs["past_key_values"] = outputs.past_buckets_states
+
+ if getattr(outputs, "state", None) is not None:
+ model_kwargs["state"] = outputs.state
+
+ # update token_type_ids with last value
+ if "token_type_ids" in model_kwargs:
+ token_type_ids = model_kwargs["token_type_ids"]
+ model_kwargs["token_type_ids"] = torch.cat(
+ [token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1
+ )
+
+ if not is_encoder_decoder:
+ # update attention mask
+ if "attention_mask" in model_kwargs:
+ attention_mask = model_kwargs["attention_mask"]
+ model_kwargs["attention_mask"] = torch.cat(
+ [
+ attention_mask,
+ attention_mask.new_ones(
+ (attention_mask.shape[0], horizon_length)
+ ),
+ ],
+ dim=-1,
+ )
+ else:
+ # update decoder attention mask
+ if "decoder_attention_mask" in model_kwargs:
+ decoder_attention_mask = model_kwargs["decoder_attention_mask"]
+ model_kwargs["decoder_attention_mask"] = torch.cat(
+ [
+ decoder_attention_mask,
+ decoder_attention_mask.new_ones(
+ (decoder_attention_mask.shape[0], horizon_length)
+ ),
+ ],
+ dim=-1,
+ )
+
+ if (
+ "cache_position" in model_kwargs
+ and model_kwargs["cache_position"] is not None
+ ):
+ model_kwargs["cache_position"] = (
+ model_kwargs["cache_position"][-1:] + horizon_length
+ )
+
+ return model_kwargs
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/utils.py b/iotdb-core/ainode/iotdb/ainode/core/model/utils.py
new file mode 100644
index 00000000000..9da8486d390
--- /dev/null
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/utils.py
@@ -0,0 +1,94 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+import importlib
+import json
+import sys
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Dict, Tuple
+
+from iotdb.ainode.core.model.model_enums import (
+ MODEL_CONFIG_FILE,
+ MODEL_WEIGHTS_FILE,
+ UriType,
+)
+
+
+def parse_uri_type(uri: str) -> UriType:
+ if uri.startswith("repo://"):
+ return UriType.REPO
+ elif uri.startswith("file://"):
+ return UriType.FILE
+ else:
+ raise ValueError(
+ f"Unsupported URI type: {uri}. Supported formats: repo:// or file://"
+ )
+
+
+def get_parsed_uri(uri: str) -> str:
+ return uri[7:] # Remove "repo://" or "file://" prefix
+
+
+@contextmanager
+def temporary_sys_path(path: str):
+ """Context manager for temporarily adding a path to sys.path"""
+ path_added = path not in sys.path
+ if path_added:
+ sys.path.insert(0, path)
+ try:
+ yield
+ finally:
+ if path_added and path in sys.path:
+ sys.path.remove(path)
+
+
+def load_model_config(config_path: Path) -> Dict:
+ with open(config_path, "r", encoding="utf-8") as f:
+ return json.load(f)
+
+
+def validate_model_files(model_dir: Path) -> Tuple[Path, Path]:
+ """Validate model files exist, return config and weights file paths"""
+ config_path = model_dir / MODEL_CONFIG_FILE
+ weights_path = model_dir / MODEL_WEIGHTS_FILE
+
+ if not config_path.exists():
+ raise ValueError(f"Model config file does not exist: {config_path}")
+ if not weights_path.exists():
+ raise ValueError(f"Model weights file does not exist: {weights_path}")
+
+ # Create __init__.py file to ensure model directory can be imported as a module
+ init_file = model_dir / "__init__.py"
+ if not init_file.exists():
+ init_file.touch()
+
+ return config_path, weights_path
+
+
+def import_class_from_path(module_name, class_path: str):
+ file_name, class_name = class_path.rsplit(".", 1)
+ module = importlib.import_module(module_name + "." + file_name)
+ return getattr(module, class_name)
+
+
+def ensure_init_file(path: Path):
+ """Ensure __init__.py file exists in the given path"""
+ init_file = path / "__init__.py"
+ if not init_file.exists():
+ init_file.touch()
diff --git a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py
index f01e1594f06..de5de968a7e 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py
@@ -20,7 +20,7 @@ from iotdb.ainode.core.constant import TSStatusCode
from iotdb.ainode.core.log import Logger
from iotdb.ainode.core.manager.cluster_manager import ClusterManager
from iotdb.ainode.core.manager.inference_manager import InferenceManager
-from iotdb.ainode.core.manager.model_manager import ModelManager
+from iotdb.ainode.core.manager.model_manager import get_model_manager
from iotdb.ainode.core.rpc.status import get_status
from iotdb.ainode.core.util.gpu_mapping import get_available_devices
from iotdb.thrift.ainode import IAINodeRPCService
@@ -47,24 +47,10 @@ from iotdb.thrift.common.ttypes import TSStatus
logger = Logger()
-def _ensure_device_id_is_available(device_id_list: list[str]) -> TSStatus:
- """
- Ensure that the device IDs in the provided list are available.
- """
- available_devices = get_available_devices()
- for device_id in device_id_list:
- if device_id not in available_devices:
- return TSStatus(
- code=TSStatusCode.INVALID_URI_ERROR.value,
- message=f"Device ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.",
- )
- return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value)
-
-
class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
def __init__(self, ainode):
self._ainode = ainode
- self._model_manager = ModelManager()
+ self._model_manager = get_model_manager()
self._inference_manager = InferenceManager()
def stop(self) -> None:
@@ -78,43 +64,55 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
def registerModel(self, req: TRegisterModelReq) -> TRegisterModelResp:
return self._model_manager.register_model(req)
- def loadModel(self, req: TLoadModelReq) -> TSStatus:
- status = self._ensure_model_is_built_in_or_fine_tuned(req.existingModelId)
- if status.code != TSStatusCode.SUCCESS_STATUS.value:
- return status
- status = _ensure_device_id_is_available(req.deviceIdList)
- if status.code != TSStatusCode.SUCCESS_STATUS.value:
- return status
- return self._inference_manager.load_model(req)
-
- def unloadModel(self, req: TUnloadModelReq) -> TSStatus:
- status = self._ensure_model_is_built_in_or_fine_tuned(req.modelId)
- if status.code != TSStatusCode.SUCCESS_STATUS.value:
- return status
- status = _ensure_device_id_is_available(req.deviceIdList)
- if status.code != TSStatusCode.SUCCESS_STATUS.value:
- return status
- return self._inference_manager.unload_model(req)
-
def deleteModel(self, req: TDeleteModelReq) -> TSStatus:
return self._model_manager.delete_model(req)
- def inference(self, req: TInferenceReq) -> TInferenceResp:
- return self._inference_manager.inference(req)
+ def showModels(self, req: TShowModelsReq) -> TShowModelsResp:
+ return self._model_manager.show_models(req)
- def forecast(self, req: TForecastReq) -> TSStatus:
- return self._inference_manager.forecast(req)
+ def loadModel(self, req: TLoadModelReq) -> TSStatus:
+ if not self._model_manager.is_model_registered(req.existingModelId):
+ return TSStatus(
+ code=TSStatusCode.MODEL_NOT_FOUND_ERROR.value,
+ message=f"Model [{req.existingModelId}] is not supported. You can use 'SHOW MODELS' to retrieve the available models.",
+ )
- def getAIHeartbeat(self, req: TAIHeartbeatReq) -> TAIHeartbeatResp:
- return ClusterManager.get_heart_beat(req)
+ available_devices = get_available_devices()
+ for device_id in req.deviceIdList:
+ if device_id not in available_devices:
+ return TSStatus(
+ code=TSStatusCode.INVALID_URI_ERROR.value,
+ message=f"Device ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.",
+ )
- def showModels(self, req: TShowModelsReq) -> TShowModelsResp:
- return self._model_manager.show_models(req)
+ return self._inference_manager.load_model(req)
+
+ def unloadModel(self, req: TUnloadModelReq) -> TSStatus:
+ if not self._model_manager.is_model_registered(req.modelId):
+ return TSStatus(
+ code=TSStatusCode.MODEL_NOT_FOUND_ERROR.value,
+ message=f"Model [{req.modelId}] is not supported. You can use 'SHOW MODELS' to retrieve the available models.",
+ )
+
+ available_devices = get_available_devices()
+ for device_id in req.deviceIdList:
+ if device_id not in available_devices:
+ return TSStatus(
+ code=TSStatusCode.INVALID_URI_ERROR.value,
+ message=f"Device ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.",
+ )
+
+ return self._inference_manager.unload_model(req)
def showLoadedModels(self, req: TShowLoadedModelsReq) -> TShowLoadedModelsResp:
- status = _ensure_device_id_is_available(req.deviceIdList)
- if status.code != TSStatusCode.SUCCESS_STATUS.value:
- return TShowLoadedModelsResp(status=status, deviceLoadedModelsMap={})
+ available_devices = get_available_devices()
+ for device_id in req.deviceIdList:
+ if device_id not in available_devices:
+ status = TSStatus(
+ code=TSStatusCode.INVALID_URI_ERROR.value,
+ message=f"Device ID [{device_id}] is not available. You can use 'SHOW AI_DEVICES' to retrieve the available devices.",
+ )
+ return TShowLoadedModelsResp(status=status, deviceLoadedModelsMap={})
return self._inference_manager.show_loaded_models(req)
def showAIDevices(self) -> TShowAIDevicesResp:
@@ -123,13 +121,14 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
deviceIdList=get_available_devices(),
)
+ def inference(self, req: TInferenceReq) -> TInferenceResp:
+ return self._inference_manager.inference(req)
+
+ def forecast(self, req: TForecastReq) -> TSStatus:
+ return self._inference_manager.forecast(req)
+
+ def getAIHeartbeat(self, req: TAIHeartbeatReq) -> TAIHeartbeatResp:
+ return ClusterManager.get_heart_beat(req)
+
def createTrainingTask(self, req: TTrainingReq) -> TSStatus:
pass
-
- def _ensure_model_is_built_in_or_fine_tuned(self, model_id: str) -> TSStatus:
- if not self._model_manager.is_built_in_or_fine_tuned(model_id):
- return TSStatus(
- code=TSStatusCode.MODEL_NOT_FOUND_ERROR.value,
- message=f"Model [{model_id}] is not a built-in or fine-tuned model. You can use 'SHOW MODELS' to retrieve the available models.",
- )
- return TSStatus(code=TSStatusCode.SUCCESS_STATUS.value)