polaudit 1.0.0__py3-none-any.whl

polaudit/__init__.py

@@ -0,0 +1,39 @@
+"""A-posteriori unimodality (AU) of annotator subgroups, and the language models used to audit it."""
+from .agreement import cohen_kappa, spearman, weighted_kappa
+from .annotation import (AnnotationError, annotate, collect_runs, ensemble, inspect_run, load_model_inputs, progress,
+                         sample_inputs)
+from .datasets import SPECS, get_spec
+from .eda import eda_figures, eda_tables, save_eda
+from .eda_dataset import dataset_tables, eda_overview_all, save_dataset_eda
+from .attribution_bootstrap import attribution_gap, attribution_table, bootstrap_gap
+from .kappa_groups import compute_kappa, save_kappa
+from .prompt_sensitivity import compare_prompts
+from .persona_kappa import compute_persona_kappa, save_persona_kappa
+from .shuffling_robustness import compute_null, save_null
+from .poles import compute_poles, save_poles
+from .model_comparison import compare_models
+from .model_eda import label_share_figure, model_eda, save_model_eda, score_histogram_figure
+from .env import Env
+from .export import csv_ready, sample_per_dimension, save_au_tables
+from .models import DEFAULT_CONFIGS, Model, ModelConfig, ModelError, ModelHub
+from .persona_annotation import (annotate_personas, collect_persona_runs, group_order, inspect_personas, load_persona_inputs,
+                                 persona_long)
+from .pipeline import DatasetResult, clean, load_results, make_tool, run_dataset
+from .prompts import PersonaPrompt, ScoringPrompt, label_from_score, save_prompts
+from .sampling import cover_per_dimension, minimal_cover
+from .scoring import score_personas, score_texts
+from .spec import DatasetSpec
+from .unimodality import PrunedAU, Unimodality, dfu_of_hist
+from .utils import find_root, read_au_csv
+
+__version__ = "1.0.0"
+__all__ = ["Unimodality", "PrunedAU", "dfu_of_hist",
+           "DatasetSpec", "SPECS", "get_spec", "run_dataset", "DatasetResult", "clean", "make_tool",
+           "save_au_tables", "csv_ready", "sample_per_dimension", "minimal_cover", "cover_per_dimension", "find_root", "read_au_csv",
+           "Env", "Model", "ModelConfig", "ModelHub", "ModelError", "DEFAULT_CONFIGS",
+           "ScoringPrompt", "PersonaPrompt", "label_from_score", "score_texts", "score_personas",
+           "annotate", "ensemble", "load_model_inputs", "progress", "inspect_run", "collect_runs", "sample_inputs",
+           "AnnotationError", "cohen_kappa", "weighted_kappa", "spearman",
+           "eda_tables", "eda_figures", "save_eda", "dataset_tables", "save_dataset_eda", "eda_overview_all",
+           "attribution_gap", "attribution_table", "bootstrap_gap", "load_results", "compute_kappa", "save_kappa", "compute_persona_kappa", "save_persona_kappa", "save_prompts",
+           "compute_null", "compare_prompts", "compare_models", "model_eda", "save_model_eda", "score_histogram_figure", "label_share_figure", "save_null", "compute_poles", "save_poles", "annotate_personas", "load_persona_inputs", "persona_long", "collect_persona_runs", "inspect_personas", "group_order"]

polaudit/agreement.py

