Sampling Script

Content-agnostic, topic-erased, judge-free quality bucketing for IFT / pretraining text corpora — with first-class Indic-language support (script routing, code-mixing, romanization, matra/conjunct complexity).

Builds a representation whose variance is driven by how something is written (fluency, richness, structure, repetition) rather than what it's about, clusters that representation into buckets, validates that the buckets track quality rather than topic, and draws a statistically-sized sample per bucket for human review.

This is a package, not a script: every stage is a small, independently swappable module, driven by one YAML config. This README is the full reference — everything the package can currently do, in one place.

• Install

• Quickstart

• What each stage does

• CLI reference

• Config reference — every field, every option

• Feature extractors available

• Discovery methods available

• Topic-erasure methods available

• Embedding backends available

• Input formats available

• Multi-GPU embedding & resume

• Outputs

• Extending qaxis

• Testing

• Updating dependencies

• Caveats

Install

This project uses uv for environment and dependency management — no manual venv, no pip install. uv.lock pins every dependency (including transitive ones) for a reproducible install; uv sync reads it and builds .venv from scratch in seconds.

uv sync                        # core pipeline (TF-IDF fallback) + pytest, ~330MB, no GPU needed
uv sync --extra embeddings     # + sentence-transformers (real style embeddings; pulls in torch)
uv sync --extra gpu            # + a CUDA-enabled torch build
uv sync --all-extras           # everything

uv sync also installs the dev dependency group (pytest) by default — that's a PEP 735 group, not a pip extra, so it never ends up in a built wheel; skip it with uv sync --no-dev for a leaner env.

If you specifically need a CUDA build of torch that doesn't match the default PyPI wheel for your platform, point uv at the right index — see uv's PyTorch integration guide, e.g. uv sync --extra gpu --index pytorch=https://download.pytorch.org/whl/cu124.

No uv? curl -LsSf https://astral.sh/uv/install.sh | sh, or see the install docs.

Quickstart

uv run qaxis init-config --output config.yaml     # a full default config to start editing
uv run qaxis run --input demo_synthia.jsonl --output-dir out

uv run transparently uses .venv (creating/updating it from uv.lock first if needed) — no source .venv/bin/activate required.

uv run qaxis run --config config.yaml --input /data/corpus/ --output-dir out \
                  --checkpoint-dir ./ckpt --devices cuda:0 cuda:1

What each stage does

qaxis run walks through seven stages, every one of them independently configurable and reported on in report.md:

CLI reference

qaxis init-config [--output PATH]
qaxis embed  [--config PATH] [--input PATH] [--checkpoint-dir DIR]
             [--devices DEV [DEV ...]] [--model NAME]
qaxis run    [--config PATH] [--input PATH] [--checkpoint-dir DIR]
             [--devices DEV [DEV ...]] [--model NAME] [--output-dir DIR]

• init-config — writes Config.default() to a YAML file (config.yaml unless --output is given). Start here, then edit the sections you care about.

• embed — pre-computes style embeddings into a resumable checkpoint. Requires --checkpoint-dir (or embedding.checkpoint_dir in the config) and sentence-transformers to be installed (uv sync --extra embeddings); exits with a clear error otherwise rather than silently falling back.

• run — the full pipeline. Requires --input and --output-dir (or the config equivalents).

Every flag other than --config is a convenience override for the handful of things people change per invocation — input.path, embedding.checkpoint_dir, output.dir, embedding.devices, embedding.model. Everything else (which features run, erasure method, cluster sweep, sample sizing, ...) only exists in the YAML config, so a run is fully reproducible from --config alone. Every run also writes run_config.yaml into the output directory with the exact config actually used.

No --config? Both embed and run fall back to Config.default() with your CLI overrides applied on top — you don't need a config file at all for a quick run with default settings.

Config reference

qaxis init-config writes every field below with its default. A config file only needs to contain the keys you want to change — see Partial overrides below.

input

embedding

features

topic

discovery

clustering

sampling

output

top-level

Partial overrides

A config file only needs the keys you're changing — everything else falls back to Config.default():

features:
    enabled: [lexical_diversity, structural]
    include_function_words: false
discovery:
    methods: [pca]

Feature extractors available

Set via config.features.enabled. List them from a shell:

uv run python -c "from qaxis.features.registry import FEATURE_REGISTRY; print(list(FEATURE_REGISTRY))"

Not in this registry (controlled separately, config.features.include_function_words):

Discovery methods available

Set via config.discovery.methods. Both run by default and are cross-checked against each other in the report.

Topic-erasure methods available

Set via config.topic.erasure_method.

Embedding backends available

Set via config.embedding.backend.

Input formats available

