data_sampler

Generate a synthetic Q&A dataset from a knowledge base — then throw most of it away.

The premise: a fine-tuning set's quality, not its raw size, drives downstream accuracy. So data_sampler generates aggressively, judges every row, and keeps only what is grounded, non-duplicated and leakage-free.

pip install winnow-sampler          # the package installs as `data_sampler`
data-sampler-build-techqa --techqa_dir /path/to/TechQA   # build the bundled public testbed
data-sampler-run --dataset_key techqa_en --num_new_questions 1500 --num_oos_questions 300

To work on it instead: pip install -e ".[dev]".

Full operating instructions are in USAGE.md, shipped inside the source distribution (pip download --no-binary :all: winnow-sampler).

The pipeline

  KB documents + seed Q&A
          │
   ┌──────▼──────┐  paraphrase · KB question-gen · out-of-scope gen · answer-gen
   │ 1. generate │──────────────────────────────► synthetic_qa.jsonl
   └──────┬──────┘
          │
   ┌──────▼──────┐  retrieve context, then judge: faithfulness, correctness,
   │ 2. score    │  completeness, answerability
   └──────┬──────┘──────────────────────────────► synthetic_qa_scored.jsonl
          │
   ┌──────▼──────┐  keep faithful AND correct (or verified out-of-scope); dedup
   │ 3. curate   │──────────────────────────────► synthetic_qa_curated.jsonl
   └──────┬──────┘
          │
   ┌──────▼──────┐  leakage-safe, OOS-capped train / held-out split
   │ 4. split    │──────────────────────────────► sft_train.jsonl + held_out.jsonl
   └─────────────┘

Three mechanisms shrink the dataset, all deliberate: the quality threshold in stage 3, question dedup in stage 3, and the out-of-scope ratio cap in stage 4.

What it gets right

Questions come from four layers, not one. Seed questions from the testbed, paraphrases of those, fresh questions invented from KB passages (single-fact, broad, and multi-hop spanning two passages), and out-of-scope questions the KB deliberately cannot answer — each generated with the previous layers' questions in a seen set so layers don't duplicate each other.

Refusal is trained, not hoped for. Out-of-scope rows get a fixed IDK answer and are verified by the judge to be genuinely unanswerable from retrieved context before they are kept, so the model learns when to decline rather than only how.

The split can't leak. A paraphrase follows its parent question: if the parent is held out, so is the paraphrase. Without that, a model evaluated on a held-out question would have trained on a reworded copy of it and scored far too well.

Everything resumes. Every stage writes JSONL and re-reads its predecessor's file, and generation caches each question layer separately, fingerprinted with the knobs that produced it — so an interrupted run picks up where it stopped, and a re-run with different knobs rebuilds what those knobs changed instead of silently reusing the old data. Answer generation is the exception: it is not cached, so a re-run re-answers every in-scope question.

Runs are auditable. Each stage drops a manifest recording its knobs and row counts, so you can see where the data shrank and why — including what it lost: failed LLM calls and dropped rows are counted, and generation fails the run rather than writing a truncated dataset when too many calls die. A cost report covers the generation stage's spend (judging is not yet metered).

Requirements

• Python 3.10+
• An Azure OpenAI resource with a chat deployment (generation and judging) and an embeddings deployment (the scoring stage retrieves the context it judges against) — copy .env.example to .env and fill it in. The file is read from the working directory or $DATA_SAMPLER_HOME; real environment variables override it.
• A file-backed knowledge base. techqa_en ships configured; build it from the public TechQA release (CDLA-Permissive-1.0) with data-sampler-build-techqa.

Programmatic API

from data_sampler import DataSampler, GenerationConfig, PipelineConfig, ScoringConfig

config = PipelineConfig(
    generation=GenerationConfig(
        dataset_key="techqa_en",
        paraphrases_per_source_entry=2,
        num_new_questions=1500,
        num_oos_questions=300,
    ),
    scoring=ScoringConfig(dataset_key="techqa_en"),
)

result = DataSampler(config).run()
print(result["split"]["counts"])

Each stage is also callable on its own — .generate(), .score(), .filter_and_dedup(), .split() — and falls back to the previous stage's file on disk when given no rows. To resume: DataSampler(config).run(from_stage="filter").

Adding a knowledge base

1. Add an entry to KB_CONFIGS in src/data_sampler/core/data/config.py with source="local", lang, domain_description and system_prompt.
2. Write a builder emitting kb_documents.jsonl and synthetic_qa_raw.jsonl into experiments/shared/<kb>/data/, following src/data_sampler/cli/build_techqa.py.

Nothing else needs editing — the new key appears in every CLI's --dataset_key choices automatically.

Layout

src/data_sampler/
├── pipeline.py      DataSampler — chains the four stages
├── config.py        stage configs + output-path resolution
├── env.py           workspace .env discovery + credential checks
├── kb.py            which KBs are usable
├── caching.py       generation-stage resumability (fingerprinted)
├── manifest.py      per-stage run manifests
├── generation/      paraphrase / question / answer generation + prompts
├── curation/        judging, filtering, dedup, splitting + judge prompts
├── cli/             one module per command (installed as console scripts)
└── core/            shared building blocks: KB configs, cleaning, chunking,
                     splitting, LLM clients, local FAISS retrieval, I/O

Tests

pytest

No network, no LLM, no KB required — every external boundary is mocked. tests/test_end_to_end.py runs all four stages over a fixture KB (tests/fixtures/mini_kb.py) with only the model and the embedder stubbed, so the chain itself is covered, not just the stages in isolation.

License

MIT — see LICENSE. The TechQA dataset itself is CDLA-Permissive-1.0 and is not redistributed here; data-sampler-build-techqa derives the testbed from your own copy.