polartox

NLP toolkit for annotator polarization research. Provides tools for synthetic annotation-data generation, Polarized Trees analysis, and systematic hyperparameter benchmarking.

Project website — overview, the four modules, and nDFU explained.

Install

pip install polartox

Tools

polartox.datagen

Builds a pool of annotators with explicit demographic identities and generates annotation datasets where every text independently gets k active dimensions that drive its disagreement:

• k = 0 — no demographic dimension explains the disagreement, providing a true unimodal negative control
• k ≥ 1 — a subset of dimensions is active, each with its own random toxic/civil lean split and a continuous intensity (alpha) controlling how strongly it pulls toward its pole

Identities' rating distributions are built by taking the elementwise product of their active-dimension shapes — signal composes rather than averages away, allowing the generated data to span the full nDFU range.

from polartox.datagen import (
    AnnotatorPool,
    DEFAULT_DIMENSIONS,
    DEFAULT_DEPTH_WEIGHTS,
    DEFAULT_INTENSITY_RANGE,
)

pool = AnnotatorPool(
    dimensions=DEFAULT_DIMENSIONS,
    scale=5,
    intensity_range=DEFAULT_INTENSITY_RANGE,
    depth_weights=DEFAULT_DEPTH_WEIGHTS,
    annotators_per_identity=10,
)

result = pool.generate_dataset(
    n_texts=100,
    n_annotators_per_text=150,
    noise=0.05,
    seed=42,
)

dataset, ground_truth = result

# dataset columns:
# text_id, annotator_id, <dimensions>, rating

# ground_truth:
# per-text active_dims, lean, and alpha

nDFU scoring is provided by the collaborative ndfu package (Pavlopoulos & Likas, 2024) rather than reimplemented here:

from ndfu import dfu, pdf

text_data = dataset[dataset["text_id"] == 0]
hist = pdf(
    text_data["rating"].tolist(),
    range(1, pool.scale + 1),
)
score = dfu(hist)

polartox.polarized_tree and polartox.pipeline

Given an annotation dataset, recursively partitions annotators by demographic dimension to identify the dimensions and intersectional subgroups associated with polarized opinions. polartox.polarized_tree (PolarizedTree) builds and inspects one text's tree; polartox.pipeline (PolarizedTreesPipeline) filters a corpus, builds a tree per text, and aggregates corpus-level metrics across them.

from polartox.pipeline import PolarizedTreesPipeline

pipe = PolarizedTreesPipeline(
    dims=list(DEFAULT_DIMENSIONS.keys()),
    scale=5,
    theta_filter=0.3,
    min_size_frac=0.03,
    h=0.15,
    max_depth=6,
    relative_h=True,
)

results = pipe.run_full_evaluation(dataset)

# results["F"], results["C"], results["P"]
# dataset-level polarization summaries

# results["diagnostics"]
# ground-truth-free diagnostics

PEG formulations (variant)

At every node the tree splits on the dimension with the highest PEG (Polarization Explanation Gain): how much of the node's polarization (its nDFU) a split explains. PEG compares the node's nDFU with its subgroups' nDFU. Three base formulations differ in which subgroup they compare against, and two combine them:

beta (default 1) only matters for harmonic: it weights avg against max and min with weights (1, beta^2, 1), so beta=1 is the plain harmonic mean. An unknown variant raises ValueError immediately. The older names beta (the harmonic mean of max and avg only) and var (now avg) no longer exist.

To compare configurations, for example the same settings with different PEG formulations, pairwise_ari(pipelines, dataset) returns the mean adjusted Rand index between their trees (1 = the same groups of annotators, about 0 = chance).

When synthetic ground truth is available, recovery metrics can also be computed:

results = pipe.run_full_evaluation(
    dataset,
    ground_truth=ground_truth,
)

results["recovery"]
# per-text Jaccard, precision, recall, and exact match

polartox.benchmark

PolarizedTreesBenchmark provides systematic configuration search and model selection when ground truth is available.

The benchmark:

1. takes a PolarizedTreesPipeline, annotations, and ground truth;
2. generates valid configurations from a search space;
3. evaluates each configuration using the requested recovery metrics;
4. ranks configurations according to a selected metric;
5. returns the best configuration, best pipeline, top configurations, and complete results;
6. can save benchmark results and a JSON report.

