add fineweb-edu dataset as calibration bakeup

auto_round/calib_dataset.py

@@ -18,6 +18,7 @@
 import multiprocessing
 import os
 import random
+import ssl
 import sys
 
 logging.getLogger("datasets").setLevel(logging.WARNING)
@@ -32,6 +33,51 @@
 
 CALIB_DATASETS = {}
 _GITHUB_CODE_CLEAN_MAX_DATASETS_VERSION = Version("3.6.0")
+_FINEWEB_EDU_HF_DATASET = "HuggingFaceFW/fineweb-edu"
+_FINEWEB_EDU_MODELSCOPE_DATASET = "AI-ModelScope/fineweb-edu"
+_FINEWEB_EDU_CONFIG = "sample-10BT"
+
+
+def _get_dataset_network_error(error):
+    """Return the network failure in an exception chain, if present."""
+    visited = set()
+    current = error
+    while current is not None and id(current) not in visited:
+        visited.add(id(current))
+        error_name = type(current).__name__.lower()
+        error_module = type(current).__module__.lower()
+        error_message = str(current).lower()
+        if isinstance(current, (ConnectionError, TimeoutError, ssl.SSLError)):
+            return current
+        if any(term in error_name for term in ("connection", "http", "proxy", "ssl", "timeout")) and any(
+            package in error_module
+            for package in ("datasets", "huggingface_hub", "httpcore", "httpx", "requests", "urllib3")
+        ):
+            return current
+        if any(
+            marker in error_message
+            for marker in ("connection timed out", "max retries exceeded", "proxy error", "ssl error")
+        ):
+            return current
+        current = current.__cause__ or current.__context__
+    return None
+
+
+def _warn_on_dataset_network_error(error, dataset_name):
+    """Give users an actionable cross-hub fallback for dataset network failures."""
+    network_error = _get_dataset_network_error(error)
+    if network_error is None:
+        return
+    logger.warning(
+        "Failed to load calibration dataset %r because of an HTTP/proxy network issue: %s. "
+        "Check your network and proxy settings, or switch to FineWeb-Edu with `--dataset fineweb-edu`. "
+        "By default it loads from Hugging Face (%s); set `AR_USE_MODELSCOPE=1` and install `modelscope` "
+        "to load its ModelScope mirror (%s).",
+        dataset_name,
+        network_error,
+        _FINEWEB_EDU_HF_DATASET,
+        _FINEWEB_EDU_MODELSCOPE_DATASET,
+    )
 
 
 def get_code_calibration_dataset(nsamples, datasets_version=None):
