Metadata-Version: 2.4
Name: winnow-sampler
Version: 0.1.0
Summary: Generate and curate synthetic Q&A datasets from a knowledge base for fine-tuning
Author: Automaise AI Lab
License-Expression: MIT
Keywords: llm,fine-tuning,synthetic-data,rag,dataset-curation
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: openai>=1.0
Requires-Dist: python-dotenv
Requires-Dist: pyyaml
Requires-Dist: tqdm
Requires-Dist: numpy
Requires-Dist: faiss-cpu
Provides-Extra: dev
Requires-Dist: pytest>=7; extra == "dev"
Requires-Dist: ruff>=0.4; extra == "dev"
Dynamic: license-file

# 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.

```bash
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](https://huggingface.co/datasets/PrimeQA/TechQA) release
  (CDLA-Permissive-1.0) with `data-sampler-build-techqa`.

## Programmatic API

```python
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

```bash
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](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.
