[AINode] Revert transformer and tokenizer dependencies update (#16394)

CRZbulabula · JackieTien97 · commit a41553888f75 · 2025-12-16T11:38:48.000+08:00
(cherry picked from commit 4429187)
diff --git a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeInferenceSQLIT.java b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeInferenceSQLIT.java
@@ -56,26 +56,12 @@ public class AINodeInferenceSQLIT {
         "CREATE TIMESERIES root.AI.s1 WITH DATATYPE=DOUBLE, ENCODING=RLE",
         "CREATE TIMESERIES root.AI.s2 WITH DATATYPE=INT32, ENCODING=RLE",
         "CREATE TIMESERIES root.AI.s3 WITH DATATYPE=INT64, ENCODING=RLE",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(1,1.0,2.0,3,4)",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(2,2.0,3.0,4,5)",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(3,3.0,4.0,5,6)",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(4,4.0,5.0,6,7)",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(5,5.0,6.0,7,8)",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(6,6.0,7.0,8,9)",
-        "insert into root.AI(timestamp,s0,s1,s2,s3) values(7,7.0,8.0,9,10)",
       };
 
   static String[] WRITE_SQL_IN_TABLE =
       new String[] {
         "CREATE DATABASE root",
         "CREATE TABLE root.AI (s0 FLOAT FIELD, s1 DOUBLE FIELD, s2 INT32 FIELD, s3 INT64 FIELD)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(1,1.0,2.0,3,4)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(2,2.0,3.0,4,5)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(3,3.0,4.0,5,6)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(4,4.0,5.0,6,7)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(5,5.0,6.0,7,8)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(6,6.0,7.0,8,9)",
-        "insert into root.AI(time,s0,s1,s2,s3) values(7,7.0,8.0,9,10)",
       };
 
   @BeforeClass
@@ -84,6 +70,24 @@ public static void setUp() throws Exception {
     EnvFactory.getEnv().initClusterEnvironment(1, 1);
     prepareData(WRITE_SQL_IN_TREE);
     prepareTableData(WRITE_SQL_IN_TABLE);
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      for (int i = 0; i < 2880; i++) {
+        statement.execute(
+            String.format(
+                "INSERT INTO root.AI(timestamp,s0,s1,s2,s3) VALUES(%d,%f,%f,%d,%d)",
+                i, (float) i, (double) i, i, i));
+      }
+    }
+    try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement statement = connection.createStatement()) {
+      for (int i = 0; i < 2880; i++) {
+        statement.execute(
+            String.format(
+                "INSERT INTO root.AI(time,s0,s1,s2,s3) VALUES(%d,%f,%f,%d,%d)",
+                i, (float) i, (double) i, i, i));
+      }
+    }
   }
 
   @AfterClass