@@ -193,22 +239,7 @@ def get_pile_dataset(
     tokenizer_function = get_tokenizer_function(
         tokenizer, seqlen, apply_chat_template=apply_chat_template, system_prompt=system_prompt
     )
-    try:
-        calib_dataset = load_dataset("NeelNanda/pile-10k", split=split)
-    except Exception as e:
-        import ssl
-
-        error_message = str(e)
-        # Check for proxy or SSL error
-        if "proxy" in error_message.lower() or isinstance(e, ssl.SSLError) or "SSL" in error_message.upper():
-            logger.error(
-                f"Network error detected, please check proxy settings. "
-                f"Error: {error_message}. Or consider using a backup dataset by `pip install modelscope` "
-                f"and set '--dataset swift/pile-val-backup' in AutoRound API."
-            )
-        else:
-            logger.error(f"Failed to load the dataset: {error_message}")
-        sys.exit(1)
+    calib_dataset = load_dataset("NeelNanda/pile-10k", split=split)
     calib_dataset = calib_dataset.shuffle(seed=seed)
     calib_dataset = calib_dataset.map(
         tokenizer_function,
@@ -221,46 +252,65 @@ def get_pile_dataset(
     return calib_dataset
 
 
-@register_dataset(["swift/pile-val-backup", "pile-val-backup"])
-def get_pile_val_dataset(
+@register_dataset([_FINEWEB_EDU_HF_DATASET, _FINEWEB_EDU_MODELSCOPE_DATASET, "fineweb-edu"])
+def get_fineweb_edu_dataset(
     tokenizer,
     seqlen,
-    dataset_name="swift/pile-val-backup",
+    dataset_name="fineweb-edu",
     split=None,
     seed=42,
     apply_chat_template=False,
     system_prompt=None,
 ):
-    """Returns a dataloader for the specified dataset and split.
+    """Return a streaming FineWeb-Edu calibration dataset from either supported hub.
 
     Args:
     tokenizer: The tokenizer to be used for tokenization.
     seqlen: The maximum sequence length.
-    data_name: The name of the dataset.
-    split: The data split to be used (e.g., "train", "test", "validation").
+    dataset_name: FineWeb-Edu alias or explicit Hugging Face/ModelScope repository.
+    split: The data split to use. FineWeb-Edu currently provides ``train``.
     seed: The random seed for shuffling the dataset.
     apply_chat_template: Whether to apply chat template in tokenization.
 
     Returns:
-    A dataloader for the specified dataset and split, using the provided tokenizer and sequence length.
+    A tokenized streaming dataset using the provided tokenizer and sequence length.
     """
-
-    split = "validation"
+    split = "train" if split is None else split
+    if isinstance(split, list):
+        if len(split) != 1:
+            raise ValueError("FineWeb-Edu supports only one split at a time.")
+        split = split[0]
+    if split != "train":
+        raise ValueError("FineWeb-Edu supports only the train split.")
 
     tokenizer_function = get_tokenizer_function(
         tokenizer, seqlen, apply_chat_template=apply_chat_template, system_prompt=system_prompt
     )
-    from transformers.utils.versions import require_version
-
-    require_version(
-        "modelscope",
-        "Loading 'swift/pile-val-backup' dataset requires modelscope to be installed, " "`pip install modelscope`",
+    use_modelscope = dataset_name == _FINEWEB_EDU_MODELSCOPE_DATASET or (
+        dataset_name == "fineweb-edu" and envs.AR_USE_MODELSCOPE
     )
-    from modelscope import MsDataset  # pylint: disable=E0401
+    if use_modelscope:
+        from transformers.utils.versions import require_version
+
+        require_version(
+            "modelscope",
+            "Loading FineWeb-Edu from ModelScope requires `modelscope`; install it with `pip install modelscope`.",
+        )
+        from modelscope import MsDataset  # pylint: disable=E0401
 
-    calib_dataset = MsDataset.load(
-        "swift/pile-val-backup", "default", split=split
-    ).to_iterable_dataset()  # , use_streaming=True
+        calib_dataset = MsDataset.load(
+            _FINEWEB_EDU_MODELSCOPE_DATASET,
+            subset_name=_FINEWEB_EDU_CONFIG,
+            split=split,
+            use_streaming=True,
+        )
+    else:
+        calib_dataset = load_dataset(
+            _FINEWEB_EDU_HF_DATASET,
+            name=_FINEWEB_EDU_CONFIG,
+            split=split,
+            streaming=True,
+        )
     calib_dataset = calib_dataset.shuffle(seed=seed).take(10000)
     calib_dataset = calib_dataset.map(tokenizer_function, batched=True)
 
@@ -1109,7 +1159,11 @@ def get_dataset(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, n
     """
     # Allow disabling subprocess mode via environment variable
     if envs.AR_DISABLE_DATASET_SUBPROCESS:
-        return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples)
+        try:
+            return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples)
+        except Exception as error:
+            _warn_on_dataset_network_error(error, dataset_name)
+            raise
 
     # Run preprocessing in a subprocess so all temporary memory is freed on exit.
     # The HuggingFace datasets cache is warmed up as a side effect.
@@ -1139,7 +1193,11 @@ def get_dataset(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, n
 
     # (Re-)load the dataset in the main process.  When the subprocess
     # succeeded the HF datasets cache makes this almost instant.
-    return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples)
+    try:
+        return _get_dataset_impl(tokenizer, seqlen, dataset_name, seed, nsamples)
+    except Exception as error:
+        _warn_on_dataset_network_error(error, dataset_name)
+        raise
 
 
 def get_dataloader(tokenizer, seqlen, dataset_name="NeelNanda/pile-10k", seed=42, bs=8, nsamples=512):

docs/step_by_step.md

@@ -63,7 +63,9 @@ pip install auto-round
 
 The [NeelNanda/pile-10k](https://huggingface.co/datasets/NeelNanda/pile-10k) in huggingface is adopted as the default
 calibration data and will be downloaded automatically from the datasets Hub. Other available datasets include:
-- `swift/pile-val-backup` from modelscope for addressing HF network issue
+- [`fineweb-edu`](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) from Hugging Face. If Hugging Face is
+  inaccessible, install `modelscope`, set `AR_USE_MODELSCOPE=1`, and use the same alias to load the
+  [ModelScope mirror](https://modelscope.cn/datasets/AI-ModelScope/fineweb-edu)
 - `BAAI/CCI3-HQ` for Chinese
 - `codeparrot/github-code-clean` for code
 - `HuggingFaceH4/ultrachat_200k` for chat data

docs/step_by_step_CN.md

@@ -61,10 +61,11 @@ pip install auto-round
 ## 2 准备标定数据集
 
 ### 默认数据集
-**对于中国大陆用户推荐使用 ModelScope 中的 swift/pile-val-backup 以解决 Huggingface 不能访问的问题**
+**如果无法访问 Hugging Face，建议安装 `modelscope`、设置 `AR_USE_MODELSCOPE=1`，并使用 `fineweb-edu` 数据集。**
 
 默认标定数据集为 Hugging Face 上的 [NeelNanda/pile-10k](https://huggingface.co/datasets/NeelNanda/pile-10k) ，该数据集会自动从 Huggingface Hub 下载。同时也支持使用以下数据集：
-- ModelScope 中的 `swift/pile-val-backup`：用于解决 HF 访问问题
+- Hugging Face 上的 [`fineweb-edu`](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu)。如果无法访问 Hugging Face，
+  可安装 `modelscope`、设置 `AR_USE_MODELSCOPE=1`，并使用同一别名从 [ModelScope 镜像](https://modelscope.cn/datasets/AI-ModelScope/fineweb-edu) 加载
 - `BAAI/CCI3-HQ`：用于中文场景
 - `codeparrot/github-code-clean`：用于代码场景
 - `HuggingFaceH4/ultrachat_200k`：用于对话数据

test/unit/common/utils/test_calib_dataset_helpers.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 """Tests for the small pure helpers in ``auto_round/calib_dataset.py``."""
 
+import ssl
+import sys
+import types
 from unittest.mock import MagicMock
 
 import pytest
@@ -63,6 +66,72 @@ class _Returned:
             CALIB_DATASETS.pop("_return_check_", None)
 
 
+# ---------------------------------------------------------------------------
+# network failures and FineWeb-Edu routing
+# ---------------------------------------------------------------------------
+class TestDatasetNetworkErrors:
+    def test_network_warning_recommends_fineweb_edu(self, monkeypatch):
+        import auto_round.calib_dataset as calib_dataset
+
+        warning = MagicMock()
+        monkeypatch.setattr(calib_dataset.logger, "warning", warning)
+
+        calib_dataset._warn_on_dataset_network_error(ssl.SSLError("proxy unavailable"), "dataset")
+
+        warning.assert_called_once()
+        message = warning.call_args.args[0]
+        assert "--dataset fineweb-edu" in message
+        assert "AR_USE_MODELSCOPE=1" in message
+
+
+class TestFineWebEduDataset:
+    @staticmethod
+    def _streaming_dataset_mock():
+        dataset = MagicMock()
+        dataset.shuffle.return_value = dataset
+        dataset.take.return_value = dataset
+        dataset.map.return_value = dataset
+        return dataset
+
+    def test_alias_loads_huggingface_sample_without_modelscope(self, monkeypatch):
+        import auto_round.calib_dataset as calib_dataset
+
+        dataset = self._streaming_dataset_mock()
+        load_dataset = MagicMock(return_value=dataset)
+        monkeypatch.setattr(calib_dataset, "load_dataset", load_dataset)
+        monkeypatch.delenv("AR_USE_MODELSCOPE", raising=False)
+        monkeypatch.setitem(sys.modules, "modelscope", None)
+
+        result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128)
+
+        assert result is dataset
+        load_dataset.assert_called_once_with(
+            "HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True
+        )
+        dataset.shuffle.assert_called_once_with(seed=42)
+        dataset.take.assert_called_once_with(10000)
+
+    def test_alias_loads_modelscope_sample_when_enabled(self, monkeypatch):
+        import transformers.utils.versions as transformers_versions
+
+        import auto_round.calib_dataset as calib_dataset
+
+        dataset = self._streaming_dataset_mock()
+        load = MagicMock(return_value=dataset)
+        modelscope = types.ModuleType("modelscope")
+        modelscope.MsDataset = types.SimpleNamespace(load=load)
+        monkeypatch.setitem(sys.modules, "modelscope", modelscope)
+        monkeypatch.setenv("AR_USE_MODELSCOPE", "1")
+        monkeypatch.setattr(transformers_versions, "require_version", MagicMock())
+
+        result = calib_dataset.get_fineweb_edu_dataset(MagicMock(), 128)
+
+        assert result is dataset
+        load.assert_called_once_with(
+            "AI-ModelScope/fineweb-edu", subset_name="sample-10BT", split="train", use_streaming=True
+        )
+
+
 # ---------------------------------------------------------------------------
 # _make_map_fingerprint
 # ---------------------------------------------------------------------------