@@ -0,0 +1,61 @@
+"""How much two sets of ratings of the same texts agree: two prompts, two models, a model and a group of people.
+
+``cohen_kappa`` is for labels (civil / toxic), ``weighted_kappa`` and ``spearman`` for scores. For all of them 1 means perfect
+agreement, 0 means no more agreement than chance would give, and a negative value means less. A pair where one side is missing
+(NaN) is left out. A value that cannot be computed (everything rated the same by both sides, no pairs left) is NaN.
+"""
+from __future__ import annotations
+
+from typing import Iterable, Sequence
+
+import numpy as np
+import pandas as pd
+
+__all__ = ["cohen_kappa", "weighted_kappa", "spearman"]
+
+
+def _pairs(a: Iterable, b: Iterable):
+    a, b = np.asarray(list(a), dtype=object), np.asarray(list(b), dtype=object)
+    if len(a) != len(b):
+        raise ValueError(f"the two sets of ratings must have the same length ({len(a)} and {len(b)})")
+    keep = ~(pd.isna(a) | pd.isna(b))
+    return a[keep], b[keep]
+
+
+def cohen_kappa(a: Iterable, b: Iterable) -> float:
+    """Cohen's kappa of two sets of labels (any values that can be compared, e.g. ``"toxic"``/``"civil"`` or True/False)."""
+    a, b = _pairs(a, b)
+    if len(a) == 0:
+        return float("nan")
+    chance = sum(float(np.mean(a == c)) * float(np.mean(b == c)) for c in set(a) | set(b))
+    return float("nan") if chance == 1 else (float(np.mean(a == b)) - chance) / (1 - chance)
+
+
+def weighted_kappa(a: Iterable, b: Iterable, scale: Sequence[int] = range(1, 6)) -> float:
+    """Quadratic weighted kappa of two sets of scores on ``scale`` (default 1-5): a disagreement of 2 points counts four times as much as
+    one of 1 point. A score outside ``scale`` raises ``ValueError``."""
+    a, b = _pairs(a, b)
+    scale = list(scale)
+    if len(a) == 0 or len(scale) < 2:
+        return float("nan")
+    position = {int(v): i for i, v in enumerate(scale)}
+    try:
+        seen = np.zeros((len(scale), len(scale)))
+        for x, y in zip(a, b):
+            seen[position[int(x)], position[int(y)]] += 1
+    except KeyError as missing:
+        raise ValueError(f"score {missing} is not on the scale {scale[0]}-{scale[-1]}") from None
+    expected = np.outer(seen.sum(axis=1), seen.sum(axis=0)) / seen.sum()
+    steps = np.arange(len(scale))
+    weight = (steps[:, None] - steps[None, :]) ** 2 / (len(scale) - 1) ** 2
+    spread = float((weight * expected).sum())
+    return float("nan") if spread == 0 else 1 - float((weight * seen).sum()) / spread
+
+
+def spearman(a: Iterable, b: Iterable) -> float:
+    """Spearman rank correlation of two sets of scores: do they put the texts in the same order?"""
+    a, b = _pairs(a, b)
+    if len(a) < 2:
+        return float("nan")
+    ra, rb = pd.Series(a.astype(float)).rank(), pd.Series(b.astype(float)).rank()
+    return float("nan") if ra.nunique() < 2 or rb.nunique() < 2 else float(ra.corr(rb))      # one side never changes: undefined

polaudit/annotation.py

