[AINode] Integrate Sundial as built-in model

integration-test/src/test/java/org/apache/iotdb/relational/it/schema/IoTDBDatabaseIT.java

@@ -579,6 +579,7 @@ public void testInformationSchema() throws SQLException {
  Arrays.asList(
  "_STLForecaster,",
  "_NaiveForecaster,",
+ "_sundial,",
  "_HoltWinters,",
  "_TimerXL,",
  "_ExponentialSmoothing,",
@@ -702,6 +703,7 @@ public void testInformationSchema() throws SQLException {
  "model_id,",
  new HashSet<>(
  Arrays.asList(
+ "_sundial,",
  "_TimerXL,",
  "_STLForecaster,",
  "_NaiveForecaster,",

iotdb-core/ainode/ainode/core/constant.py

@@ -152,6 +152,9 @@ class BuiltInModelType(Enum):
  # timerxl
  TIMER_XL = "_timerxl"
 
+ # sundial
+ SUNDIAL = "_sundial"
+
  @classmethod
  def values(cls) -> List[str]:
  values = []
@@ -259,7 +262,13 @@ class AttributeName(Enum):
  ATTENTION_DROPOUT = "attention_dropout"
  INITIALIZER_RANGE = "initializer_range"
  MAX_POSITION_EMBEDDINGS = "max_position_embeddings"
- TIMERXL_CKPT_PATH = "ckpt_path"
+ 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

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

@@ -52,6 +52,17 @@ def infer(self, full_data, predict_length=96, **_):
  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 ('=', '|'):
+ data = data.byteswap().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, **_):
@@ -94,6 +105,8 @@ def infer(self, full_data, window_interval=None, window_step=None, **kwargs):
 def _get_strategy(model_id, model):
  if model_id == '_timerxl':
  return TimerXLStrategy(model)
+ if model_id == '_sundial':
+ return SundialStrategy(model)
  if model_id.startswith('_'):
  return BuiltInStrategy(model)
  return RegisteredStrategy(model)

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

@@ -30,12 +30,14 @@
 
 from ainode.TimerXL.models import timer_xl
 from ainode.TimerXL.models.configuration_timer import TimerxlConfig
+from ainode.core.model.sundial import modeling_sundial
 from ainode.core.config import AINodeDescriptor
 from ainode.core.constant import AttributeName, BuiltInModelType
 from ainode.core.exception import InferenceModelInternalError
 from ainode.core.exception import WrongAttributeTypeError, NumericalRangeException, StringRangeException, \
  ListRangeException, BuiltInModelNotSupportError
 from ainode.core.log import Logger
+from ainode.core.model.sundial.configuration_sundial import SundialConfig
 
 logger = Logger()
 
@@ -57,6 +59,8 @@ def get_model_attributes(model_id: str):
  attribute_map = stray_attribute_map
  elif model_id == BuiltInModelType.TIMER_XL.value:
  attribute_map = timerxl_attribute_map
+ elif model_id == BuiltInModelType.SUNDIAL.value:
+ attribute_map = sundial_attribute_map
  else:
  raise BuiltInModelNotSupportError(model_id)
  return attribute_map
@@ -99,6 +103,8 @@ def fetch_built_in_model(model_id, inference_attributes):
  model = STRAYModel(attributes)
  elif model_id == BuiltInModelType.TIMER_XL.value:
  model = timer_xl.Model(TimerxlConfig.from_dict(attributes))
+ elif model_id == BuiltInModelType.SUNDIAL.value:
+ model = modeling_sundial.SundialForPrediction(SundialConfig.from_dict(attributes))
  else:
  raise BuiltInModelNotSupportError(model_id)
 
@@ -321,6 +327,100 @@ def parse_attribute(input_attributes: Dict[str, str], attribute_map: Dict[str, A
  raise e
  return attributes
 
+sundial_attribute_map = {
+ AttributeName.INPUT_TOKEN_LEN.value: IntAttribute(
+ name=AttributeName.INPUT_TOKEN_LEN.value,
+ default_value=16,
+ default_low=1,
+ default_high=5000
+ ),
+ AttributeName.HIDDEN_SIZE.value: IntAttribute(
+ name=AttributeName.HIDDEN_SIZE.value,
+ default_value=768,
+ default_low=1,
+ default_high=5000
+ ),
+ AttributeName.INTERMEDIATE_SIZE.value: IntAttribute(
+ name=AttributeName.INTERMEDIATE_SIZE.value,
+ default_value=3072,
+ default_low=1,
+ default_high=5000
+ ),
+ AttributeName.OUTPUT_TOKEN_LENS.value: ListAttribute(
+ name=AttributeName.OUTPUT_TOKEN_LENS.value,
+ default_value=[720],
+ value_type=int
+ ),
+ AttributeName.NUM_HIDDEN_LAYERS.value: IntAttribute(
+ name=AttributeName.NUM_HIDDEN_LAYERS.value,
+ default_value=12,
+ default_low=1,
+ default_high=16
+ ),
+ AttributeName.NUM_ATTENTION_HEADS.value: IntAttribute(
+ name=AttributeName.NUM_ATTENTION_HEADS.value,
+ default_value=12,
+ default_low=1,
+ default_high=192
+ ),
+ AttributeName.HIDDEN_ACT.value: StringAttribute(
+ name=AttributeName.HIDDEN_ACT.value,
+ default_value="silu",
+ value_choices=["relu", "gelu", "silu", "tanh"],
+ ),
+ AttributeName.USE_CACHE.value: BooleanAttribute(
+ name=AttributeName.USE_CACHE.value,
+ default_value=True,
+ ),
+ AttributeName.ROPE_THETA.value: IntAttribute(
+ name=AttributeName.ROPE_THETA.value,
+ default_value=10000,
+ default_low=1000,
+ default_high=50000
+ ),
+ AttributeName.DROPOUT_RATE.value: FloatAttribute(
+ name=AttributeName.DROPOUT_RATE.value,
+ default_value=0.1,
+ default_low=0.0,
+ default_high=1.0
+ ),
+ AttributeName.INITIALIZER_RANGE.value: FloatAttribute(
+ name=AttributeName.INITIALIZER_RANGE.value,
+ default_value=0.02,
+ default_low=0.0,
+ default_high=1.0
+ ),
+ AttributeName.MAX_POSITION_EMBEDDINGS.value: IntAttribute(
+ name=AttributeName.MAX_POSITION_EMBEDDINGS.value,
+ default_value=10000,
+ default_low=1,
+ default_high=50000
+ ),
+ AttributeName.FLOW_LOSS_DEPTH.value: IntAttribute(
+ name=AttributeName.FLOW_LOSS_DEPTH.value,
+ default_value=3,
+ default_low=1,
+ default_high=50
+ ),
+ AttributeName.NUM_SAMPLING_STEPS.value: IntAttribute(
+ name=AttributeName.NUM_SAMPLING_STEPS.value,
+ default_value=50,
+ default_low=1,
+ default_high=5000
+ ),
+ AttributeName.DIFFUSION_BATCH_MUL.value: IntAttribute(
+ name=AttributeName.DIFFUSION_BATCH_MUL.value,
+ default_value=4,
+ default_low=1,
+ default_high=5000
+ ),
+ AttributeName.CKPT_PATH.value: StringAttribute(
+ name=AttributeName.CKPT_PATH.value,
+ default_value=os.path.join(os.getcwd(), AINodeDescriptor().get_config().get_ain_models_dir(), 'weights',
+ 'sundial'),
+ value_choices=['']
+ )
+}
 
 timerxl_attribute_map = {
  AttributeName.INPUT_TOKEN_LEN.value: IntAttribute(
@@ -391,8 +491,8 @@ def parse_attribute(input_attributes: Dict[str, str], attribute_map: Dict[str, A
  default_low=1,
  default_high=50000
  ),
- AttributeName.TIMERXL_CKPT_PATH.value: StringAttribute(
- name=AttributeName.TIMERXL_CKPT_PATH.value,
+ AttributeName.CKPT_PATH.value: StringAttribute(
+ name=AttributeName.CKPT_PATH.value,
  default_value=os.path.join(os.getcwd(), AINodeDescriptor().get_config().get_ain_models_dir(), 'weights',
  'timerxl', 'model.safetensors'),
  value_choices=['']

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

@@ -0,0 +1,17 @@
+# 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.
+#

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

@@ -0,0 +1,66 @@
+# 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,
+ ckpt_path: str = None, # weight path
+ **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
+ self.ckpt_path = ckpt_path
+
+ super().__init__(
+ **kwargs,
+ )