Set via config.input.format.

Multi-GPU embedding & resume

For large corpora, split embedding from analysis. Sync the embeddings extra first:

uv sync --extra embeddings
uv run qaxis embed --input /data/corpus/ --checkpoint-dir ./ckpt --devices cuda:0 cuda:1 cuda:2 cuda:3
uv run qaxis run   --input /data/corpus/ --checkpoint-dir ./ckpt --output-dir out

embed writes shard files (shard_00000.npz, ...) keyed by doc_id. Re-running the exact same command — after a crash, a timeout, or to add more GPUs — skips every doc_id already embedded and only fills what's missing. run does the same automatically if pointed at an incomplete checkpoint. Omit --devices to auto-detect every visible CUDA device.

Outputs

<output-dir>/
  report.md              walkthrough of every stage with embedded plots
  plots/*.png
  review_sheet.csv        medoid + LQAS-sized random sample per bucket
  bucket_manifest.csv     bucket sizes, stability, quality scores, sample sizes
  assignments.csv         every record's bucket + quality scores
  run_config.yaml         the exact config used, for reproducibility

report.md walks Stage 0 (script/code-mixing) through Stage 7 (sampling plan), including the topic-erasure before/after probe accuracy and the ARI/AMI topic-vs-quality validation — the two numbers that tell you whether the buckets are actually tracking quality rather than domain.

In review_sheet.csv, only role == "random" rows should feed an accept/reject decision — role == "medoid" rows are for orientation only and are deliberately excluded from the statistical sample.

Extending qaxis

New input format (parquet, a database export, ...): subclass qaxis.io.base.CorpusLoader, implement load(path) -> pd.DataFrame producing the required columns (see that module's docstring), add it to qaxis.io.registry.LOADER_REGISTRY, and set input.format in the config.

New feature extractor: subclass qaxis.features.base.FeatureExtractor, implement compute_row(ctx) -> dict (and fit(df, config) if it needs a corpus-level pass, like the perplexity extractor's per-script language model), add it to qaxis.features.registry.FEATURE_REGISTRY, and list its name in config.features.enabled. It's automatically available to the corruption-contrast discovery stage too, with no extra wiring.

New topic-erasure method (true LEACE, INLP): add a function to qaxis.topic.erasure.ERASURE_METHODS and set topic.erasure_method.

New embedding backend (a real content-controlled style model like CISR/LUAR): subclass qaxis.embeddings.base.EmbeddingBackend, register it in qaxis.embeddings.registry.BACKEND_REGISTRY, set embedding.backend.

New discovery method: add its name to qaxis.discovery.registry.DISCOVERY_METHODS and call it from qaxis.pipeline.run alongside the existing two.

Testing

uv run pytest tests/

pytest comes from the default dev dependency group, so a plain uv sync (or letting uv run sync implicitly) already has it — no extra install step. The smoke tests run the full pipeline on the bundled demo corpus using the TF-IDF/CPU fallback path, so they need no GPU and no optional dependencies:

• test_run_on_single_file / test_run_on_directory — the pipeline completes and produces every expected artifact for both input shapes.

• test_reduced_feature_and_discovery_config — a config enabling only a subset of features and discovery methods still runs, and the report reflects it.

• test_topic_erasure_ablation_changes_result — erasure_method: none retains ~100% of representation variance against topic, residualize much less; if this regresses, erasure broke.

• test_config_round_trip — Config survives a to_yaml / from_yaml cycle unchanged.

Updating dependencies

uv add <package>                          # add a new core dependency
uv add --optional embeddings <package>    # add to an extras group instead
uv add --group dev <package>              # add a dev-only dependency
uv lock --upgrade                          # re-resolve everything to latest compatible versions
uv sync                                    # apply uv.lock to .venv

Caveats

• Judge/LLM score columns, if present in the input, are read exactly once, purely as a post-hoc sanity check after clustering — never as an input to features, erasure, or discovery.

• residualize topic erasure is a one-shot linear approximation of LEACE/INLP. The pipeline validates it with a before/after linear topic-probe; if accuracy doesn't fall close to chance, iterate erasure (or write a stronger method and register it) before trusting the buckets.

• Indic/code-mixing signals (script routing, romanization heuristic, matra/conjunct density) are dependency-free reference implementations. Swap in GlotLID, IndicXlit, and a trained KenLM per language-script for production use.

• Any cluster with bootstrap Jaccard stability < 0.6 is not a reliable bucket — the report flags these; merge, resplit, or route to individual review.

• Human quality judgments have limited value in predicting downstream training utility (DataComp-LM, arXiv:2406.11794). Validate any pruning rule against a real downstream signal before relying on it.