@@ -0,0 +1,218 @@
+"""Run ONE model over the texts of the datasets and save every reply (the baseline: no persona).
+
+``annotate(model, inputs, out_dir, datasets=[...])`` is called once per model (and, if wanted, once per dataset), from the
+notebook, so the order is controlled there and a problem shows up where it happens. Everything of a dataset is kept in its
+own folder, ``out_dir/<dataset>/<run>/``::
+
+    <model>.csv         one row per text: dataset, item_id, text, model, run, batch, score, label   (one file per model)
+    all_models.csv      the models side by side, one row per text: <model>_score, <model>_label for every model run so far
+    raw/<model>.raw.jsonl   one line per batch: the prompt, the reply, the parsed scores (the source of truth)
+
+and one line per (model, dataset) pass is added to ``out_dir/run_log.csv``. ``all_models.csv`` is refreshed after every pass.
+If a run stops (an API error, a reply that cannot be parsed) everything finished so far is kept, and calling ``annotate``
+again carries on at the first missing batch.
+"""
+from __future__ import annotations
+
+import json
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Callable, Dict, Iterable, List, Optional
+
+import pandas as pd
+from tqdm.auto import tqdm
+
+from .datasets import SPECS
+from .models import DEFAULT_CONFIGS, Model
+from .prompts import ScoringPrompt
+
+__all__ = ["annotate", "ensemble", "load_model_inputs", "progress", "inspect_run", "collect_runs", "sample_inputs", "AnnotationError"]
+
+ENSEMBLE_FILE = "all_models.csv"
+MODEL_ORDER = list(DEFAULT_CONFIGS)                  # the order of the models in all_models.csv; other names follow alphabetically
+
+LOG_COLUMNS = ["time_utc", "run", "dataset", "model", "provider", "model_id", "temperature", "max_tokens", "prompt_version",
+               "item_word", "batch_size", "texts", "batches", "batches_resumed", "model_calls", "seconds", "status"]
+
+
+class AnnotationError(RuntimeError):
+    """A batch could not be annotated; what was finished before it is saved."""
+
+
+def load_model_inputs(path: Path) -> Dict[str, pd.DataFrame]:
+    """The texts to annotate, per dataset: each text once, in the order of the file (``item_id``, ``text``)."""
+    table = pd.read_csv(path, dtype={"item_id": str})
+    return {dataset: part.drop_duplicates("item_id")[["item_id", "text"]].reset_index(drop=True)
+            for dataset, part in table.groupby("dataset", sort=False)}
+
+
+def _paths(out_dir: Path, dataset: str, run: str, model: str):
+    folder = Path(out_dir) / dataset / run
+    return folder, folder / "raw" / f"{model}.raw.jsonl", folder / f"{model}.csv"
+
+
+def ensemble(out_dir: Path, dataset: str, run: str = "baseline", models: Optional[Iterable[str]] = None) -> pd.DataFrame:
+    """The models of one dataset side by side: one row per text, ``<model>_score`` and ``<model>_label`` for each model that has
+    a saved file (or for ``models``). Written to ``<dataset>/<run>/all_models.csv`` and returned."""
+    folder = Path(out_dir) / dataset / run
+    found = [p.stem for p in folder.glob("*.csv") if p.name != ENSEMBLE_FILE]
+    names = list(models) if models is not None else sorted(found, key=lambda m: (MODEL_ORDER.index(m) if m in MODEL_ORDER else len(MODEL_ORDER), m))
+    merged = None
+    for name in names:
+        part = pd.read_csv(folder / f"{name}.csv", dtype={"item_id": str})
+        part = part[["dataset", "item_id", "text", "score", "label"]].rename(columns={"score": f"{name}_score", "label": f"{name}_label"})
+        merged = part if merged is None else merged.merge(part.drop(columns=["dataset", "text"]), on="item_id", how="outer")
+    if merged is None:
+        return pd.DataFrame()
+    merged.to_csv(folder / ENSEMBLE_FILE, index=False, encoding="utf-8")
+    return merged
+
+
+def _read_batches(raw_path: Path) -> Dict[int, dict]:
+    done: Dict[int, dict] = {}
+    if raw_path.exists():
+        for line in raw_path.read_text(encoding="utf-8").splitlines():
+            record = json.loads(line)
+            if record["status"] == "ok":
+                done[record["batch"]] = record
+    return done
+
+
+def _write_csv(dataset: str, run: str, model: str, done: Dict[int, dict], csv_path: Path) -> None:
+    """The csv is rebuilt from the saved batches only (each keeps its own texts), never from the current input."""
+    rows = [{"dataset": dataset, "item_id": item_id, "text": text, "model": model, "run": run,
+             "score": parsed["score"], "label": parsed["label"]}
+            for _, record in sorted(done.items())
+            for item_id, text, parsed in zip(record["item_ids"], record["texts"], record["parsed"])]
+    pd.DataFrame(rows, columns=["dataset", "item_id", "text", "model", "run", "score", "label"]).to_csv(csv_path, index=False, encoding="utf-8")
+
+
+def _log(out_dir: Path, **fields) -> None:
+    path = Path(out_dir) / "run_log.csv"
+    line = pd.DataFrame([{c: fields.get(c) for c in LOG_COLUMNS}], columns=LOG_COLUMNS)
+    line.to_csv(path, mode="a", header=not path.exists(), index=False, encoding="utf-8")
+
+
+def annotate(model: Model, inputs: Dict[str, pd.DataFrame], out_dir: Path, prompt: Callable[[str, str], str], run: str = "baseline",
+             batch_size: int = 10, parse_retries: int = 2, limit: Optional[int] = None, datasets: Optional[Iterable[str]] = None,
+             item_words: Optional[Dict[str, str]] = None, scale_max: int = 5, toxic_threshold: float = 3,
+             show_progress: bool = True) -> pd.DataFrame:
+    """Score (1-5) and label every text of every dataset with one model, batch by batch. Returns one row per dataset.
+
+    ``limit`` keeps only the first texts of each dataset (a pilot). ``item_words`` overrides the word used for a text in the
+    prompt (default: the word of the dataset, e.g. "comment"). A batch whose reply does not have exactly one score per text is
+    asked again ``parse_retries`` times and then raises ``AnnotationError`` (the reply is kept in the raw file).
+    ``prompt`` is your function ``prompt(texts, item) -> str``: ``texts`` is the batch already numbered ("1. ...\n2. ..."), ``item``
+    the word for a text, and the string it returns is sent to the model. The reply must have a line ``score=<n>, label=<toxic/civil>``
+    per text. Use another ``run`` name when you change the prompt. ``show_progress`` shows a bar per dataset and model (texts done,
+    time elapsed and left, and the batch being asked).
+    """
+    if not callable(prompt):
+        raise TypeError("prompt must be a function prompt(texts, item) -> str")
+    words = {name: spec.item_word for name, spec in SPECS.items()}
+    words.update(item_words or {})
+    summary = []
+    for dataset in (list(datasets) if datasets is not None else list(inputs)):
+        table = inputs[dataset].head(limit) if limit else inputs[dataset]
+        table = table.reset_index(drop=True)
+        scoring = ScoringPrompt(item=words.get(dataset, "text"), scale_max=scale_max, toxic_threshold=toxic_threshold, write=prompt)
+        folder, raw_path, csv_path = _paths(out_dir, dataset, run, model.name)
+        raw_path.parent.mkdir(parents=True, exist_ok=True)
+        done = _read_batches(raw_path)
+        batches = [table.iloc[i:i + batch_size] for i in range(0, len(table), batch_size)]
+        started, calls, resumed, status = time.time(), 0, 0, "complete"
+        bar = tqdm(total=len(table), desc=f"{dataset} / {model.name}", unit="text", disable=not show_progress)
+        try:
+            for number, batch in enumerate(batches, start=1):
+                ids = batch["item_id"].tolist()
+                if number in done:
+                    if done[number]["item_ids"] != ids:
+                        raise AnnotationError(f"{model.name}/{dataset}: saved batch {number} holds other texts than the input; use another run name")
+                    resumed += 1
+                    bar.update(len(batch))
+                    continue
+                first = str(batch["text"].iloc[0]).replace("\n", " ")
+                bar.set_postfix_str(f"batch {number}/{len(batches)}: {first[:40]}...", refresh=True)
+                question = scoring.build(batch["text"].tolist())
+                for attempt in range(1, parse_retries + 2):
+                    reply = model.complete(question)
+                    calls += 1
+                    parsed = scoring.parse(reply)
+                    if len(parsed) == len(batch):
+                        break
+                else:
+                    with open(raw_path, "a", encoding="utf-8") as f:
+                        f.write(json.dumps({"batch": number, "item_ids": ids, "prompt": question, "reply": reply, "attempts": attempt,
+                                            "status": "unparseable"}, ensure_ascii=False) + "\n")
+                    raise AnnotationError(f"{model.name}/{dataset}: batch {number}: expected {len(batch)} scores, received {len(parsed)} "
+                                          f"after {attempt} attempts; the reply is in {raw_path.name}")
+                record = {"batch": number, "item_ids": ids, "texts": batch["text"].tolist(), "prompt": question, "reply": reply,
+                          "parsed": parsed, "attempts": attempt,
+                          "status": "ok", "time_utc": datetime.now(timezone.utc).isoformat(timespec="seconds")}
+                with open(raw_path, "a", encoding="utf-8") as f:
+                    f.write(json.dumps(record, ensure_ascii=False) + "\n")
+                done[number] = record
+                bar.update(len(batch))
+        except Exception as exc:
+            status = f"failed: {type(exc).__name__}: {exc}"[:300]
+            raise
+        finally:
+            bar.close()
+            _write_csv(dataset, run, model.name, done, csv_path)
+            try:
+                ensemble(out_dir, dataset, run)                          # all_models.csv follows the model that was just run
+            except Exception as exc:                                     # never hide the error that stopped the run
+                print(f"warning: all_models.csv of {dataset}/{run} could not be refreshed: {exc!r}")
+            cfg = model.config
+            _log(out_dir, time_utc=datetime.now(timezone.utc).isoformat(timespec="seconds"), run=run, dataset=dataset, model=model.name,
+                 provider=cfg.provider, model_id=cfg.model_id, temperature=cfg.temperature, max_tokens=cfg.max_tokens,
+                 prompt_version=getattr(prompt, "__name__", "custom"), item_word=words.get(dataset, "text"), batch_size=batch_size, texts=len(table),
+                 batches=len(batches), batches_resumed=resumed, model_calls=calls, seconds=round(time.time() - started, 1), status=status)
+        summary.append({"dataset": dataset, "texts": len(table), "batches": len(batches), "resumed": resumed, "model calls": calls,
+                        "seconds": round(time.time() - started, 1)})
+    return pd.DataFrame(summary).set_index("dataset")
+
+
+def progress(inputs: Dict[str, pd.DataFrame], out_dir: Path, models: Iterable[str], run: str = "baseline") -> pd.DataFrame:
+    """For each dataset and model: how many of the texts have a saved score."""
+    rows = {}
+    for dataset, table in inputs.items():
+        row = {"texts": len(table)}
+        for name in models:
+            _, raw_path, _ = _paths(out_dir, dataset, run, name)
+            row[name] = sum(len(r["item_ids"]) for r in _read_batches(raw_path).values())
+        rows[dataset] = row
+    return pd.DataFrame(rows).T
+
+
+def inspect_run(out_dir: Path, model: str, run: str = "baseline") -> pd.DataFrame:
+    """A first look at what one model returned, per dataset: how many texts, the scores it gave, the share it called toxic, and
+    whether any batch had to be asked again or could not be parsed. Meant for checking a model before running the next one."""
+    rows = {}
+    for csv_path in sorted(Path(out_dir).glob(f"*/{run}/{model}.csv")):
+        table = pd.read_csv(csv_path)
+        raw_path = csv_path.parent / "raw" / f"{model}.raw.jsonl"
+        records = [json.loads(line) for line in raw_path.read_text(encoding="utf-8").splitlines()] if raw_path.exists() else []
+        row = {"texts": len(table), "mean score": round(float(table["score"].mean()), 2) if len(table) else None,
+               "toxic %": round(100 * float((table["label"] == "toxic").mean())) if len(table) else None}
+        row.update({f"score {s}": int((table["score"] == s).sum()) for s in range(1, 6)})
+        row["batches asked again"] = sum(1 for r in records if r["status"] == "ok" and r["attempts"] > 1)
+        row["batches not parsed"] = sum(1 for r in records if r["status"] == "unparseable")
+        rows[csv_path.parent.parent.name] = row
+    return pd.DataFrame(rows).T if rows else pd.DataFrame()
+
+
+def sample_inputs(inputs: Dict[str, pd.DataFrame], n: int, seed: int = 42) -> Dict[str, pd.DataFrame]:
+    """``n`` random texts of every dataset (all of them when a dataset has fewer), in the order of the input. The same ``n`` and
+    ``seed`` always give the same texts, so every prompt and every model sees the same subset."""
+    return {name: table.sample(n=min(n, len(table)), random_state=seed).sort_index().reset_index(drop=True)
+            for name, table in inputs.items()}
+
+
+def collect_runs(out_dir: Path, run: str = "baseline") -> pd.DataFrame:
+    """All saved scores of a run: one row per dataset, model and text."""
+    files = sorted(p for p in Path(out_dir).glob(f"*/{run}/*.csv") if p.name != ENSEMBLE_FILE)
+    if not files:
+        return pd.DataFrame(columns=["dataset", "item_id", "text", "model", "run", "score", "label"])
+    return pd.concat([pd.read_csv(f, dtype={"item_id": str}) for f in files], ignore_index=True)

