Metadata-Version: 2.4
Name: sol-metric
Version: 0.1.1
Summary: SOL: distributional text evaluation with double-sliced Wasserstein
Maintainer: Gregor Kornhardt
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/GregorKornhardt/SOL
Project-URL: Source, https://github.com/GregorKornhardt/SOL
Project-URL: Issues, https://github.com/GregorKornhardt/SOL/issues
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: numpy>=1.23.5
Provides-Extra: torch
Requires-Dist: torch>=2.2.2; extra == "torch"
Provides-Extra: text
Requires-Dist: torch>=2.2.2; extra == "text"
Requires-Dist: transformers<5,>=4.51.3; extra == "text"
Requires-Dist: packaging>=20; extra == "text"
Provides-Extra: dev
Requires-Dist: build>=1.2; extra == "dev"
Requires-Dist: twine>=5; extra == "dev"
Requires-Dist: ruff<1,>=0.11; extra == "dev"
Requires-Dist: mypy<2,>=1.15; extra == "dev"
Requires-Dist: cffconvert>=2; extra == "dev"
Dynamic: license-file

# SOL - Sliced Optimal Transport for Language

SOL measures the gap between two text
corpora with a double-sliced Wasserstein distance over contextual token
embeddings. Lower is closer; identical corpora score zero.

