Metadata-Version: 2.4
Name: tinyzchunk
Version: 0.2.2
Summary: GPU-free chunker distilled from zChunk (EN + PT-BR). Weights fetched from HuggingFace on first use.
Author: Carlo Moro
License: Apache-2.0
Project-URL: Homepage, https://huggingface.co/cnmoro/tinyzchunk
Project-URL: Repository, https://github.com/cnmoro/tinyzchunk
Keywords: chunking,rag,text-splitter,nlp,retrieval
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.20
Requires-Dist: huggingface_hub>=0.20

# tinyzchunk

A **GPU-free, distilled chunker** for RAG pipelines, born from the
[zChunk algorithm](https://github.com/zchunk) (LLM log-probability-based
chunking). The original zChunk needs a large LLM (Llama-3.1-70B) to decide where
to insert semantic split tokens. Here we **distill that teacher into two tiny
MLPs** that read nothing but raw characters — no GPU, no tokenizer, no model
download at runtime. Just numpy.

Specialized for **English and Brazilian Portuguese** (trained on EN + PT-BR
corpora and validated on real-world event-organisation documents), but the
features are character-level and work for any Latin-script text.

## How it works

**Teachers (LLM, run once, need a GPU):**
1. **Logprob teacher** (`scripts/teacher.py`): prompts Qwen2.5-7B to act as a
   "chunker" (repeat text while inserting `段` section splits / `顿` sentence
   splits), then reads the log-probability of those tokens at every position.
   Used to label prose (paragraph/sentence boundaries).
2. **Generation teacher** (`scripts/teacher_gen.py`): asks the LLM to reproduce
   a document inserting `【SPLIT】` markers between self-contained units, then
   aligns the markers back to character positions. Used to label structured
   documents (Q&A roteiros, schedules, bios, FAQs) where the logprob trick is
   unreliable.

**Students (distilled, CPU-only):**
- `tinyzchunk/line_model.py` — a tiny MLP over a **window of line features** that
  predicts, per line, whether a new unit starts there. This is the primary
  boundary detector: it reliably finds Q&A pairs, section headers, schedule
  entries, list items and bios by content alone (question marks, date/time
  fields, title-case headers, sentence completion, blank-line structure).
- `tinyzchunk/model.py` — a tiny char-level MLP used as a fallback to break
  overly-long prose segments at sentence/paragraph boundaries.

Both use the same vectorized feature extractor (`tinyzchunk/features.py`,
~60 content features per character) and run with numpy only.

```
                ┌──────────────────────────┐
  corpus ─────► │ LLM teachers (GPU, once) │──► boundary labels
                └──────────────────────────┘      │
                                                  ▼
                ┌──────────────────────────┐   ┌──────────────────────┐
  text ───────► │ char features (numpy)    │──►│ tiny MLPs (numpy)    │──► chunks
                └──────────────────────────┘   └──────────────────────┘
```

## Install & use

```bash
pip install tinyzchunk
python -m tinyzchunk document.txt
cat document.txt | python -m tinyzchunk
```

```python
from tinyzchunk import Chunker

chunker = Chunker()                       # fetches weights from HF on first use
chunks = chunker.chunk(long_document)     # -> list[str]

# fetch the latest weights explicitly
chunker = Chunker.from_pretrained("cnmoro/tinyzchunk")

# tuning (defaults bias toward under-splitting)
chunker = Chunker(big_threshold=0.7, small_threshold=0.7,
                  max_chunk_chars=2500, min_chunk_chars=100)
```

The pip package is weight-free: the weights are fetched from HuggingFace on
first use (cached), so model downloads are trackable.  To run the training /
distillation scripts instead, install from source with
`pip install -r requirements.txt`.

## Publishing & download tracking

The weights are published as a HuggingFace model repo:
**https://huggingface.co/cnmoro/tinyzchunk**

- HF tracks model downloads automatically (shown on the model page + via the
  `downloads` field of `HfApi().model_info(...)`).
- `Chunker.from_pretrained("cnmoro/tinyzchunk")` fetches the weights through
  `huggingface_hub`, so every `from_pretrained` call increments the counter.
- The package is also publishable to PyPI (`pyproject.toml` included); PyPI
  tracks installs/downloads independently.

## Reproducing the distillation

```bash
# 1. structured + wrapped training corpus (NOT the target docs)
python scripts/build_struct_corpus.py
python scripts/build_wrapped_corpus.py

# 2. label them with the LLM teachers (GPU)
python scripts/teacher.py  --in data/corpus.jsonl --out data/labels/labels.jsonl
python scripts/teacher_gen.py --in data/struct_corpus_prio.jsonl \
       --out data/struct_labels/labels.jsonl --limit 180

# 3. train the tiny students, export weights
python scripts/train.py      # char model
python scripts/train_line.py # line model

# 4. evaluate (real-world docs are held out)
python scripts/eval_heldout.py
```

## Evaluation

The 28 real event-organisation documents (Q&A roteiros, event
schedules, mini-bios, sectioned prose, FAQs, contact lists) are **never used in
training**. Quality is judged manually (`data/heldout_out/*.chunks.txt`) plus an
objective boundary-F1 against a generation-teacher reference.

Best results on the target documents (big_threshold=0.5, small_threshold=0.7):

| document type | boundary F1 vs teacher |
|---|---|
| event schedules | 0.92–0.98 |
| Q&A roteiros | 0.82–1.00 |
| FAQs, sectioned prose | 0.95 / 0.88 |
| contact-block lists | 0.98 |
| dense long-form legal documents | 32 chunks (was 152) |

**Noise robustness.** The chunker is trained on documents with simulated PDF/OCR
noise (narrow mid-word wrapping, page numbers, form feeds, OCR-mangled words) —
`scripts/build_noisy.py`.  On a held-out noisy eval it scores **0.97 boundary
F1**, up from 0.13 without noise training.

**Generalization to other real HuggingFace datasets** (`data/hf_eval_corpus.jsonl`
and PT-BR `Canarim-Instruct`): Q&A pairs in CDC FAQ, sections in PT Wikipedia,
paragraphs in news, sections in FCC regulations, and noisy variants — ~3% of
chunks are fragments.

Speed: ~120–170k chars/s on CPU (≈10–30 ms per document), weights ≈1 MB total,
numpy-only.

Known limits (honest): a couple of short documents under-split, and the
generation-teacher reference is itself inconsistent for some documents (labelled
coarsely, so F1 is a lower bound on actual quality).

For guaranteed-perfect chunking of any document, `scripts/teacher_gen.py`
provides the LLM mode (needs a GPU):
```
python scripts/teacher_gen.py --in doc.jsonl --out labels.jsonl
```

## Layout

```
tinyzchunk/              the library (numpy-only)
  chunker.py             chunk() API; line-model primary + char-model fallback
  features.py            vectorized per-character content features
  line_model.py          line-level unit-start model (numpy inference)
  model.py               char-level boundary model (numpy inference)
  labels.py              teacher labels -> boundary labels
  weights.npz, line_weights.npz   distilled weights (~60KB + ~400KB)
  __main__.py            CLI
scripts/                 corpus builders, teachers, training, evaluation
```