@@ -109,6 +113,29 @@ public void callInferenceTestInTable() throws SQLException {
   }
 
   public void callInferenceTest(Statement statement) throws SQLException {
+    // SQL0: Invoke timer-sundial and timer-xl to inference, the result should success
+    try (ResultSet resultSet =
+        statement.executeQuery(
+            "CALL INFERENCE(sundial, \"select s1 from root.AI\", generateTime=true, predict_length=720)")) {
+      ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+      checkHeader(resultSetMetaData, "Time,output0");
+      int count = 0;
+      while (resultSet.next()) {
+        count++;
+      }
+      assertEquals(720, count);
+    }
+    try (ResultSet resultSet =
+        statement.executeQuery(
+            "CALL INFERENCE(timer_xl, \"select s2 from root.AI\", generateTime=true, predict_length=256)")) {
+      ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+      checkHeader(resultSetMetaData, "Time,output0");
+      int count = 0;
+      while (resultSet.next()) {
+        count++;
+      }
+      assertEquals(256, count);
+    }
     // SQL1: user-defined model inferences multi-columns with generateTime=true
     String sql1 =
         "CALL INFERENCE(identity, \"select s0,s1,s2,s3 from root.AI\", generateTime=true)";
@@ -171,15 +198,15 @@ public void callInferenceTest(Statement statement) throws SQLException {
     //      assertEquals(3, count);
     //    }
 
-    try (ResultSet resultSet = statement.executeQuery(sql4)) {
-      ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
-      checkHeader(resultSetMetaData, "Time,output0");
-      int count = 0;
-      while (resultSet.next()) {
-        count++;
-      }
-      assertEquals(6, count);
-    }
+    //    try (ResultSet resultSet = statement.executeQuery(sql4)) {
+    //      ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+    //      checkHeader(resultSetMetaData, "Time,output0");
+    //      int count = 0;
+    //      while (resultSet.next()) {
+    //        count++;
+    //      }
+    //      assertEquals(6, count);
+    //    }
   }
 
   @Test
@@ -219,6 +246,29 @@ public void errorCallInferenceTest(Statement statement) {
   public void selectForecastTestInTable() throws SQLException {
     try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
         Statement statement = connection.createStatement()) {
+      // SQL0: Invoke timer-sundial and timer-xl to forecast, the result should success
+      try (ResultSet resultSet =
+          statement.executeQuery(
+              "SELECT * FROM FORECAST(model_id=>'sundial', input=>(SELECT time,s1 FROM root.AI) ORDER BY time, output_length=>720)")) {
+        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+        checkHeader(resultSetMetaData, "time,s1");
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(720, count);
+      }
+      try (ResultSet resultSet =
+          statement.executeQuery(
+              "SELECT * FROM FORECAST(model_id=>'timer_xl', input=>(SELECT time,s2 FROM root.AI) ORDER BY time, output_length=>256)")) {
+        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+        checkHeader(resultSetMetaData, "time,s2");
+        int count = 0;
+        while (resultSet.next()) {
+          count++;
+        }
+        assertEquals(256, count);
+      }
       // SQL1: user-defined model inferences multi-columns with generateTime=true
       String sql1 =
           "SELECT * FROM FORECAST(model_id=>'identity', input=>(SELECT time,s0,s1,s2,s3 FROM root.AI) ORDER BY time, output_length=>7)";
@@ -280,15 +330,15 @@ public void selectForecastTestInTable() throws SQLException {
       //        assertEquals(3, count);
       //      }
 
-      try (ResultSet resultSet = statement.executeQuery(sql4)) {
-        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
-        checkHeader(resultSetMetaData, "time,s0");
-        int count = 0;
-        while (resultSet.next()) {
-          count++;
-        }
-        assertEquals(6, count);
-      }
+      //      try (ResultSet resultSet = statement.executeQuery(sql4)) {
+      //        ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
+      //        checkHeader(resultSetMetaData, "time,s0");
+      //        int count = 0;
+      //        while (resultSet.next()) {
+      //          count++;
+      //        }
+      //        assertEquals(6, count);
+      //      }
     }
   }
 }
diff --git a/iotdb-core/ainode/ainode/core/manager/inference_manager.py b/iotdb-core/ainode/ainode/core/manager/inference_manager.py
@@ -84,7 +84,8 @@ class TimerXLStrategy(InferenceStrategy):
     def infer(self, full_data, predict_length=96, **_):
         data = full_data[1][0]
         if data.dtype.byteorder not in ("=", "|"):
-            data = data.byteswap().newbyteorder()
+            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)
@@ -96,7 +97,8 @@ 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()
+            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(
@@ -270,7 +272,7 @@ def _run(
             full_data = deserializer(raw)
             inference_attrs = extract_attrs(req)
 
-            predict_length = int(inference_attrs.get("predict_length", 96))
+            predict_length = int(inference_attrs.pop("predict_length", 96))
             if (
                 predict_length
                 > AINodeDescriptor().get_config().get_ain_inference_max_predict_length()
@@ -307,7 +309,8 @@ def _run(
                 # TODO: TSBlock -> Tensor codes should be unified
                 data = full_data[1][0]
                 if data.dtype.byteorder not in ("=", "|"):
-                    data = data.byteswap().newbyteorder()
+                    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")
                 if model_id == "sundial":
@@ -340,7 +343,9 @@ def _run(
                 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, **inference_attrs)
+                outputs = strategy.infer(
+                    full_data, predict_length=predict_length, **inference_attrs
+                )
 
             # construct response
             status = get_status(TSStatusCode.SUCCESS_STATUS)
diff --git a/iotdb-core/ainode/pyproject.toml b/iotdb-core/ainode/pyproject.toml
@@ -55,10 +55,39 @@ packages = [
 ]
 
 [tool.poetry.dependencies]
-python = ">=3.9, <3.13"
-numpy = "^1.21.4"
-pandas = ">=2.2.0"
-torch = ">=2.7.1"
+python = ">=3.11,<=3.13.5"
+
+# Core scientific stack
+numpy = [
+    { version = "^2.3.2", python = ">=3.10" },
+    { version = "^1.26.4", python = ">=3.9,<3.10" }
+]
+scipy = [
+    { version = "^1.12.0", python = ">=3.10" },
+    { version = "^1.11.4", python = ">=3.9,<3.10" }
+]
+pandas = "^2.3.2"
+scikit-learn = [
+    { version = "^1.7.1", python = ">=3.10" },
+    { version = "^1.5.2", python = ">=3.9,<3.10" }
+]
+statsmodels = "^0.14.5"
+sktime = "0.38.5"
+
+# ---- DL / HF stack ----
+torch = ">=2.7.0"
+torchmetrics = "^1.8.0"
+transformers = "==4.40.1"
+tokenizers = ">=0.19,<0.20"
+huggingface_hub = "^0.34.4"
+safetensors = "^0.6.2"
+einops = "^0.8.1"
+
+# ---- Optimizers / utils ----
+optuna = "^4.4.0"
+psutil = "^7.0.0"
+requests = "^2.32.5"
+dynaconf = "^3.2.11"
 thrift = ">=0.14.0"
 dynaconf = "^3.1.11"
 requests = "^2.31.0"