[Paper](https://arxiv.org/abs/2610.06513) ·
[Website](https://gregorkornhardt.github.io/SOL/) ·
[API](https://github.com/GregorKornhardt/SOL/blob/main/docs/API.md) ·
[Models](https://github.com/GregorKornhardt/SOL/blob/main/docs/MODELS.md) ·
[Changelog](https://github.com/GregorKornhardt/SOL/blob/main/CHANGELOG.md) ·
[Contributing](https://github.com/GregorKornhardt/SOL/blob/main/CONTRIBUTING.md)

## Install

```bash
pip install 'sol-metric[text]'    # compare texts (PyTorch + Transformers)
pip install 'sol-metric[torch]'   # score existing embeddings on a GPU
pip install sol-metric            # score existing embeddings with NumPy only
```

Python 3.10 or newer. For a specific CUDA version, install
[PyTorch](https://pytorch.org/get-started/locally/) first; add `accelerate` to
load large checkpoints straight onto the GPU.

## Use

```python
from sol_metric import SOL

reference = ["The train arrived at noon.", "It rained throughout the night."]
generated = ["The afternoon train was late.", "Heavy rain fell overnight."]
print(SOL().from_texts(reference, generated))
```

The default encoder is Dream 7B, as in the paper: a 15 GB download that needs
a GPU with about 16 GB. Any Hugging Face model can be used instead:

```python
from sol_metric import SOL, TextEncoder

encoder = TextEncoder('gpt2-large')
metric = SOL()
print(metric.from_texts(reference, generated, encoder=encoder))
```

Features can be saved and reused, for example one reference against several
generated corpora:

```python
from tempfile import TemporaryDirectory
from sol_metric import Features

with TemporaryDirectory() as tmp:
    metric.featurize_texts(reference, encoder=encoder, path=f'{tmp}/reference').close()
    with Features.load(f'{tmp}/reference') as ref:
        with metric.featurize_texts(generated, encoder=encoder) as gen:
            print(metric.from_features(ref, gen))
```

Existing token embeddings need only NumPy. Stack the token states of all
documents in one array and mark where each document starts and ends:

```python
import numpy as np
from sol_metric import SOL

rng = np.random.default_rng(0)
x = rng.normal(size=(12, 16))  # corpus 1: 12 token states of dimension 16
x_offsets = [0, 4, 7, 12]      # 3 documents: rows 0-3, 4-6 and 7-11
y = rng.normal(size=(10, 16))  # corpus 2: 10 token states
y_offsets = [0, 5, 10]         # 2 documents: rows 0-4 and 5-9
print(SOL().from_tokens(x, x_offsets, y, y_offsets))
```

Good to know:

- Results are reproducible: the random directions and GP draws are seeded
  (`seed`, `gp_seed`, both 0 by default). Change them to measure the spread
  over seeds, as the paper does with nine seed pairs.
- The paper's encoders (Dream 7B, GPT-2 Large, OLMo 2 1B, Qwen3 0.6B) are pinned
  to the paper's checkpoint commits; pass `revision='main'` for the latest.
- Encoders run in FP16 on a GPU and FP32 on a CPU. Each document is encoded
  whole up to the model's limit (2,048 tokens for Dream, 1,024 for GPT-2 Large);
  longer documents are truncated with a warning.
- A GPU is used when PyTorch sees one; pass `device='cpu'` to force the CPU.
- `metric.last_run` records the score, encoder, revision and all settings.

## Reproduce the paper

The repository bundles LM1B and OpenWebText comparisons, encoded with GPT-2
Large. The 5k OpenWebText pair is the paper's seed-0 subset for MDLM.

```bash
python examples/compare_corpora.py owt                             # 32 texts per corpus
python examples/compare_corpora.py owt --validate --device cuda:0  # 5,000 per corpus
```

| 5k pair | Generated by | Expected score (FP16) |
| --- | --- | ---: |
| LM1B | vMF flow model | 0.961872093 |
| OpenWebText | MDLM, exact sampling | 2.167415507 |

`--validate` re-encodes both corpora and checks the score against the paper's
evaluation code on the same texts; this package matches it to 1e-10. On an RTX
5090 the OpenWebText check takes about 2.5 minutes. The paper's table reports
2.167356 because it used GPT-2 token IDs instead of the decoded text. The 5k
texts (20 MB) are in the Git repository, not the PyPI package; see the
[validation manifest](https://github.com/GregorKornhardt/SOL/blob/main/examples/validation_data/validation.json).

## Settings

| `SOL(...)` | Default | Meaning |
| --- | --- | --- |
| `n_directions`, `n_projections` | 8,192 for Dream, else 1,024 | token directions and GP slices |
| `n_quantiles` | 64 | quantile bins per direction |
| `lengthscale` | 0.1 | RBF length scale of the GP slices |
| `seed`, `gp_seed` | 0, 0 | random directions and GP draws |
| `scale` | hidden dimension | factor on the squared distance |
| `device` | `'auto'` | CUDA when available, else NumPy |

Memory is handled automatically: features stay in RAM while they fit and go
to temporary files otherwise, and batch sizes never change scores; if memory
runs out, the error names the batch size to lower. All settings are in the
[API reference](https://github.com/GregorKornhardt/SOL/blob/main/docs/API.md).

## What SOL computes

Each document is the uniform measure μᵢ on its token states in ℝᵈ, and a corpus
is the uniform distribution over its documents. A unit direction ξ and a
Gaussian-process draw g map every document to one number,
M(μᵢ) = ⟨g, Q((e_ξ)\#μᵢ)⟩, the GP-weighted integral of its projected quantile
function. SOL averages one-dimensional Wasserstein distances over both:

```math
\mathrm{SOL}(\boldsymbol\mu,\boldsymbol\nu)^2
= d\;\mathbb E_{\xi\sim\mathcal U(\mathbb S^{d-1}),\,g\sim\gamma}
\Big[\mathrm W_2^2\big((M_{\xi,g})_\#\boldsymbol\mu,\ (M_{\xi,g})_\#\boldsymbol\nu\big)\Big]
```

with d the hidden dimension and γ a zero-mean GP with RBF covariance. The
estimator uses `n_directions` × `n_projections` slices, exact quantile-bin
integrals and exact transport for unequal corpus sizes.

## Cite

```bibtex
@misc{kornhardt2026solmeasuringgapstext,
      title={SOL: Measuring Gaps between Text Distributions by Double Sliced Wasserstein Metrics},
      author={Gregor Kornhardt and Moritz Piening and Jannis Chemseddine and Gabriele Steidl},
      year={2026},
      eprint={2610.06513},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2610.06513},
}
```

## License

Apache-2.0. Model checkpoints and corpora keep their own licenses.
