[cherry-pick] update tool_helpers version and add unittest (#9114)

README.md

@@ -137,7 +137,7 @@ Unified Checkpoint 大模型存储格式在模型参数分布上支持动态扩
 ### pip 安装
 
 ```shell
-pip install --upgrade paddlenlp==3.0.0b0
+pip install --upgrade paddlenlp==3.0.0b1
 ```
 
 或者可通过以下命令安装最新 develop 分支代码：

legacy/model_zoo/ernie-1.0/data_tools/helpers.cpp

@@ -32,7 +32,7 @@ using namespace std;
 
 const int32_t LONG_SENTENCE_LEN = 512;
 
-void build_blending_indices(py::array_t<uint8_t>& dataset_index,
+void build_blending_indices(py::array_t<int16_t>& dataset_index,
                             py::array_t<int64_t>& dataset_sample_index,
                             const py::array_t<double>& weights,
                             const int32_t num_datasets,
@@ -73,7 +73,7 @@ void build_blending_indices(py::array_t<uint8_t>& dataset_index,
     }
 
     // Populate the indices.
-    dataset_index_ptr[sample_idx] = static_cast<uint8_t>(max_error_index);
+    dataset_index_ptr[sample_idx] = static_cast<int16_t>(max_error_index);
     dataset_sample_index_ptr[sample_idx] = current_samples[max_error_index];
 
     // Update the total samples.

paddlenlp/__init__.py

@@ -19,7 +19,7 @@ from datetime import datetime
 PADDLENLP_STABLE_VERSION = "PADDLENLP_STABLE_VERSION"
 
 
-__version__ = "3.0.0b0.post"
+__version__ = "3.0.0b1.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")
 else:

paddlenlp/data/blendable_dataset.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import hashlib
+import importlib.metadata
 import os
 import time
 
@@ -45,8 +46,18 @@ class BlendableDataset(paddle.io.Dataset):
         # Build indicies.
         def _build_indices():
             start_time = time.time()
-            assert num_datasets < 255
-            dataset_index = np.zeros(self.size, dtype=np.uint8)
+
+            tool_helpers_version = importlib.metadata.version("tool_helpers")
+            if tool_helpers_version > "0.1.1":
+                assert (
+                    num_datasets < 32767
+                ), f"Detect num_datasets({num_datasets})>=32767. Currently, num_datasets should be less than 32767."
+                dataset_index = np.zeros(self.size, dtype=np.int16)
+            else:
+                assert (
+                    num_datasets < 255
+                ), f"Detect num_datasets:({num_datasets})>=255. When 'tool_helpers<=0.1.1', num_datasets should be less than 255. To support num_datasets greater than 255, please upgrade `tool_helpers>=0.1.2`."
+                dataset_index = np.zeros(self.size, dtype=np.uint8)
             dataset_sample_index = np.zeros(self.size, dtype=np.int64)
 
             from tool_helpers import helpers

requirements.txt

@@ -20,7 +20,7 @@ uvicorn
 typer
 rich
 safetensors
-tool_helpers==0.1.1 ; platform_system == "Linux"
+tool_helpers>=0.1.1 ; platform_system == "Linux"
 aistudio-sdk>=0.1.3
 jinja2
 regex

setup.py

@@ -110,7 +110,7 @@ def show():
         f.write(content)
 
 
-__version__ = "3.0.0b0.post"
+__version__ = "3.0.0b1.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")
 else:

tests/data/test_blendable_dataset.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed 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.metadata
+import unittest
+
+import numpy as np
+from parameterized import parameterized
+
+
+def build_blending_indices_python(dataset_index, dataset_sample_index, weights, num_datasets, size, verbose):
+    """
+    Given multiple datasets and a weighting array, build samples such that it follows those weights.
+
+    Parameters:
+    - dataset_index: NumPy array to store the dataset index for each sample.
+    - dataset_sample_index: NumPy array to store the sample index within each dataset.
+    - weights: NumPy array of weights for each dataset.
+    - num_datasets: Integer, the number of datasets.
+    - size: Integer, the total number of samples to generate.
+    - verbose: Boolean, whether to print verbose output.
+    """
+    if verbose:
+        print("> building indices for blendable datasets ...")
+
+    # Initialize buffer for number of samples used for each dataset.
+    current_samples = np.zeros(num_datasets, dtype=np.int64)
+
+    # For each sample:
+    for sample_idx in range(size):
+        # Determine where the max error in sampling is happening.
+        sample_idx_double = max(sample_idx, 1)
+        max_error_index = 0
+        max_error = weights[0] * sample_idx_double - current_samples[0]
+        for dataset_idx in range(1, num_datasets):
+            error = weights[dataset_idx] * sample_idx_double - current_samples[dataset_idx]
+            if error > max_error:
+                max_error = error
+                max_error_index = dataset_idx
+
+        # Populate the indices.
+        dataset_index[sample_idx] = max_error_index
+        dataset_sample_index[sample_idx] = current_samples[max_error_index]
+
+        # Update the total samples.
+        current_samples[max_error_index] += 1
+
+    # Print info
+    if verbose:
+        print(" > sample ratios:")
+        for dataset_idx in range(num_datasets):
+            ratio = current_samples[dataset_idx] / size
+            print(f"   dataset {dataset_idx}, input: {weights[dataset_idx]}, achieved: {ratio}")
+
+
+def skip_if_version_not_equal(version="0.1.1", package_name="tool_helpers"):
+    try:
+        importlib.import_module(package_name)
+    except ImportError:
+        return True, f"package<{package_name}> not found, so to skip this test"
+    package_version = importlib.metadata.version(package_name)
+    if package_version != version:
+        return True, f"{package_name} version must be equal to {version}, but got {package_version}!"
+    return False, f"{package_name} version is ok!"
+
+
+class TestToolHelpers(unittest.TestCase):
+    def _test_build_blending_indices(
+        self, num_datasets=128, size=8192, dataset_index_dtype="uint8", verbose=False, seed=42, assert_true=True
+    ):
+        if isinstance(dataset_index_dtype, str):
+            dataset_index_dtype = np.dtype(dataset_index_dtype)
+        assert dataset_index_dtype in [np.uint8, np.int16], "dataset_index_dtype must be uint8 or int16!"
+
+        np.random.seed(seed)
+        random_numbers = np.random.rand(num_datasets)
+        random_numbers[0] = 200
+        weights = random_numbers / random_numbers.sum()
+        weights = weights.astype(np.float64)
+
+        # for ground truth, so we use np.int32
+        python_dataset_index = np.zeros(size, dtype=np.int32)
+        python_dataset_sample_index = np.zeros(size, dtype=np.int64)
+        build_blending_indices_python(
+            python_dataset_index, python_dataset_sample_index, weights, num_datasets, size, verbose
+        )
+
+        from tool_helpers import helpers
+
+        c_dataset_index = np.zeros(size, dtype=dataset_index_dtype)
+        c_dataset_sample_index = np.zeros(size, dtype=np.int64)
+        helpers.build_blending_indices(c_dataset_index, c_dataset_sample_index, weights, num_datasets, size, verbose)
+
+        assert_func = self.assertTrue if assert_true else self.assertFalse
+        assert_func(np.all(python_dataset_index == c_dataset_index.astype(python_dataset_index.dtype)))
+        self.assertTrue(
+            np.all(python_dataset_sample_index == c_dataset_sample_index.astype(python_dataset_sample_index.dtype))
+        )
+
+    @parameterized.expand(
+        [
+            (128, 8192, "uint8", False, 42, True),
+            (1024, 8192, "uint8", False, 42, False),
+            (128, 8192, "int16", False, 42, False),
+            (1024, 8192, "int16", False, 42, False),
+        ]
+    )
+    @unittest.skipIf(*skip_if_version_not_equal(version="0.1.1", package_name="tool_helpers"))
+    def test_build_blending_indices_version_0_1_1(
+        self, num_datasets=128, size=8192, dataset_index_dtype="uint8", verbose=False, seed=42, assert_true=True
+    ):
+        self._test_build_blending_indices(num_datasets, size, dataset_index_dtype, verbose, seed, assert_true)
+
+    @parameterized.expand(
+        [
+            (128, 8192, "uint8", False, 42, True),
+            (1024, 8192, "uint8", False, 42, False),
+            (128, 8192, "int16", False, 42, True),
+            (1024, 8192, "int16", False, 42, True),
+        ]
+    )
+    @unittest.skipIf(*skip_if_version_not_equal(version="0.1.2", package_name="tool_helpers"))
+    def test_build_blending_indices_version_0_1_2(
+        self, num_datasets=128, size=8192, dataset_index_dtype="uint8", verbose=False, seed=42, assert_true=True
+    ):
+        self._test_build_blending_indices(num_datasets, size, dataset_index_dtype, verbose, seed, assert_true)