polaudit/attribution_bootstrap.py

@@ -0,0 +1,86 @@
+"""Do non-demographic characteristics attribute a larger share of the testable texts than demographic ones?
+
+Nothing here asks a model for anything, and nothing changes the AU analysis: it reads the final results after the removal of thin groups.
+
+For one dataset the unit is a **polarized text with at least four ratings** (a text needs two groups of two ratings to be compared). For every
+characteristic of the dataset, such a text is
+
+* **testable** when at least two groups have two or more ratings on it (after the removal of thin groups), and
+* **attributed** when it is AU for that characteristic.
+
+The characteristics gender, race and age are *demographic* (D); every other one (politics or ideology, traditionalism, religion, education, occupation) is
+*non-demographic* (N). The pooled attribution rate of a type is attributed / testable summed over its characteristics, and the gap is ``N - D`` in
+percentage points. The texts are resampled with replacement (``reps`` times, fixed ``seed``); in every resample both rates and the gap are computed again from
+the same drawn texts, so the two types move together. The 95% interval is the 2.5th to 97.5th percentile of the gaps.
+"""
+from __future__ import annotations
+
+from typing import Dict, Optional, Sequence
+
+import numpy as np
+import pandas as pd
+
+from .eda import _texts
+from .names import ordered_datasets
+from .pipeline import DatasetResult
+
+__all__ = ["attribution_table", "bootstrap_gap", "attribution_gap"]
+
+DEMOGRAPHIC = ("gender", "race", "age")
+COLUMNS = ["testable D", "attributed D", "testable N", "attributed N"]
+
+
+def attribution_table(result: DatasetResult, demographic: Sequence[str] = DEMOGRAPHIC, min_text_ratings: Optional[int] = None) -> pd.DataFrame:
+    """One row per polarized text with at least ``min_text_ratings`` ratings (default: ``min_ratings * min_groups`` of the dataset, four with the default rules): in how
+    many demographic (D) and non-demographic (N) characteristics it is testable and in how many it is attributed (AU), from the final results after the removal of
+    thin groups. ``demographic`` names the characteristics that count as demographic (every other one is non-demographic)."""
+    texts = _texts(result)
+    if min_text_ratings is None:
+        min_text_ratings = result.tool.min_ratings * result.tool.min_groups
+    keep = texts.index[texts["polarised"] & (texts["ratings"] >= min_text_ratings)]
+    table = pd.DataFrame(0, index=keep, columns=COLUMNS, dtype=int)
+    for dimension, pr in result.pruned.items():
+        kind = "D" if dimension in demographic else "N"
+        rows = pr.result.set_index("item_id")
+        testable = (rows["polarised"] & (rows["kept_groups"].map(len) >= result.tool.min_groups)).reindex(keep, fill_value=False)
+        attributed = rows["au"].reindex(keep, fill_value=False)
+        table[f"testable {kind}"] += testable.to_numpy().astype(int)
+        table[f"attributed {kind}"] += attributed.to_numpy().astype(int)
+    return table
+
+
+def _rates(sums: np.ndarray):
+    """The pooled rates (D, N) of a vector of the four sums; NaN when a type has nothing testable."""
+    with np.errstate(divide="ignore", invalid="ignore"):
+        return sums[1] / sums[0], sums[3] / sums[2]
+
+
+def bootstrap_gap(table: pd.DataFrame, reps: int = 10_000, seed: int = 42) -> dict:
+    """The observed gap ``N - D`` of the pooled attribution rates (percentage points), its 95% percentile interval over ``reps`` resamples of the texts, and
+    the share of the resamples in which the gap is zero or negative."""
+    values = table[COLUMNS].to_numpy(dtype=float)
+    n = len(values)
+    rate_d, rate_n = _rates(values.sum(axis=0))
+    rng = np.random.default_rng(seed)
+    gaps = np.empty(reps)
+    for b in range(reps):
+        counts = np.bincount(rng.integers(0, n, n), minlength=n)                  # how often each text is drawn
+        d, m = _rates(counts @ values)
+        gaps[b] = 100 * (m - d)
+    finite = gaps[np.isfinite(gaps)]
+    return {"texts": n, "testable D": int(values[:, 0].sum()), "attributed D": int(values[:, 1].sum()), "rate D %": 100 * rate_d,
+            "testable N": int(values[:, 2].sum()), "attributed N": int(values[:, 3].sum()), "rate N %": 100 * rate_n,
+            "N - D (pp)": 100 * (rate_n - rate_d), "95% interval low": float(np.percentile(finite, 2.5)), "95% interval high": float(np.percentile(finite, 97.5)),
+            "share of resamples with N - D <= 0 (%)": 100 * float(np.mean(finite <= 0)), "resamples": reps, "resamples without a rate": int(reps - len(finite))}
+
+
+def attribution_gap(results: Dict[str, DatasetResult], reps: int = 10_000, seed: int = 42, demographic: Sequence[str] = DEMOGRAPHIC,
+                    min_text_ratings: Optional[int] = None) -> pd.DataFrame:
+    """``bootstrap_gap`` for every dataset (in the order Attitudes, POPQUORN, MHS, Kumar). Every dataset starts from the same ``seed``, so its result does
+    not depend on the other datasets."""
+    rows = {name: bootstrap_gap(attribution_table(results[name], demographic, min_text_ratings), reps, seed) for name in ordered_datasets(results)}
+    table = pd.DataFrame.from_dict(rows, orient="index").rename_axis("dataset")
+    for column in ("rate D %", "rate N %", "N - D (pp)", "95% interval low", "95% interval high"):
+        table[column] = table[column].round(2)
+    table["share of resamples with N - D <= 0 (%)"] = table["share of resamples with N - D <= 0 (%)"].round(2)
+    return table

