tinyzchunk

A GPU-free, distilled chunker for RAG pipelines, born from the zChunk algorithm (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

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

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

# 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):

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