[AINode] Integrate moirai2 as builtin forecasting model

LICENSE

@@ -349,4 +349,14 @@ The chronos-forecasting is open source software licensed under the Apache Licens
 Project page: https://github.com/amazon-science/chronos-forecasting
 License: https://github.com/amazon-science/chronos-forecasting/blob/main/LICENSE
 
---------------------------------------------------------------------------------
+--------------------------------------------------------------------------------
+
+The following files include code modified from uni2ts project.
+
+./iotdb-core/ainode/iotdb/ainode/core/model/moirai2/*
+
+The uni2ts is open source software licensed under the Apache License 2.0
+Project page: https://github.com/SalesforceAIResearch/uni2ts
+License: https://github.com/SalesforceAIResearch/uni2ts/blob/main/LICENSE.txt
+
+--------------------------------------------------------------------------------

integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java

@@ -56,7 +56,9 @@ public class AINodeTestUtils {
               new AbstractMap.SimpleEntry<>(
                   "timer_xl", new FakeModelInfo("timer_xl", "timer", "builtin", "active")),
               new AbstractMap.SimpleEntry<>(
-                  "chronos2", new FakeModelInfo("chronos2", "t5", "builtin", "active")))
+                  "chronos2", new FakeModelInfo("chronos2", "t5", "builtin", "active")),
+              new AbstractMap.SimpleEntry<>(
+                  "moirai2", new FakeModelInfo("moirai2", "moirai", "builtin", "active")))
           .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 
   public static final Map<String, FakeModelInfo> BUILTIN_MODEL_MAP;

iotdb-core/ainode/iotdb/ainode/core/inference/pool_scheduler/basic_pool_scheduler.py

@@ -59,6 +59,15 @@ def _estimate_shared_pool_size_by_total_mem(
     # Seize memory usage for each model
     mem_usages: Dict[str, float] = {}
     for model_info in all_models:
+        if model_info.model_type not in MODEL_MEM_USAGE_MAP:
+            logger.error(
+                f"[Inference] Model type '{model_info.model_type}' not found in MODEL_MEM_USAGE_MAP. "
+                f"Available types: {list(MODEL_MEM_USAGE_MAP.keys())}"
+            )
+            raise KeyError(
+                f"Model type '{model_info.model_type}' not found in MODEL_MEM_USAGE_MAP. "
+                f"Please add memory usage configuration for '{model_info.model_type}' in constant.py"
+            )
         mem_usages[model_info.model_id] = (
             MODEL_MEM_USAGE_MAP[model_info.model_type] * INFERENCE_EXTRA_MEMORY_RATIO
         )

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

@@ -596,7 +596,7 @@ def predict(
             context_tensor = torch.cat([context_tensor, prediction], dim=-1)[
                 ..., -self.model_context_length :
             ]
-            (batch_size, n_quantiles, context_length) = context_tensor.shape
+            batch_size, n_quantiles, context_length = context_tensor.shape
 
             with torch.no_grad():
                 # Reshape (batch, n_quantiles, context_length) -> (batch * n_quantiles, context_length)

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

@@ -145,4 +145,17 @@ def __repr__(self):
             "AutoModelForCausalLM": "model.Chronos2Model",
         },
     ),
+    "moirai2": ModelInfo(
+        model_id="moirai2",
+        category=ModelCategory.BUILTIN,
+        state=ModelStates.INACTIVE,
+        model_type="moirai",
+        pipeline_cls="pipeline_moirai2.Moirai2Pipeline",
+        repo_id="Salesforce/moirai-2.0-R-small",
+        auto_map={
+            "AutoConfig": "configuration_moirai2.Moirai2Config",
+            "AutoModelForCausalLM": "modeling_moirai2.Moirai2ForPrediction",
+        },
+        transformers_registered=True,
+    ),
 }

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

@@ -61,7 +61,6 @@ def load_model(model_info: ModelInfo, **model_kwargs) -> Any:
 
 def load_model_from_transformers(model_info: ModelInfo, **model_kwargs):
     device_map = model_kwargs.get("device_map", "cpu")
-    trust_remote_code = model_kwargs.get("trust_remote_code", True)
     train_from_scratch = model_kwargs.get("train_from_scratch", False)
 
     model_path = os.path.join(
@@ -107,11 +106,9 @@ def load_model_from_transformers(model_info: ModelInfo, **model_kwargs):
             model_cls = AutoModelForCausalLM
 
     if train_from_scratch:
-        model = model_cls.from_config(config_cls, trust_remote_code=trust_remote_code)
+        model = model_cls.from_config(config_cls)
     else:
-        model = model_cls.from_pretrained(
-            model_path, trust_remote_code=trust_remote_code
-        )
+        model = model_cls.from_pretrained(model_path)
 
     return BACKEND.move_model(model, device_map)
 

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

@@ -0,0 +1,22 @@
+# 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.
+#
+# This file is part of the Apache IoTDB project.
+#
+# This file includes code modified from the uni2ts project (https://github.com/salesforce/uni2ts).
+# The original code is licensed under the Apache License 2.0.
+#

iotdb-core/ainode/iotdb/ainode/core/model/moirai2/common/__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/iotdb/ainode/core/model/moirai2/common/torch_util.py

@@ -0,0 +1,130 @@
+# 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 Optional
+
+import numpy as np
+import torch
+from jaxtyping import Bool, Float, Int
+
+numpy_to_torch_dtype_dict = {
+    bool: torch.bool,
+    np.uint8: torch.uint8,
+    np.int8: torch.int8,
+    np.int16: torch.int16,
+    np.int32: torch.int32,
+    np.int64: torch.int64,
+    np.float16: torch.float16,
+    np.float32: torch.float32,
+    np.float64: torch.float64,
+    np.complex64: torch.complex64,
+    np.complex128: torch.complex128,
+}
+
+
+def packed_attention_mask(
+    sample_id: Int[torch.Tensor, "*batch seq_len"],
+) -> Bool[torch.Tensor, "*batch seq_len seq_len"]:
+    sample_id = sample_id.unsqueeze(-1)
+    attention_mask = sample_id.eq(sample_id.mT)
+    return attention_mask
+
+
+def packed_causal_attention_mask(
+    sample_id: Int[torch.Tensor, "*batch seq_len"],
+    time_id: Int[torch.Tensor, "*batch seq_len"],
+) -> Bool[torch.Tensor, "*batch seq_len seq_len"]:
+    attention_mask = packed_attention_mask(sample_id)
+    expanded_id1 = time_id.unsqueeze(-2)
+    expanded_id2 = time_id.unsqueeze(-1)
+    compare_res = expanded_id1 <= expanded_id2
+    attention_mask = attention_mask * compare_res
+    return attention_mask
+
+
+def mask_fill(
+    tensor: Float[torch.Tensor, "*batch dim"],
+    mask: Bool[torch.Tensor, "*batch"],
+    value: Float[torch.Tensor, "dim"],
+) -> Float[torch.Tensor, "*batch dim"]:
+    mask = mask.unsqueeze(-1)
+    return tensor * ~mask + value * mask
+
+
+def safe_div(
+    numer: torch.Tensor,
+    denom: torch.Tensor,
+) -> torch.Tensor:
+    return numer / torch.where(
+        denom == 0,
+        1.0,
+        denom,
+    )
+
+
+def size_to_mask(
+    max_size: int,
+    sizes: Int[torch.Tensor, "*batch"],
+) -> Bool[torch.Tensor, "*batch max_size"]:
+    mask = torch.arange(max_size, device=sizes.device)
+    return torch.lt(mask, sizes.unsqueeze(-1))
+
+
+def fixed_size(
+    value: Float[torch.Tensor, "*batch max_size"],
+) -> Int[torch.Tensor, "*batch"]:
+    sizes = torch.ones_like(value[..., 0], dtype=torch.long) * value.shape[-1]
+    return sizes
+
+
+def sized_mean(
+    value: Float[torch.Tensor, "*batch max_size"],
+    sizes: Optional[Int[torch.Tensor, "*batch"]],
+    dim: Optional[int | tuple[int, ...]] = None,
+    keepdim: bool = False,
+    size_keepdim: bool = False,
+    correction: int = 0,
+) -> Float[torch.Tensor, "..."]:
+    value = value * size_to_mask(value.shape[-1], sizes)
+    div_val = safe_div(
+        value.sum(dim=-1).sum(dim, keepdim=keepdim),
+        torch.clamp(sizes.sum(dim, keepdim=keepdim) - correction, min=0),
+    )
+    if size_keepdim:
+        div_val = div_val.unsqueeze(-1)
+    return div_val
+
+
+def masked_mean(
+    value: Float[torch.Tensor, "..."],
+    mask: Bool[torch.Tensor, "..."],
+    dim: Optional[int | tuple[int, ...]] = None,
+    keepdim: bool = False,
+    correction: int = 0,
+) -> Float[torch.Tensor, "..."]:
+    return safe_div(
+        (value * mask).sum(dim=dim, keepdim=keepdim),
+        torch.clamp(mask.float().sum(dim, keepdim=keepdim) - correction, min=0),
+    )
+
+
+def unsqueeze_trailing_dims(x: torch.Tensor, shape: torch.Size) -> torch.Tensor:
+    if x.ndim > len(shape) or x.shape != shape[: x.ndim]:
+        raise ValueError
+    dim = (...,) + (None,) * (len(shape) - x.ndim)
+    return x[dim]

iotdb-core/ainode/iotdb/ainode/core/model/moirai2/configuration_moirai2.py

@@ -0,0 +1,61 @@
+# 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, Tuple
+
+from transformers import PretrainedConfig
+
+
+class Moirai2Config(PretrainedConfig):
+    model_type = "moirai2"
+
+    def __init__(
+        self,
+        d_model: int = 384,
+        d_ff: int = 1024,
+        num_layers: int = 6,
+        patch_size: int = 16,
+        max_seq_len: int = 512,
+        attn_dropout_p: float = 0.0,
+        dropout_p: float = 0.0,
+        scaling: bool = True,
+        num_predict_token: int = 4,
+        quantile_levels: Tuple[float, ...] = (
+            0.1,
+            0.2,
+            0.3,
+            0.4,
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+        ),
+        **kwargs,
+    ):
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.num_layers = num_layers
+        self.patch_size = patch_size
+        self.max_seq_len = max_seq_len
+        self.attn_dropout_p = attn_dropout_p
+        self.dropout_p = dropout_p
+        self.scaling = scaling
+        self.num_predict_token = num_predict_token
+        self.quantile_levels = quantile_levels
+        super().__init__(**kwargs)