wtpsplit 2.2.0__tar.gz → 2.2.2__tar.gz

wtpsplit-2.2.2/PKG-INFO

@@ -0,0 +1,620 @@
+Metadata-Version: 2.4
+Name: wtpsplit
+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
+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"
+Provides-Extra: aitune
+Requires-Dist: aitune>=0.3.0; extra == "aitune"
+Requires-Dist: requests; extra == "aitune"
+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>
+
+This repository allows you to segment text into sentences or other semantic units. It implements the models from:
+
+- **SaT** &mdash; [Segment Any Text: A Universal Approach for Robust, Efficient and Adaptable Sentence Segmentation](https://arxiv.org/abs/2406.16678) by Markus Frohmann, Igor Sterner, Benjamin Minixhofer, Ivan Vulić and Markus Schedl (**state-of-the-art, encouraged**).
+- **WtP** &mdash; [Where’s the Point? Self-Supervised Multilingual Punctuation-Agnostic Sentence Segmentation](https://aclanthology.org/2023.acl-long.398/) by Benjamin Minixhofer, Jonas Pfeiffer and Ivan Vulić (*previous version, maintained for reproducibility*).
+
+The namesake WtP is maintained for consistency. Our new followup SaT provides robust, efficient and adaptable sentence segmentation across 85 languages at higher performance and less compute cost. Check out the **state-of-the-art** results in 8 distinct corpora and 85 languages demonstrated in our [Segment any Text paper](https://arxiv.org/abs/2406.16678).
+
+![System Figure](./configs/system-fig.png)
+
+## Installation
+
+```bash
+pip install wtpsplit
+```
+Or one of the following for ONNX support:
+```bash
+pip install wtpsplit[onnx-gpu]
+pip install wtpsplit[onnx-cpu]
+```
+
+## Usage
+
+```python
+from wtpsplit import SaT
+
+sat = SaT("sat-3l")
+# optionally run on GPU for better performance
+# also supports TPUs via e.g. sat.to("xla:0"), in that case pass `pad_last_batch=True` to sat.split
+sat.half().to("cuda")
+
+sat.split("This is a test This is another test.")
+# returns ["This is a test ", "This is another test."]
+
+# do this instead of calling sat.split on every text individually for much better performance
+sat.split(["This is a test This is another test.", "And some more texts..."])
+# returns an iterator yielding lists of sentences for every text
+
+# use our '-sm' models for general sentence segmentation tasks
+sat_sm = SaT("sat-3l-sm")
+sat_sm.half().to("cuda") # optional, see above
+sat_sm.split("this is a test this is another test")
+# returns ["this is a test ", "this is another test"]
+
+# use trained lora modules for strong adaptation to language & domain/style
+sat_adapted = SaT("sat-3l", style_or_domain="ud", language="en")
+sat_adapted.half().to("cuda") # optional, see above
+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! 🚀
+
+```python
+sat = SaT("sat-3l-sm", ort_providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+```
+
+```python
+>>> from wtpsplit import SaT
+>>> texts = ["This is a sentence. This is another sentence."] * 1000
+
+# PyTorch GPU
+>>> model_pytorch = SaT("sat-3l-sm")
+>>> model_pytorch.half().to("cuda");
+>>> %timeit list(model_pytorch.split(texts))
+# 144 ms ± 252 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
+# quite fast already, but...
+
+# onnxruntime GPU
+>>> model_ort = SaT("sat-3l-sm", ort_providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+>>> %timeit list(model_ort.split(texts))
+# 94.9 ms ± 165 μs per loop (mean ± std. dev. of 7 runs, 10 loops each
+# ...this should be ~50% faster! (tested on RTX 3090)
+```
+
+If you wish to use LoRA in combination with an ONNX model:
+
+- Run `scripts/export_to_onnx_sat.py` with `use_lora: True` and an appropriate `output_dir: <OUTPUT_DIR>`.
+  - If you have a local LoRA module, use `lora_path`.
+  - If you wish to load a LoRA module from the HuggingFace hub, use `style_or_domain` and `language`.
+- Load the ONNX model with merged LoRA weights:
+  `sat = SaT(<OUTPUT_DIR>, onnx_providers=["CUDAExecutionProvider", "CPUExecutionProvider"])`
+
+## Available Models
+
+If you need a general sentence segmentation model, use `-sm` models (e.g., `sat-3l-sm`)
+For speed-sensitive applications, we recommend 3-layer models (`sat-3l` and `sat-3l-sm`). They provide a great tradeoff between speed and performance.
+The best models are our 12-layer models: `sat-12l` and `sat-12l-sm`.
+
+| Model                                                                        | English Score | Multilingual Score |
+| :--------------------------------------------------------------------------- | ------------: | -----------------: |
+| [sat-1l](https://huggingface.co/segment-any-text/sat-1l)                        |          88.5 |               84.3 |
+| [sat-1l-sm](https://huggingface.co/segment-any-text/sat-1l-sm)                  |          88.2 |               87.9 |
+| [sat-3l](https://huggingface.co/segment-any-text/sat-3l)                        |          93.7 |               89.2 |
+| [sat-3l-lora](https://huggingface.co/segment-any-text/sat-3l/tree/main/loras)   |          96.7 |               94.8 |
+| [sat-3l-sm](https://huggingface.co/segment-any-text/sat-3l-sm)                  |          96.5 |               93.5 |
+| [sat-6l](https://huggingface.co/segment-any-text/sat-6l)                        |          94.1 |               89.7 |
+| [sat-6l-sm](https://huggingface.co/segment-any-text/sat-6l-sm)                  |          96.9 |               95.1 |
+| [sat-9l](https://huggingface.co/segment-any-text/sat-9l)                        |          94.3 |               90.3 |
+| [sat-12l](https://huggingface.co/segment-any-text/sat-12l)                      |          94.0 |               90.4 |
+| [sat-12l-lora](https://huggingface.co/segment-any-text/sat-12l/tree/main/loras) |          97.3 |               95.9 |
+| [sat-12l-sm](https://huggingface.co/segment-any-text/sat-12l-sm)                |          97.4 |               96.0 |
+
+The scores are macro-average F1 score across all available datasets for "English", and macro-average F1 score across all datasets and languages for "Multilingual". "adapted" means adapation via LoRA; check out the [paper](https://arxiv.org/abs/2406.16678) for details.
+
+For comparison, here the English scores of some other tools:
+
+| Model                                                    | English Score |
+| :------------------------------------------------------- | ------------: |
+| PySBD                                                    |          69.6 |
+| SpaCy (sentencizer; monolingual)                         |          92.9 |
+| SpaCy (sentencizer; multilingual)                        |          91.5 |
+| Ersatz                                                   |          91.4 |
+| Punkt (`nltk.sent_tokenize`)                           |          92.2 |
+| [WtP (3l)](https://huggingface.co/benjamin/wtp-canine-s-3l) |          93.9 |
+
+Note that this library also supports previous [`WtP`](https://arxiv.org/abs/2305.18893) models.
+You can use them in essentially the same way as `SaT`models:
+
+```python
+from wtpsplit import WtP
+
+wtp = WtP("wtp-bert-mini")
+# similar functionality as for SaT models
+wtp.split("This is a test This is another test.")
+```
+
+For more details on WtP and reproduction details, see the [WtP doc](./README_WTP.md).
+
+## Paragraph Segmentation
+
+Since SaT are trained to predict newline probablity, they can segment text into paragraphs in addition to sentences.
+
+```python
+# returns a list of paragraphs, each containing a list of sentences
+# adjust the paragraph threshold via the `paragraph_threshold` argument.
+sat.split(text, do_paragraph_segmentation=True)
+```
+
+## (NEW! v2.2+) Length-Constrained Segmentation
+
+Control segment lengths with `min_length` and `max_length` parameters. This is useful when you need segments within specific size limits (e.g., for embedding models, storage, or downstream processing).
+
+### Basic Usage
+
+```python
+from wtpsplit import SaT
+
+sat = SaT("sat-3l-sm")
+
+text = (
+    "In the beginning God created the heaven and the earth. "
+    "And the earth was without form, and void; and darkness was upon the face of the deep. "
+    "And the Spirit of God moved upon the face of the waters. "
+    "And God said, Let there be light: and there was light. "
+    "And God saw the light, that it was good: and God divided the light from the darkness. "
+    "And God called the light Day, and the darkness he called Night. "
+    "And the evening and the morning were the first day."
+)
+
+# Split with a maximum segment length of 120 characters
+segments = sat.split(text, max_length=120)
+for i, s in enumerate(segments):
+    print(f"[{len(s):3d} chars] {s}")
+# [ 55 chars] In the beginning God created the heaven and the earth. 
+# [ 86 chars] And the earth was without form, and void; and darkness was upon the face of the deep. 
+# [112 chars] And the Spirit of God moved upon the face of the waters. And God said, Let there be light: and there was light. 
+# [ 86 chars] And God saw the light, that it was good: and God divided the light from the darkness. 
+# [115 chars] And God called the light Day, and the darkness he called Night. And the evening and the morning were the first day.
+
+assert "".join(segments) == text  # text is perfectly preserved
+
+# Enforce both min and max length
+sat.split(text, min_length=80, max_length=200)
+
+# Use the greedy algorithm for minimally faster (but less optimal) results
+sat.split(text, max_length=120, algorithm="greedy")
+```
+
+### Priors for Length Preference
+
+Use priors to influence segment length distribution. Available priors:
+
+| Prior | Best For |
+|-------|----------|
+| `"uniform"` (default) | Just enforce max_length, let model decide |
+| `"gaussian"` | Prefer segments around a target length (intuitive) |
+| `"lognormal"` | Right-skewed preference (more tolerant of longer segments) |
+| `"clipped_polynomial"` | Must be very close to target length |
+
+```python
+# Gaussian prior (recommended): prefer segments around target_length
+sat.split(text, max_length=100, prior_type="gaussian", 
+          prior_kwargs={"target_length": 50, "spread": 10})
+
+# Log-normal prior: right-skewed (more tolerant of longer segments)
+sat.split(text, max_length=100, prior_type="lognormal", 
+          prior_kwargs={"target_length": 70, "spread": 25})
+
+# Clipped polynomial: hard cutoff at ±spread from target
+sat.split(text, max_length=100, prior_type="clipped_polynomial", 
+          prior_kwargs={"target_length": 60, "spread": 25})
+```
+
+### Language-Aware Defaults
+
+Pass `lang_code` to use language-specific defaults for `target_length` and `spread` (based on language-specific corpus statistics):
+
+```python
+# German has longer average sentences → auto-uses target_length=90, spread=35
+sat.split(text, max_length=150, prior_type="gaussian", 
+          prior_kwargs={"lang_code": "de"})
+
+# Chinese has shorter sentences → auto-uses target_length=45, spread=15
+sat.split(text, max_length=100, prior_type="gaussian", 
+          prior_kwargs={"lang_code": "zh"})
+```
+
+When using LoRA with a language, this happens automatically:
+
+```python
+sat = SaT("sat-3l", style_or_domain="ud", language="de")
+sat.split(text, max_length=150, prior_type="gaussian")  # auto-uses German defaults
+```
+
+### How It Works
+
+The Viterbi algorithm finds globally optimal segmentation points that balance:
+- The model's sentence boundary predictions (where natural splits occur)
+- Your length preferences (via the prior; if provided)
+
+**Text Reconstruction:**
+```python
+# With constraints (max_length or min_length):
+original_text = "".join(segments)  # segments may contain newlines
+
+# Without constraints (SaT default with split_on_input_newlines=True):
+original_text = "\n".join(segments)
+```
+
+> **Note**: When using length constraints, segments may contain newlines. If you want to remove them, you can just post-process the output.
+
+> **Note**: When `max_length` is set, the `threshold` parameter is ignored. The Viterbi/greedy algorithms use raw model probabilities directly instead of threshold-based filtering.
+
+For more details, see the [Length Constraints Documentation](./docs/LENGTH_CONSTRAINTS.md).
+
+## Adaptation
+
+SaT can be domain- and style-adapted via LoRA. We provide trained LoRA modules for Universal Dependencies, OPUS100, Ersatz, and TED (i.e., ASR-style transcribed speeches) sentence styles in 81 languages for `sat-3l`and `sat-12l`. Additionally, we provide LoRA modules for legal documents (laws and judgements) in 6 languages, code-switching in 4 language pairs, and tweets in 3 languages. For details, we refer to our [paper](https://arxiv.org/abs/2406.16678).
+
+We also provided verse segmentation modules for 16 genres for `sat-12-no-limited-lookahead`.
+
+Load LoRA modules like this:
+
+```python
+
+# requires both lang_code and style_or_domain
+# for available ones, check the <model_repository>/loras folder
+sat_lora = SaT("sat-3l", style_or_domain="ud", language="en")
+sat_lora.split("Hello this is a test But this is different now Now the next one starts looool")
+# now for a highly distinct domain
+sat_lora_distinct = SaT("sat-12l", style_or_domain="code-switching", language="es-en")
+sat_lora_distinct.split("in the morning over there cada vez que yo decía algo él me decía algo")
+```
+
+You can also freely adapt the segmentation threshold, with a higher threshold leading to more conservative segmentation:
+
+```python
+
+sat.split("This is a test This is another test.", threshold=0.4)
+# works similarly for lora; but thresholds are higher
+sat_lora.split("Hello this is a test But this is different now Now the next one starts looool", threshold=0.7)
+```
+
+## Advanced Usage
+
+### Get the newline or sentence boundary probabilities for a text:
+
+```python
+# returns newline probabilities (supports batching!)
+sat.predict_proba(text)
+```
+
+### Load a SaT model in [HuggingFace `transformers`](https://github.com/huggingface/transformers):
+
+```python
+# import library to register the custom models 
+import wtpsplit.models
+from transformers import AutoModelForTokenClassification
+
+model = AutoModelForTokenClassification.from_pretrained("segment-any-text/sat-3l-sm") # or some other model name; see https://huggingface.co/segment-any-text
+```
+
+### Adapt to your own corpus via LoRA
+
+Our models can be efficiently adapted via LoRA in a powerful way. Only 10-100 training segmented training sentences should already improve performance considerably. To do so:
+
+Clone the repository and install requirements:
+
+```
+git clone https://github.com/segment-any-text/wtpsplit
+cd wtpsplit
+pip install -r requirements.txt
+pip install adapters==0.2.1 --no-dependencies
+cd ..
+```
+
+1. Create data in this format:
+
+```python
+import torch
+
+torch.save(
+    {
+        "language_code": {
+            "sentence": {
+                "dummy-dataset": {
+                    "meta": {
+                        "train_data": ["train sentence 1", "train sentence 2"],
+                    },
+                    "data": [
+                        "test sentence 1",
+                        "test sentence 2",
+                    ]
+                }
+            }
+        }
+    },
+    "dummy-dataset.pth"
+)
+```
+
+Note that there should not be any newlines within individual sentences! This now raises an error. Instead, each entry of a list should be a sentence, and there should be no "\n" characters. So your corpus should already be well-split.
+
+2. Create/adapt config; provide base model via `model_name_or_path` and training data .pth via `text_path`:
+
+`configs/lora/lora_dummy_config.json`
+
+We recommend starting using this config, and adapting `model_name_or_path`, `output_dir`, and `text_path` if needed.
+You may also wish to adapt other aspects such as `adapter_config` and batch sizes, but this is more experimental.
+
+3. Train LoRA:
+
+```
+python3 wtpsplit/train/train_lora.py configs/lora/lora_dummy_config.json
+```
+
+4. Once training is done, provide your saved module's path to SaT:
+
+```python
+
+sat_lora_adapted = SaT("model-used", lora_path="dummy_lora_path")
+sat_lora_adapted.split("Some domains-specific or styled text")
+```
+
+**Important:** Use the **same model variant** for inference as for training (e.g. `sat-12l-sm` and `sat-12l` have different configs; an adapter trained on one cannot be loaded on the other).
+
+Adjust the dataset name, language and model in the above to your needs.
+
+## Reproducing the paper
+
+`configs/` contains the configs for the runs from the paper for base and sm models as well as LoRA modules. Launch training for each of them like this:
+
+```
+python3 wtpsplit/train/train.py configs/<config_name>.json
+python3 wtpsplit/train/train_sm.py configs/<config_name>.json
+python3 wtpsplit/train/train_lora.py configs/<config_name>.json
+```
+
+In addition:
+
+- `wtpsplit/data_acquisition` contains the code for obtaining evaluation data and raw text from the mC4 corpus.
+- `wtpsplit/evaluation` contains the code for:
+  - evaluation (i.e. sentence segmentation results) via `intrinsic.py`.
+  - short-sequence evaluation (i.e. sentence segmentation results for pairs/k-mers of sentences) via `intrinsic_pairwise.py`.
+  - LLM baseline evaluation (`llm_sentence.py`), legal baseline evaluation (`legal_baselines.py`)
+  - baseline (PySBD, nltk, etc.) evaluation results in `intrinsic_baselines.py` and `intrinsic_baselines_multi.py`
+  - Raw results in JSON format are also in `evaluation_results/`
+  - Statistical significane testing code and results ara in `stat_tests/`
+  - punctuation annotation experiments in `punct_annotation.py` and `punct_annotation_wtp.py` (WtP only)
+  - extrinsic evaluation on Machine Translation in `extrinsic.py` (WtP only)
+
+Ensure to install packages from `requirements.txt` beforehand.
+
+## Supported Languages
+
+<details>
+  <summary>Table with supported languages</summary>
+
+| iso | Name            |
+| :-- | :-------------- |
+| af  | Afrikaans       |
+| am  | Amharic         |
+| ar  | Arabic          |
+| az  | Azerbaijani     |
+| be  | Belarusian      |
+| bg  | Bulgarian       |
+| bn  | Bengali         |
+| ca  | Catalan         |
+| ceb | Cebuano         |
+| cs  | Czech           |
+| cy  | Welsh           |
+| da  | Danish          |
+| de  | German          |
+| el  | Greek           |
+| en  | English         |
+| eo  | Esperanto       |
+| es  | Spanish         |
+| et  | Estonian        |
+| eu  | Basque          |
+| fa  | Persian         |
+| fi  | Finnish         |
+| fr  | French          |
+| fy  | Western Frisian |
+| ga  | Irish           |
+| gd  | Scottish Gaelic |
+| gl  | Galician        |
+| gu  | Gujarati        |
+| ha  | Hausa           |
+| he  | Hebrew          |
+| hi  | Hindi           |
+| hu  | Hungarian       |
+| hy  | Armenian        |
+| id  | Indonesian      |
+| ig  | Igbo            |
+| is  | Icelandic       |
+| it  | Italian         |
+| ja  | Japanese        |
+| jv  | Javanese        |
+| ka  | Georgian        |
+| kk  | Kazakh          |
+| km  | Central Khmer   |
+| kn  | Kannada         |
+| ko  | Korean          |
+| ku  | Kurdish         |
+| ky  | Kirghiz         |
+| la  | Latin           |
+| lt  | Lithuanian      |
+| lv  | Latvian         |
+| mg  | Malagasy        |
+| mk  | Macedonian      |
+| ml  | Malayalam       |
+| mn  | Mongolian       |
+| mr  | Marathi         |
+| ms  | Malay           |
+| mt  | Maltese         |
+| my  | Burmese         |
+| ne  | Nepali          |
+| nl  | Dutch           |
+| no  | Norwegian       |
+| pa  | Panjabi         |
+| pl  | Polish          |
+| ps  | Pushto          |
+| pt  | Portuguese      |
+| ro  | Romanian        |
+| ru  | Russian         |
+| si  | Sinhala         |
+| sk  | Slovak          |
+| sl  | Slovenian       |
+| sq  | Albanian        |
+| sr  | Serbian         |
+| sv  | Swedish         |
+| ta  | Tamil           |
+| te  | Telugu          |
+| tg  | Tajik           |
+| th  | Thai            |
+| tr  | Turkish         |
+| uk  | Ukrainian       |
+| ur  | Urdu            |
+| uz  | Uzbek           |
+| vi  | Vietnamese      |
+| xh  | Xhosa           |
+| yi  | Yiddish         |
+| yo  | Yoruba          |
+| zh  | Chinese         |
+| zu  | Zulu            |
+
+</details>
+
+For details, please see our [Segment any Text paper](https://arxiv.org/abs/2406.16678).
+
+## Community Ports
+
+- **Rust**: [wtsplit-rs](https://github.com/19h/wtsplit-rs) by [@19h](https://github.com/19h)
+
+*Note: Community ports are independently maintained and may have different feature sets or update schedules.*
+
+## Citations
+
+For the `SaT` models, please kindly cite our paper:
+
+```
+@inproceedings{frohmann-etal-2024-segment,
+    title = "Segment Any Text: A Universal Approach for Robust, Efficient and Adaptable Sentence Segmentation",
+    author = "Frohmann, Markus  and
+      Sterner, Igor  and
+      Vuli{\'c}, Ivan  and
+      Minixhofer, Benjamin  and
+      Schedl, Markus",
+    editor = "Al-Onaizan, Yaser  and
+      Bansal, Mohit  and
+      Chen, Yun-Nung",
+    booktitle = "Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing",
+    month = nov,
+    year = "2024",
+    address = "Miami, Florida, USA",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2024.emnlp-main.665",
+    pages = "11908--11941"
+}
+
+```
+
+For the library and the WtP models, please cite:
+
+```
+@inproceedings{minixhofer-etal-2023-wheres,
+    title = "Where{'}s the Point? Self-Supervised Multilingual Punctuation-Agnostic Sentence Segmentation",
+    author = "Minixhofer, Benjamin  and
+      Pfeiffer, Jonas  and
+      Vuli{\'c}, Ivan",
+    booktitle = "Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
+    month = jul,
+    year = "2023",
+    address = "Toronto, Canada",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2023.acl-long.398",
+    pages = "7215--7235"
+}
+```
+
+## Acknowledgments
+
+This research was funded in whole or in part by the Austrian Science Fund (FWF): P36413, P33526, and DFH-23, and by the State of Upper Austria and the Federal Ministry of Education, Science, and Research, through grants LIT-2021-YOU-215. In addition, Ivan Vulic and Benjamin Minixhofer have been supported through the Royal Society University Research Fellowship ‘Inclusive and Sustainable Language Technology for a Truly Multilingual World’ (no 221137) awarded to Ivan Vulić. This research has also been supported with Cloud TPUs from Google’s TPU Research Cloud (TRC). This work was also supported by compute credits from a Cohere For AI Research Grant, these grants are designed to support academic partners conducting research with the goal of releasing scientific artifacts and data for good projects. We also thank Simone Teufel for fruitful discussions.
+
+---
+
+For any questions, please create an issue or send an email to markus.frohmann@gmail.com, and I will get back to you as soon as possible.

{wtpsplit-2.2.0 → wtpsplit-2.2.2}/README.md

@@ -51,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.0 → 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"]