wtpsplit 2.2.1__tar.gz → 2.2.2__tar.gz

{wtpsplit-2.2.1 → wtpsplit-2.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wtpsplit
-Version: 2.2.1
+Version: 2.2.2
 Summary: Universal Robust, Efficient and Adaptable Sentence Segmentation
 Home-page: https://github.com/segment-any-text/wtpsplit
 Author: Markus Frohmann, Igor Sterner, Benjamin Minixhofer
@@ -23,6 +23,9 @@ Provides-Extra: onnx-gpu
 Requires-Dist: onnxruntime-gpu>=1.13.1; extra == "onnx-gpu"
 Provides-Extra: onnx-cpu
 Requires-Dist: onnxruntime>=1.13.1; extra == "onnx-cpu"
+Provides-Extra: aitune
+Requires-Dist: aitune>=0.3.0; extra == "aitune"
+Requires-Dist: requests; extra == "aitune"
 Dynamic: author
 Dynamic: author-email
 Dynamic: description
@@ -88,6 +91,44 @@ sat_adapted.split("This is a test This is another test.")
 # returns ['This is a test ', 'This is another test']
 ```
 
+## Faster PyTorch inference (`torch.compile` / TorchInductor)
+
+For repeated inference without exporting to ONNX, compile the PyTorch model with [TorchInductor](https://pytorch.org/docs/stable/torch.compiler.html) via `optimize()` (PyTorch 2.0+). This is **optional** - the first `split` after `optimize()` can be slow while graphs are built. Call `optimize()` **after** `to()` / `half()` so the compiled graph matches *your* device and dtype.
+
+```python
+from wtpsplit import SaT
+
+sat = SaT("sat-3l-sm")
+sat.half().to("cuda")
+sat.optimize()  # backend="inductor" (default); dynamic shapes on
+
+sat.split("This is a test This is another test.")
+```
+
+- **SaT** and **WtP** PyTorch checkpoints only — not available with `ort_providers` / ONNX.
+- `backend` synonyms: `"inductor"`, `"torchinductor"`.
+- Chunk length and the last batch size change between calls, so `dynamic=True` is the default. `mode="reduce-overhead"` (CUDA graphs) is faster only when every forward uses the same shapes; otherwise use the default or `mode="max-autotune-no-cudagraphs"`.
+- Optional on NVIDIA Ampere+: `torch.set_float32_matmul_precision("high")` before inference (faster fp32 matmuls).
+
+### NVIDIA AITune (auto backend selection, CUDA)
+
+[NVIDIA AITune](https://github.com/ai-dynamo/aitune) benchmarks backends (TensorRT, Torch-TensorRT, Torch Inductor, …) on your GPU and picks a fast path. Requires Linux, CUDA, and a separate install:
+
+```bash
+pip install wtpsplit[aitune] --extra-index-url https://pypi.nvidia.com
+```
+
+```python
+sat = SaT("sat-3l-sm")
+sat.half().to("cuda")
+# first_wins tries TensorRT → Torch-TensorRT → Inductor when available
+sat.optimize(backend="aitune")
+# or only Inductor via AITune (faster tuning, no TensorRT dependency):
+# sat.optimize(backend="aitune", aitune_strategy="inductor_only", aitune_max_batches=4)
+
+sat.split("This is a test This is another test.")
+```
+
 ## ONNX Support
 
 🚀 You can now enable even faster ONNX inference for `sat` and `sat-sm` models! 🚀

wtpsplit-2.2.1/wtpsplit.egg-info/PKG-INFO → wtpsplit-2.2.2/README.md

@@ -1,40 +1,3 @@
-Metadata-Version: 2.4
-Name: wtpsplit
-Version: 2.2.1
-Summary: Universal Robust, Efficient and Adaptable Sentence Segmentation
-Home-page: https://github.com/segment-any-text/wtpsplit
-Author: Markus Frohmann, Igor Sterner, Benjamin Minixhofer
-Author-email: markus.frohmann@gmail.com
-License: MIT
-Requires-Python: >=3.9
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: transformers>=4.22.2
-Requires-Dist: huggingface-hub
-Requires-Dist: numpy>=1.0
-Requires-Dist: scikit-learn>=1
-Requires-Dist: tqdm
-Requires-Dist: skops
-Requires-Dist: pandas>=1
-Requires-Dist: mosestokenizer
-Provides-Extra: adapters
-Requires-Dist: adapters>=1.0.1; extra == "adapters"
-Provides-Extra: onnx-gpu
-Requires-Dist: onnxruntime-gpu>=1.13.1; extra == "onnx-gpu"
-Provides-Extra: onnx-cpu
-Requires-Dist: onnxruntime>=1.13.1; extra == "onnx-cpu"
-Dynamic: author
-Dynamic: author-email
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: license
-Dynamic: license-file
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
 <h1 align="center">wtpsplit🪓</h1>
 <h3 align="center">Segment any Text - Robustly, Efficiently, Adaptably⚡</h3>
 
@@ -88,6 +51,44 @@ sat_adapted.split("This is a test This is another test.")
 # returns ['This is a test ', 'This is another test']
 ```
 
+## Faster PyTorch inference (`torch.compile` / TorchInductor)
+
+For repeated inference without exporting to ONNX, compile the PyTorch model with [TorchInductor](https://pytorch.org/docs/stable/torch.compiler.html) via `optimize()` (PyTorch 2.0+). This is **optional** - the first `split` after `optimize()` can be slow while graphs are built. Call `optimize()` **after** `to()` / `half()` so the compiled graph matches *your* device and dtype.
+
+```python
+from wtpsplit import SaT
+
+sat = SaT("sat-3l-sm")
+sat.half().to("cuda")
+sat.optimize()  # backend="inductor" (default); dynamic shapes on
+
+sat.split("This is a test This is another test.")
+```
+
+- **SaT** and **WtP** PyTorch checkpoints only — not available with `ort_providers` / ONNX.
+- `backend` synonyms: `"inductor"`, `"torchinductor"`.
+- Chunk length and the last batch size change between calls, so `dynamic=True` is the default. `mode="reduce-overhead"` (CUDA graphs) is faster only when every forward uses the same shapes; otherwise use the default or `mode="max-autotune-no-cudagraphs"`.
+- Optional on NVIDIA Ampere+: `torch.set_float32_matmul_precision("high")` before inference (faster fp32 matmuls).
+
+### NVIDIA AITune (auto backend selection, CUDA)
+
+[NVIDIA AITune](https://github.com/ai-dynamo/aitune) benchmarks backends (TensorRT, Torch-TensorRT, Torch Inductor, …) on your GPU and picks a fast path. Requires Linux, CUDA, and a separate install:
+
+```bash
+pip install wtpsplit[aitune] --extra-index-url https://pypi.nvidia.com
+```
+
+```python
+sat = SaT("sat-3l-sm")
+sat.half().to("cuda")
+# first_wins tries TensorRT → Torch-TensorRT → Inductor when available
+sat.optimize(backend="aitune")
+# or only Inductor via AITune (faster tuning, no TensorRT dependency):
+# sat.optimize(backend="aitune", aitune_strategy="inductor_only", aitune_max_batches=4)
+
+sat.split("This is a test This is another test.")
+```
+
 ## ONNX Support
 
 🚀 You can now enable even faster ONNX inference for `sat` and `sat-sm` models! 🚀

{wtpsplit-2.2.1 → wtpsplit-2.2.2}/pyproject.toml

@@ -4,7 +4,11 @@ target-version = ["py39", "py310", "py311", "py312"]
 
 [tool.ruff]
 line-length = 120
-lint.ignore = ["E741"]
+
+[tool.ruff.lint]
+# Pre-0.16 default. Ruff 0.16 enables 413 rules by default; CI installs latest ruff.
+select = ["E4", "E7", "E9", "F"]
+ignore = ["E741"]
 
 [tool.ruff.lint.per-file-ignores]
 "test.py" = ["E501"]

{wtpsplit-2.2.1 → wtpsplit-2.2.2}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="wtpsplit",
-    version="2.2.1",
+    version="2.2.2",
     packages=find_packages(),
     description="Universal Robust, Efficient and Adaptable Sentence Segmentation",
     long_description=open("README.md", encoding="utf-8").read(),
@@ -25,6 +25,7 @@ setup(
         "adapters": ["adapters>=1.0.1"],
         "onnx-gpu": ["onnxruntime-gpu>=1.13.1"],
         "onnx-cpu": ["onnxruntime>=1.13.1"],
+        "aitune": ["aitune>=0.3.0", "requests"],
     },
     url="https://github.com/segment-any-text/wtpsplit",
     package_data={"wtpsplit": ["data/*"]},

{wtpsplit-2.2.1 → wtpsplit-2.2.2}/wtpsplit/__init__.py

@@ -30,7 +30,7 @@ from wtpsplit.utils.constraints import (
 )
 from wtpsplit.utils.priors import create_prior_function
 
-__version__ = "2.2.1"
+__version__ = "2.2.2"
 
 # suppress docopt syntax warnings (triggered in Python 3.14+)
 warnings.filterwarnings("ignore", category=SyntaxWarning, module="docopt")
@@ -228,6 +228,27 @@ class WtP:
         assert hasattr(self, "model")
         return getattr(self.model, name)
 
+    def optimize(
+        self,
+        backend: str = "inductor",
+        *,
+        mode: str = None,
+        fullgraph: bool = False,
+        dynamic: bool = True,
+        **compile_kwargs,
+    ):
+        """Compile the PyTorch model with ``torch.compile`` (TorchInductor by default). Not available for ONNX."""
+        if not isinstance(self.model, PyTorchWrapper):
+            raise ValueError(
+                "optimize() only applies to PyTorch models loaded from a checkpoint. "
+                "It does not apply to ONNX Runtime (set ort_providers=None) or to ad-hoc model objects "
+                "that are not wrapped as expected."
+            )
+        self.model.optimize(
+            backend=backend, mode=mode, fullgraph=fullgraph, dynamic=dynamic, **compile_kwargs
+        )
+        return self
+
     def predict_proba(
         self,
         text_or_texts,
@@ -908,6 +929,27 @@ class SaT:
         assert hasattr(self, "model")
         return getattr(self.model, name)
 
+    def optimize(
+        self,
+        backend: str = "inductor",
+        *,
+        mode: str = None,
+        fullgraph: bool = False,
+        dynamic: bool = True,
+        **compile_kwargs,
+    ):
+        """Compile the PyTorch model with ``torch.compile`` (TorchInductor by default). Not available for ONNX."""
+        if not isinstance(self.model, PyTorchWrapper):
+            raise ValueError(
+                "optimize() only applies to PyTorch models loaded from a checkpoint. "
+                "It does not apply to ONNX Runtime (set ort_providers=None) or to ad-hoc model objects "
+                "that are not wrapped as expected."
+            )
+        self.model.optimize(
+            backend=backend, mode=mode, fullgraph=fullgraph, dynamic=dynamic, **compile_kwargs
+        )
+        return self
+
     def predict_proba(
         self,
         text_or_texts,

wtpsplit-2.2.2/wtpsplit/aitune_integration.py

@@ -0,0 +1,279 @@
+"""Optional NVIDIA AITune integration for :meth:`PyTorchWrapper.optimize`."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def _require_aitune():
+    try:
+        import aitune.torch as ait  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            "backend='aitune' requires NVIDIA AITune. Install with:\n"
+            "  pip install wtpsplit[aitune] --extra-index-url https://pypi.nvidia.com"
+        ) from e
+    return ait
+
+
+def _build_calibration_dataset(model, device) -> list[dict[str, Any]]:
+    """Synthetic unbatched samples covering the chunk lengths ``extract`` actually feeds.
+
+    AITune stacks these on dim 0, so each tensor is one example: ``(seq,)`` or
+    ``(seq, num_hashes)``. Lengths stay inside the position-embedding window and
+    on the character-model downsampling grid.
+    """
+    import torch
+
+    config = model.config
+    model_type = getattr(config, "model_type", "") or ""
+    rate = int(getattr(config, "downsampling_rate", 1) or 1)
+    lengths = tuple(length for length in (64, 128, 256, 512) if length % rate == 0) or (rate,)
+    uses_lang = getattr(config, "language_adapter", "off") in {"on", "shared"}
+    samples: list[dict[str, Any]] = []
+
+    for seq_len in lengths:
+        if "xlm" in model_type:
+            # CLS/SEP are added around the chunk; stay within the 512 position window.
+            length = min(seq_len + 2, 512)
+            sample = {
+                "input_ids": torch.randint(0, 50000, (length,), dtype=torch.long, device=device),
+                "attention_mask": torch.ones(length, dtype=torch.float32, device=device),
+            }
+        else:
+            num_hashes = getattr(config, "num_hash_functions", 8)
+            num_buckets = getattr(config, "num_hash_buckets", 10000)
+            sample = {
+                "hashed_ids": torch.randint(0, num_buckets, (seq_len, num_hashes), dtype=torch.long, device=device),
+                "attention_mask": torch.ones(seq_len, dtype=torch.float32, device=device),
+            }
+        if uses_lang and "xlm" not in model_type:
+            # Batched to shape (batch,) by AITune, matching ``extract``.
+            sample["language_ids"] = torch.zeros((), dtype=torch.long, device=device)
+        samples.append(sample)
+
+    return samples
+
+
+def _backend_cls(*names: str):
+    for name in names:
+        factory = _try_backend("aitune.torch.backend", name)
+        if factory is not None:
+            return factory
+    raise ImportError(
+        "Could not import an AITune TorchInductor backend "
+        f"(tried {', '.join(names)}). Upgrade with: pip install -U 'aitune>=0.4'"
+    )
+
+
+def _strategy_cls(*names: str):
+    for name in names:
+        factory = _try_backend("aitune.torch.tune_strategy", name)
+        if factory is not None:
+            return factory
+    raise ImportError(
+        "Could not import an AITune tune strategy "
+        f"(tried {', '.join(names)}). Upgrade with: pip install -U 'aitune>=0.4'"
+    )
+
+
+def _inductor_backend(*, mode: str | None, fullgraph: bool, dynamic: bool):
+    """TorchInductor backend, using the current JIT class and falling back to older names."""
+    backend_cls = _backend_cls("TorchInductorJitBackend", "TorchInductorBackend")
+    config_cls = _try_backend("aitune.torch.backend", "TorchInductorJitBackendConfig")
+    if config_cls is None:
+        return backend_cls()
+    config_kwargs: dict[str, Any] = {"fullgraph": fullgraph, "dynamic": dynamic}
+    if mode is not None:
+        config_kwargs["mode"] = mode
+    try:
+        return backend_cls(config=config_cls(**config_kwargs))
+    except TypeError:
+        return backend_cls()
+
+
+def _resolve_strategy(
+    strategy_name: str | None,
+    backends: list | None,
+    *,
+    mode: str | None,
+    fullgraph: bool,
+    dynamic: bool,
+):
+    first_wins = _strategy_cls("FirstWinsStrategy")
+    one_backend = _strategy_cls("OneBackendStrategy")
+
+    if backends is not None:
+        return first_wins(backends=backends)
+
+    key = (strategy_name or "first_wins").lower().replace("-", "_")
+    inductor = _inductor_backend(mode=mode, fullgraph=fullgraph, dynamic=dynamic)
+    if key in ("inductor", "inductor_only", "torch_inductor"):
+        return one_backend(backend=inductor)
+
+    available = _default_aitune_backends(inductor)
+    if key in ("highest_throughput", "best", "max_throughput"):
+        strategy_cls = _strategy_cls("MaxThroughputStrategy", "HighestThroughputStrategy")
+        try:
+            return strategy_cls(backends=available)
+        except TypeError:
+            return strategy_cls()
+
+    return first_wins(backends=available)
+
+
+def _default_aitune_backends(inductor_backend) -> list:
+    """Backends tried in order for ``first_wins`` (TensorRT first when importable)."""
+    backends = []
+    for class_name in ("TensorRTBackend", "TorchTensorRTJitBackend"):
+        factory = _try_backend("aitune.torch.backend", class_name)
+        if factory is not None:
+            backends.append(factory())
+    backends.append(inductor_backend)
+    return backends
+
+
+def _try_backend(module_name: str, class_name: str):
+    try:
+        import importlib
+
+        mod = importlib.import_module(module_name)
+        return getattr(mod, class_name)
+    except Exception:
+        return None
+
+
+def pop_aitune_kwargs(compile_kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Extract AITune-specific keys from ``compile_kwargs`` (mutates the dict)."""
+    keys = (
+        "aitune_batch_sizes",
+        "aitune_max_batches",
+        "aitune_strategy",
+        "aitune_backends",
+        "aitune_dry_run",
+        "aitune_calibration",
+    )
+    out = {}
+    for key in keys:
+        if key in compile_kwargs:
+            out[key] = compile_kwargs.pop(key)
+    return out
+
+
+def apply_aitune(model, *, mode: str | None = None, fullgraph: bool = False, dynamic: bool = True, **aitune_kwargs):
+    """Wrap ``model`` with AITune and run ahead-of-time tuning. Returns the wrapped module."""
+    import logging
+
+    import torch
+
+    ait = _require_aitune()
+    from aitune.torch.dataloader import DynamicShapeDataset
+
+    logger = logging.getLogger(__name__)
+
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "AITune requires a CUDA GPU. Move the model with `.to('cuda')` before `optimize(backend='aitune')`."
+        )
+
+    device = getattr(model, "device", None)
+    if device is None or getattr(device, "type", None) != "cuda":
+        raise RuntimeError(
+            "AITune requires the model on CUDA. Call `.to('cuda')` before `optimize(backend='aitune')`."
+        )
+
+    model.eval()
+    calibration = aitune_kwargs.pop("aitune_calibration", None)
+    dataset = (
+        calibration if calibration is not None else DynamicShapeDataset(_build_calibration_dataset(model, device))
+    )
+    batch_sizes = aitune_kwargs.pop("aitune_batch_sizes", None) or [1, 2]
+    max_batches = aitune_kwargs.pop("aitune_max_batches", 4)
+    strategy_name = aitune_kwargs.pop("aitune_strategy", None)
+    backends = aitune_kwargs.pop("aitune_backends", None)
+    dry_run = aitune_kwargs.pop("aitune_dry_run", False)
+    if len(set(batch_sizes)) < 2:
+        logger.warning(
+            "AITune needs at least two batch sizes to mark the batch axis as dynamic; "
+            "using [1, 2] so later split() batch sizes still match."
+        )
+        batch_sizes = [1, 2]
+
+    strategy = _resolve_strategy(
+        strategy_name, backends, mode=mode, fullgraph=fullgraph, dynamic=dynamic
+    )
+    name = f"{model.__class__.__module__}.{model.__class__.__qualname__}".replace("/", "_")
+    wrapped = ait.Module(model, name=name, strategy=strategy)
+
+    ait.tune(
+        wrapped,
+        dataset,
+        batch_sizes=batch_sizes,
+        max_num_batches_per_batch_size=max_batches,
+        device=device,
+        dry_run=dry_run,
+    )
+    if not dry_run:
+        _install_metadata_fallback(wrapped)
+    return wrapped
+
+
+def _install_metadata_fallback(wrapped):
+    """Allow AITune dynamic-shape metadata to match runtime samples.
+
+    AITune records dynamic axes during tuning, but its runtime lookup can still
+    require exact ``SampleMetadata`` dictionary membership. SaT/WtP chunk lengths
+    vary naturally, so fall back to the compatible ``TensorSpec.matches`` check
+    before reporting that no backend was found.
+    """
+    import types
+
+    from aitune.torch.module.sample_metadata import SampleMetadata
+
+    tuned = getattr(wrapped, "_self_wrapper", None)
+    if tuned is None or not hasattr(tuned, "_backends"):
+        return
+
+    def metadata_matches(expected, actual):
+        if expected == actual:
+            return True
+        if expected.llm_phase != actual.llm_phase:
+            return False
+        if expected.other_data != actual.other_data:
+            return False
+
+        expected_specs = {spec.name: spec for spec in expected.tensor_specs}
+        actual_specs = {spec.name: spec for spec in actual.tensor_specs}
+        if expected_specs.keys() != actual_specs.keys():
+            return False
+
+        for name, actual_spec in actual_specs.items():
+            expected_spec = expected_specs[name]
+            if expected_spec.dtype != actual_spec.dtype:
+                return False
+            if not expected_spec.matches(actual_spec):
+                return False
+        return True
+
+    def safe_call_backend_with_dynamic_match(self, sample):
+        args, kwargs = sample
+        # AITune 0.6+ takes one named-input dict. Older releases took (args, kwargs).
+        if hasattr(self, "_forward_signature"):
+            forward_inputs = self._forward_signature.normalize(args, kwargs)
+            sample_metadata = SampleMetadata.from_inputs(
+                forward_inputs.arguments, strict=self._config.strict_mode
+            )
+        else:
+            sample_metadata = SampleMetadata.from_inputs(args, kwargs, strict=self._config.strict_mode)
+        backend = self._backends.get(sample_metadata)
+        if backend is None:
+            for candidate_metadata, candidate_backend in self._backends.items():
+                if metadata_matches(candidate_metadata, sample_metadata):
+                    backend = candidate_backend
+                    break
+        if backend is None:
+            raise RuntimeError(self.ERROR_NO_BACKEND_FOUND.format(sample_metadata))
+        return backend.infer(*args, **kwargs)
+
+    tuned.safe_call_backend = types.MethodType(safe_call_backend_with_dynamic_match, tuned)
+    tuned._backend_func = tuned.safe_call_backend