polaudit/datasets/__init__.py

@@ -0,0 +1,13 @@
+"""The datasets: one module each, with its loader and the definition of its annotator groups."""
+from . import attitudes, kumar, mhs, popquorn
+
+SPECS = {spec.name: spec for spec in (popquorn.SPEC, attitudes.SPEC, mhs.SPEC, kumar.SPEC)}
+
+
+def get_spec(name: str):
+    if name not in SPECS:
+        raise KeyError(f"unknown dataset {name!r}; known: {sorted(SPECS)}")
+    return SPECS[name]
+
+
+__all__ = ["SPECS", "get_spec", "popquorn", "attitudes", "mhs", "kumar"]

polaudit/datasets/attitudes.py

@@ -0,0 +1,69 @@
+"""Attitudes: 1-5 ratings of 626 tweets (`toyou`); the annotator attributes are lists parallel to the ratings."""
+from __future__ import annotations
+
+import ast
+import os
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+from ..spec import DatasetSpec
+from .groups import age_generations
+
+PARALLEL = {"rating": "toyou", "gender": "annotatorGender", "race": "annotatorRace", "age": "annotatorAge",
+            "politics": "annotatorPolitics", "traditionalism": "traditionalism"}
+
+
+def _as_list(x):
+    if isinstance(x, str):
+        try:
+            x = ast.literal_eval(x)
+        except Exception:
+            return []
+    return list(x) if isinstance(x, (list, tuple, np.ndarray)) else []
+
+
+def load(root: Path, path: Optional[Path] = None) -> pd.DataFrame:
+    """One row per rating. ``path`` (or ``POLAUDIT_ATTITUDES_FILE``) overrides the default location.
+    NOTE: despite its .csv extension, the file is a pickled DataFrame."""
+    path = Path(path or os.environ.get("POLAUDIT_ATTITUDES_FILE") or os.environ.get("POLALIGN_ATTITUDES_FILE") or Path(root) / "Thesis" / "Code" / "attitudes_embedded.csv")
+    raw = pd.read_pickle(path)
+    lists = {name: raw[col].map(_as_list) for name, col in PARALLEL.items()}
+    lengths = pd.DataFrame({name: s.map(len) for name, s in lists.items()})
+    if not ((lengths.nunique(axis=1) == 1).all() and (lengths["rating"] > 0).all()):
+        raise ValueError("ratings and annotator attributes do not line up in the attitudes file")
+    rows = [(idx, raw.at[idx, "tweet"]) + tuple(lists[name].at[idx][k] for name in PARALLEL)
+            for idx in raw.index for k in range(lengths.at[idx, "rating"])]
+    return pd.DataFrame(rows, columns=["item_id", "text"] + list(PARALLEL))
+
+
+def _more_than_three_ratings(table: pd.DataFrame) -> pd.DataFrame:
+    """Texts with only 3 annotations are removed."""
+    return table[table.groupby("item_id")["rating"].transform("size") > 3]
+
+
+def _politics_group(score: pd.Series) -> pd.Series:
+    return pd.cut(pd.to_numeric(score, errors="coerce"), [-10, -0.25, 0.25, 10], labels=["left", "center", "right"]).astype(object)
+
+
+def _traditionalism_group(score: pd.Series) -> pd.Series:
+    return pd.cut(pd.to_numeric(score, errors="coerce"), [0, 3, 5], labels=["not_high", "high"]).astype(object)
+
+
+SPEC = DatasetSpec(
+    name="attitudes",
+    load=load,
+    dimensions={
+        "gender": ("gender", {"man": "man", "woman": "woman"}),
+        "race": ("race", {"white": "white", "black": "black"}),
+        "age": ("age", age_generations),
+        "politics": ("politics", _politics_group),
+        "traditionalism": ("traditionalism", _traditionalism_group),
+    },
+    prefilter=_more_than_three_ratings,
+    item_word="tweet",
+    description="Attitudes, 1-5 `toyou` ratings",
+    rating_column="toyou",                       # no annotator id in the data
+)

