wtpsplit

The prior_type determines how the algorithm prefers segment lengths. Use the following guide to choose the best one for your task:

Prior Configuration Details

All priors support an optional lang_code in prior_kwargs to use language-aware defaults (e.g., "zh", "de", "en").

1. Uniform: All lengths are equally good up to max_length. Hard cutoff at the limit.
2. Gaussian: Symmetric bell curve. Requires target_length and spread. Best for embedding models (e.g., targeting 512 chars).
3. Clipped Polynomial: A parabola that clips to zero at ±spread from target_length. More aggressive enforcement than Gaussian.
4. Log-Normal: Asymmetric/Right-skewed. More tolerant of longer segments than shorter ones.

# Prefer ~50 character segments using Gaussian prior
segments = sat.split(
    text,
    max_length=100,
    prior_type="gaussian",
    prior_kwargs={"target_length": 50, "spread": 15}
)

You can adapt SaT to a specific corpus using LoRA with as few as 10-100 segmented training sentences.

1. Setup Environment

Clone the repository and install the required dependencies, including the specific adapters version:

git clone https://github.com/segment-any-text/wtpsplit
cd wtpsplit
pip install -r requirements.txt
pip install adapters==0.2.1 --no-dependencies

2. Prepare Training Data

Create a .pth file containing your training data. The data must be structured as a nested dictionary. Important: Individual sentences must not contain newline characters (\n).

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"
)

3. Configure and Train

Create a configuration JSON file (e.g., configs/lora/lora_dummy_config.json) specifying the model_name_or_path, output_dir, and text_path (the path to your .pth file).

Run the training script:

python3 wtpsplit/train/train_lora.py configs/lora/lora_dummy_config.json

4. Inference with Adapted Model

Load the trained adapter using the lora_path argument in the SaT constructor:

sat_lora_adapted = SaT("model-used", lora_path="dummy_lora_path")
sat_lora_adapted.split("Some domains-specific or styled text")

SaT (Segment Any Text) can be domain- and style-adapted using LoRA (Low-Rank Adaptation). This allows for strong adaptation to specific languages, domains (like legal documents), or styles (like tweets or ASR transcriptions).

To use a pre-trained LoRA module, you must provide both the lang_code and the style_or_domain. Available modules can be found in the <model_repository>/loras folder.

Note: When performing inference, ensure you use the same model variant (e.g., sat-12l) that was used during training, as adapters are not cross-compatible between different model architectures.

# 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 adapt WtP to a custom dataset using the adapt.py script.

1. Prepare Data: Create a .pth file containing your training and test data in the expected nested dictionary format (Language -> Task -> Dataset -> Data).
2. Run Adaptation: Use the wtpsplit/evaluation/adapt.py script.
3. Load Mixture: Use skops to load the generated .skops mixture file into the WtP constructor.

Example Workflow:

# 1. Create dummy data
import torch
torch.save({
    "en": {
        "sentence": {
            "my-dataset": {
                "meta": {"train_data": ["train sentence 1"]},
                "data": ["test sentence 1"]
            }
        }
    }
}, "dummy-dataset.pth")

# 2. Run adaptation via CLI
# python3 wtpsplit/evaluation/adapt.py --model_path=benjamin/wtp-bert-mini --eval_data_path dummy-dataset.pth --include_langs=en

# 3. Load and use the mixture
import skops.io as sio
from wtpsplit import WtP

wtp = WtP(
    "wtp-bert-mini",
    mixtures=sio.load(
        "path/to/mixture.skops",
        ["numpy.float32", "numpy.float64", "sklearn.linear_model._logistic.LogisticRegression"],
    ),
)
wtp.split("your text here", lang_code="en", style="my-dataset")