add compat with hf-hub >1 (#176)

* add compat with hf-hub >1

* fix import order to ensure tests go through

Author: markus583

.gitignore

@@ -1,4 +1,5 @@
-.venv
+.venv/
+.venv-*/
 build
 dist
 wandb
@@ -23,5 +24,4 @@ data_subset/**
 *.txt
 *.log
 xlmr-*/**
-**/checkpoint-*/**
-.venv/
+**/checkpoint-*/**

setup.py

@@ -5,13 +5,15 @@
     version="2.2.0",
     packages=find_packages(),
     description="Universal Robust, Efficient and Adaptable Sentence Segmentation",
+    long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type="text/markdown",
     author="Markus Frohmann, Igor Sterner, Benjamin Minixhofer",
     author_email="markus.frohmann@gmail.com",
     python_requires=">=3.9",
     install_requires=[
         # "onnxruntime>=1.13.1", # can make conflicts between onnxruntime and onnxruntime-gpu
         "transformers>=4.22.2",
-        "huggingface-hub<1.0",  # v1.0 has breaking changes (HfFolder removed)
+        "huggingface-hub",
         "numpy>=1.0",
         "scikit-learn>=1",
         "tqdm",

wtpsplit/__init__.py

@@ -7,13 +7,17 @@
 from pathlib import Path
 from typing import Literal
 
+import numpy as np
+
+# Import skops before transformers: skops enumerates trusted types at import time
+# and can otherwise trigger transformers' lazy vision submodules (torchvision), which
+# wtpsplit does not use for segmentation.
+import skops.io as sio
+
 # avoid the "None of PyTorch, TensorFlow, etc. have been found" warning.
 with contextlib.redirect_stderr(open(os.devnull, "w")):
     import transformers  # noqa
 
-import numpy as np
-import skops.io as sio
-
 from huggingface_hub import hf_hub_download
 from transformers import AutoConfig, AutoModelForTokenClassification, AutoTokenizer
 from transformers.utils.hub import cached_file

wtpsplit/train/adaptertrainer.py

@@ -1,67 +1,56 @@
 import os
-from typing import Dict
+import re
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
+from packaging import version
 from torch import nn
+from torch.utils.data import Dataset
+
+from transformers import PreTrainedModel, Trainer, __version__
+from transformers.configuration_utils import PretrainedConfig
+from transformers.data.data_collator import DataCollator
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_callback import (
+    TrainerCallback,
+    TrainerControl,
+    TrainerState,
+)
+from transformers.trainer_utils import EvalPrediction
+from transformers.training_args import TrainingArguments
+from transformers.utils import (
+    CONFIG_NAME,
+    is_accelerate_available,
+    is_apex_available,
+    logging,
+)
 
-from transformers import PreTrainedModel
-from transformers.modeling_utils import unwrap_model
-from transformers.trainer import (
+from wtpsplit.train.hf_compat import is_torch_tpu_available
+from wtpsplit.train.transformers_trainer_imports import (
     ALL_LAYERNORM_LAYERS,
-    WEIGHTS_NAME,
     DataLoader,
     EvalLoopOutput,
     IterableDatasetShard,
-    List,
-    Optional,
+    WEIGHTS_NAME,
     deepspeed_init,
     denumpify_detensorize,
     find_batch_size,
     get_parameter_names,
     has_length,
     is_sagemaker_mp_enabled,
-    is_torch_tpu_available,
     nested_concat,
     nested_numpify,
     nested_truncate,
+    unwrap_model,
 )
-
 from wtpsplit.train.utils import Model
 
 if is_torch_tpu_available(check_device=False):
     import torch_xla.core.xla_model as xm  # noqa: F401
     import torch_xla.debug.metrics as met  # noqa: F401
     import torch_xla.distributed.parallel_loader as pl  # noqa: F401
 
-import re
-from typing import Callable, Tuple, Union
-
-from packaging import version
-from torch.utils.data import Dataset
-from transformers import Trainer, __version__
-from transformers.configuration_utils import PretrainedConfig
-from transformers.data.data_collator import DataCollator
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase
-from transformers.trainer_callback import (
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-)
-from transformers.trainer_utils import (
-    EvalPrediction,
-)
-from transformers.training_args import TrainingArguments
-from transformers.utils import (
-    CONFIG_NAME,
-    is_accelerate_available,
-    is_apex_available,
-    is_torch_tpu_available,
-    logging,
-)
-
-from adapters.composition import AdapterCompositionBlock, Fuse
-
 if is_apex_available():
     pass
 
@@ -74,6 +63,7 @@
 
     if version.parse(accelerate_version) >= version.parse("0.16"):
         pass
+
 logger = logging.get_logger(__name__)
 
 TRAINING_ARGS_NAME = "training_args.bin"
@@ -121,8 +111,12 @@ def __init__(
             model_frozen = getattr(self.model.backbone.base_model, "model_frozen", False)
         else:
             model_frozen = False
+        # Default: same as pre-lazy-import behaviour when `model_frozen` is False (e.g. full finetune path).
+        self.train_adapter_fusion = False
         if model_frozen and self.model.backbone.active_adapters:
-            # Check if training AdapterFusion
+            # Check if training AdapterFusion (requires optional `adapters` package)
+            from adapters.composition import AdapterCompositionBlock, Fuse
+
             self.train_adapter_fusion = (
                 isinstance(self.model.backbone.active_adapters, Fuse)
                 or isinstance(self.model.backbone.active_adapters, AdapterCompositionBlock)
@@ -304,9 +298,7 @@ def evaluation_loop(
         if args.deepspeed and not self.deepspeed:
             # XXX: eval doesn't have `resume_from_checkpoint` arg but we should be able to do eval
             # from the checkpoint eventually
-            deepspeed_engine, _, _ = deepspeed_init(
-                self, num_training_steps=0, resume_from_checkpoint=None, inference=True
-            )
+            deepspeed_engine, _, _ = deepspeed_init(self, num_training_steps=0, inference=True)
             self.model = deepspeed_engine.module
             self.model_wrapped = deepspeed_engine
             self.deepspeed = deepspeed_engine

wtpsplit/train/hf_compat.py

@@ -0,0 +1,18 @@
+"""TPU / XLA checks compatible with **transformers 4.x and 5.x**.
+
+v5 removed ``is_torch_tpu_available`` from ``transformers.trainer``; ``is_torch_xla_available``
+exists on both and matches the old behaviour.
+"""
+
+from transformers.utils.import_utils import is_torch_xla_available
+
+
+def is_torch_tpu_available(check_device: bool = True) -> bool:
+    """Same idea as the old ``transformers.trainer.is_torch_tpu_available``.
+
+    - ``check_device=False``: ``torch_xla`` is importable (optional imports).
+    - ``check_device=True`` (default): current process is on a TPU.
+    """
+    if check_device:
+        return is_torch_xla_available(check_is_tpu=True)
+    return is_torch_xla_available()

wtpsplit/train/train.py

@@ -23,7 +23,8 @@
 from torchinfo import summary
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, HfArgumentParser, TrainingArguments, set_seed
-from transformers.trainer import is_torch_tpu_available
+
+from wtpsplit.train.hf_compat import is_torch_tpu_available
 
 import wandb
 from wtpsplit.models import (

wtpsplit/train/trainer.py

@@ -1,36 +1,33 @@
 import os
-from typing import Dict
+from typing import Dict, List, Optional
 
 import numpy as np
 import torch
 import transformers
 from torch import nn
 from torch.optim.lr_scheduler import LambdaLR
 from transformers import PreTrainedModel
-from transformers.trainer import (
+
+from wtpsplit.train.hf_compat import is_torch_tpu_available
+from wtpsplit.train.transformers_trainer_imports import (
     ALL_LAYERNORM_LAYERS,
-    TRAINING_ARGS_NAME,
-    WEIGHTS_NAME,
     DataLoader,
     EvalLoopOutput,
     IterableDatasetShard,
-    List,
-    Optional,
-    ShardedDDPOption,
+    TRAINING_ARGS_NAME,
+    WEIGHTS_NAME,
     deepspeed_init,
     denumpify_detensorize,
     find_batch_size,
     get_parameter_names,
     has_length,
     is_sagemaker_mp_enabled,
-    is_torch_tpu_available,
     logger,
     nested_concat,
     nested_numpify,
     nested_truncate,
+    unwrap_model,
 )
-from transformers.modeling_utils import unwrap_model
-
 from wtpsplit.train.utils import Model
 
 if is_torch_tpu_available(check_device=False):
@@ -108,19 +105,19 @@ def create_optimizer(self):
 
             optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)
 
-            if self.sharded_ddp == ShardedDDPOption.SIMPLE:
+            _sharded = getattr(self, "sharded_ddp", None)
+            if _sharded is not None and getattr(_sharded, "name", str(_sharded)) == "SIMPLE":
                 raise NotImplementedError()
-            else:
-                self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
-                if optimizer_cls.__name__ == "Adam8bit":
-                    import bitsandbytes
+            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
+            if optimizer_cls.__name__ == "Adam8bit":
+                import bitsandbytes
 
-                    manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
+                manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
 
-                    for module in opt_model.modules():
-                        if isinstance(module, nn.Embedding):
-                            manager.register_module_override(module, "weight", {"optim_bits": 32})
-                            logger.debug(f"bitsandbytes: will optimize {module} in fp32")
+                for module in opt_model.modules():
+                    if isinstance(module, nn.Embedding):
+                        manager.register_module_override(module, "weight", {"optim_bits": 32})
+                        logger.debug(f"bitsandbytes: will optimize {module} in fp32")
 
         if is_sagemaker_mp_enabled():
             raise NotImplementedError()
@@ -226,9 +223,7 @@ def evaluation_loop(
         if args.deepspeed and not self.deepspeed:
             # XXX: eval doesn't have `resume_from_checkpoint` arg but we should be able to do eval
             # from the checkpoint eventually
-            deepspeed_engine, _, _ = deepspeed_init(
-                self, num_training_steps=0, resume_from_checkpoint=None, inference=True
-            )
+            deepspeed_engine, _, _ = deepspeed_init(self, num_training_steps=0, inference=True)
             self.model = deepspeed_engine.module
             self.model_wrapped = deepspeed_engine
             self.deepspeed = deepspeed_engine

wtpsplit/train/transformers_trainer_imports.py

@@ -0,0 +1,47 @@
+"""Stable imports for custom ``Trainer`` subclasses on **transformers 4.x and 5.x**.
+
+In v5, many symbols are no longer re-exported from ``transformers.trainer``; importing from the
+same public submodules used internally works on both versions (``integrations.deepspeed``,
+``pytorch_utils``, ``trainer_pt_utils``, etc.).
+"""
+
+from transformers.integrations.deepspeed import deepspeed_init
+from transformers.modeling_utils import unwrap_model
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
+from transformers.trainer import (
+    DataLoader,
+    EvalLoopOutput,
+    IterableDatasetShard,
+    TRAINING_ARGS_NAME,
+    WEIGHTS_NAME,
+    logger,
+)
+from transformers.trainer_pt_utils import (
+    find_batch_size,
+    get_parameter_names,
+    nested_concat,
+    nested_numpify,
+    nested_truncate,
+)
+from transformers.trainer_utils import denumpify_detensorize, has_length
+from transformers.utils import is_sagemaker_mp_enabled
+
+__all__ = [
+    "ALL_LAYERNORM_LAYERS",
+    "DataLoader",
+    "EvalLoopOutput",
+    "IterableDatasetShard",
+    "TRAINING_ARGS_NAME",
+    "WEIGHTS_NAME",
+    "deepspeed_init",
+    "denumpify_detensorize",
+    "find_batch_size",
+    "get_parameter_names",
+    "has_length",
+    "is_sagemaker_mp_enabled",
+    "logger",
+    "nested_concat",
+    "nested_numpify",
+    "nested_truncate",
+    "unwrap_model",
+]

wtpsplit/train/utils.py

@@ -50,6 +50,9 @@ def forward(
             else:
                 reduced_attention_mask = (input_ids != 0).to(torch.long)
 
+        # Transformers 4.51+ Trainer may pass this; HF models ignore it, custom backbones must not receive it.
+        kwargs.pop("num_items_in_batch", None)
+
         output = dict(
             self.backbone.forward(
                 input_ids=input_ids,