polaudit/datasets/groups.py

@@ -0,0 +1,13 @@
+"""Group definitions shared by several datasets."""
+from __future__ import annotations
+
+import pandas as pd
+
+__all__ = ["age_generations"]
+
+
+def age_generations(ages: pd.Series) -> pd.Series:
+    """Age in years -> ``gen z`` (under 25), ``millenial`` (25-44), ``gen x+`` (45 and over). Missing stays missing.
+    (The spelling ``millenial`` is the one used in all the results.)"""
+    return pd.cut(pd.to_numeric(ages, errors="coerce"), [0, 25, 45, 200], right=False,
+                  labels=["gen z", "millenial", "gen x+"]).astype(object)

polaudit/datasets/kumar.py

@@ -0,0 +1,42 @@
+"""Kumar et al.: 0-4 toxicity scores (shifted to 1-5) of 107,620 comments, 5 raters each, with the rater's attributes."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pandas as pd
+
+from ..spec import DatasetSpec
+
+COLUMNS = ["toxic_score", "gender", "race", "age_range", "political_affilation", "religion_important"]
+
+
+def load(root: Path) -> pd.DataFrame:
+    """The file is JSON lines: one comment per line with its list of ratings (about 450 MB, read line by line)."""
+    rows = []
+    with open(Path(root) / "datasets" / "kumar" / "toxicity_ratings.json", encoding="utf-8") as f:
+        for line in f:
+            comment = json.loads(line)
+            for rating in comment["ratings"]:
+                rows.append((comment["comment_id"], comment["comment"]) + tuple(rating.get(c) for c in COLUMNS))
+    out = pd.DataFrame(rows, columns=["item_id", "text"] + COLUMNS)
+    out["rating"] = out["toxic_score"] + 1                                  # 0-4 -> 1-5
+    return out
+
+
+SPEC = DatasetSpec(
+    name="kumar",
+    load=load,
+    dimensions={
+        "gender": ("gender", {"Male": "Male", "Female": "Female"}),
+        "race": ("race", {"White": "White", "Black or African American": "Black", "Asian": "Asian"}),
+        "age": ("age_range", {"Under 18": "gen z", "18 - 24": "gen z", "25 - 34": "millenial", "35 - 44": "millenial",
+                              "45 - 54": "gen x+", "55 - 64": "gen x+", "65 or older": "gen x+"}),
+        "politics": ("political_affilation", {"Liberal": "liberal", "Independent": "independent", "Conservative": "conservative"}),
+        "religion": ("religion_important", {"Not important": "not_high", "Not too important": "not_high",
+                                            "Somewhat important": "not_high", "Very important": "high"}),
+    },
+    item_word="comment",
+    description="Kumar, toxic_score + 1 (1-5)",
+    rating_column="toxic_score",                 # no rater id in the data
+)

polaudit/datasets/mhs.py

@@ -0,0 +1,47 @@
+"""Measuring Hate Speech (Hugging Face `ucberkeley-dlab/measuring-hate-speech`): the `hatespeech` item, No / Unclear / Yes -> 1-3."""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import pandas as pd
+
+from ..spec import DatasetSpec
+from .groups import age_generations
+
+URL = "https://huggingface.co/datasets/ucberkeley-dlab/measuring-hate-speech/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet"
+COLUMNS = ["comment_id", "annotator_id", "text", "hatespeech", "annotator_gender", "annotator_age", "annotator_ideology"]
+
+
+def load(root: Path, path: Optional[Path] = None) -> pd.DataFrame:
+    """Reads ``datasets/mhs/train.parquet`` when it exists (or ``path``), otherwise downloads the file from Hugging Face."""
+    local = Path(path) if path else Path(root) / "datasets" / "mhs" / "train.parquet"
+    raw = pd.read_parquet(local if local.exists() else URL, columns=COLUMNS)
+    out = raw.rename(columns={"comment_id": "item_id"})
+    out["rating"] = out["hatespeech"].astype(int) + 1                       # 1 = No, 2 = Unclear, 3 = Yes
+    return out
+
+
+def _at_least_four_ratings(table: pd.DataFrame) -> pd.DataFrame:
+    """A comment needs 4 ratings to have 2 + 2; the others cannot be AU."""
+    return table[table.groupby("item_id")["rating"].transform("size") >= 4]
+
+
+IDEOLOGY = {"extremely_liberal": "liberal", "liberal": "liberal", "slightly_liberal": "liberal", "neutral": "neutral",
+            "slightly_conservative": "conservative", "conservative": "conservative", "extremely_conservative": "conservative"}
+
+SPEC = DatasetSpec(
+    name="mhs",
+    load=load,
+    dimensions={
+        "gender": ("annotator_gender", {"male": "male", "female": "female"}),
+        "age": ("annotator_age", age_generations),
+        "ideology": ("annotator_ideology", IDEOLOGY),
+    },
+    scale_max=3,
+    prefilter=_at_least_four_ratings,
+    item_word="comment",
+    description="Measuring Hate Speech, `hatespeech` item (1 = No, 2 = Unclear, 3 = Yes)",
+    annotator_column="annotator_id",
+    rating_column="hatespeech",
+)

polaudit/datasets/popquorn.py

@@ -0,0 +1,39 @@
+"""POPQUORN: 1-5 offensiveness ratings of 1,500 texts, with the rater's gender, race, age band, education and occupation."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+
+from ..spec import DatasetSpec
+
+# the ten age bands of the file, grouped like the generations of DICES
+AGE_BANDS = {"18-24": "gen z",
+             "25-29": "millenial", "30-34": "millenial", "35-39": "millenial", "40-44": "millenial",
+             "45-49": "gen x+", "50-54": "gen x+", "54-59": "gen x+", "60-64": "gen x+", ">65": "gen x+"}
+
+
+def load(root: Path) -> pd.DataFrame:
+    raw = pd.read_csv(Path(root) / "datasets" / "popquorn" / "popquorn.csv")
+    return raw.rename(columns={"instance_id": "item_id", "offensiveness": "rating"})
+
+
+def _keep(*values):
+    return {v: v for v in values}               # keep these values as they are; every other value is removed
+
+
+SPEC = DatasetSpec(
+    name="popquorn",
+    load=load,
+    dimensions={
+        "gender": ("gender", _keep("Man", "Woman")),
+        "race": ("race", _keep("White", "Black or African American", "Asian")),
+        "age": ("age", AGE_BANDS),
+        "education": ("education", _keep("College degree", "High school diploma or equivalent", "Graduate degree")),
+        "occupation": ("occupation", _keep("Employed", "Self-employed", "Unemployed", "Retired", "Student", "Homemaker")),
+    },
+    item_word="text",
+    description="POPQUORN, 1-5 offensiveness",
+    annotator_column="user_id",
+    rating_column="offensiveness",
+)