The default search space contains 3,240 valid configurations: 648 settings of the other hyperparameters times the five PEG formulations (max, avg, min, mean, harmonic). beta is not searched; it only weights avg inside harmonic (default 1.0) and can be added to a custom search space like any other parameter.

The benchmark is fully configurable. Users can provide their own search space, choose full or random search, set the number of sampled configurations and seed, select recovery metrics, and specify the metric and direction used for model selection. Long searches can run in parallel (n_jobs), save progress and resume (checkpoint_dir, checkpoint_every), and evaluate several corpora at once (text_groups, with metrics averaged over corpora). Every pipeline is built before the first evaluation, so an invalid setting fails immediately.

A simple example:

from polartox.benchmark import PolarizedTreesBenchmark

benchmark = PolarizedTreesBenchmark(
    pipeline=pipe,
    annotations=dataset,
    ground_truth=ground_truth,
    strategy="random",
    n_runs=100,
    seed=42,
    metrics=["jaccard", "precision", "recall"],
    selection_metric="jaccard",
)

benchmark.run()

best_config = benchmark.get_best_config()
best_score = benchmark.get_best_score()
top_configs = benchmark.get_top_configs()
best_pipeline = benchmark.get_best_pipeline()

The benchmark can also be used with a custom search space:

custom_search_space = {
    "theta_filter": [0.2, 0.3],
    "min_size_frac": [0.03, 0.05],
    "max_depth": [4, 6],
    "variant": ["harmonic"],
    "h": [0.10, 0.15],
    "relative_h": [True],
    "theta_stop": [0.10, 0.15],
}

benchmark = PolarizedTreesBenchmark(
    pipeline=pipe,
    annotations=dataset,
    ground_truth=ground_truth,
    search_space=custom_search_space,
    strategy="full",
)

The selected configuration can subsequently be applied to new annotation data without ground truth. In this inference setting, Polarized Trees reports F, C, and P together with the associated diagnostics.

polartox.peg_comparison

PEGComparison compares the PEG formulations (max, avg, min, mean, harmonic) on the same corpora, everything else held fixed: recovery and tree shape per formulation, and the adjusted Rand index between the trees they build. Ground truth is required; the settings can come from PolarizedTreesBenchmark.top_configs_.

from polartox import PEGComparison

comparison = PEGComparison.from_benchmark(benchmark)    # corpora, dims, scale
comparison.run(benchmark.top_configs_.head(3))          # each row with every formulation
comparison.overview()                                   # recovery and tree shape
comparison.ari_matrix()                                 # do they split annotators the same way?

Reproducibility

The repository includes a complete synthetic benchmark workflow under benchmarks/.

The workflow is:

Synthetic data generation
          ↓
Fixed datasets + ground truth
          ↓
Configuration search
          ↓
Recovery-based model selection
          ↓
Selected Polarized Trees pipeline
          ↓
Inference on unseen data

The treesbenchmark.ipynb notebook serves both as a runnable demonstration of PolarizedTreesBenchmark and as the experimental workflow used to obtain the configurations and results reported in the paper.

The benchmark documentation, configuration space, recovery results, and unseen-corpus evaluation are described in benchmarks/README.md.

Notebooks

End-to-end, on real annotation data:

• DICES_polarized_trees_end_to_end.ipynb — the complete workflow on the DICES-350/990 safety datasets: filtering, tree construction, F/C/P, and inference without ground truth.

Component demos, one workflow step at a time (synthetic data):

• data_gen/datagen_demo.ipynb — generate synthetic annotation data with known polarization ground truth (polartox.datagen).
• polarized_tree/polarized_tree_demo.ipynb — build and inspect one text's tree directly (PolarizedTree), without a pipeline or a corpus.
• polarized_trees/trees_demo.ipynb — run the full pipeline over a corpus and compute F/C/P (PolarizedTreesPipeline).
• peg_comparison/peg_comparison_demo.ipynb — compare the PEG formulations on the same corpora (PEGComparison).
• benchmarks/notebooks/datasetdemo.ipynb — generate the fixed synthetic benchmark corpora.
• benchmarks/notebooks/treesbenchmark.ipynb — hyperparameter search, model selection, and paper reproducibility (PolarizedTreesBenchmark).

Documentation and Paper

Full API documentation and demos are available on GitHub.

For the full methodology and paper-specific details, see polarized_trees.pdf.

Changelog

See the CHANGELOG for release history.