Metadata-Version: 2.4
Name: snipex
Version: 0.1.0
Summary: Single-GPU, oNe-week Idea-to-Prototype EXecution.
Author: Seria
License-Expression: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: THIRD_PARTY_NOTICES.md
Provides-Extra: llm
Requires-Dist: numpy>=1.26; extra == "llm"
Requires-Dist: pyarrow>=21; extra == "llm"
Requires-Dist: rustbpe>=0.1; extra == "llm"
Requires-Dist: tcurve>=0.2.3; extra == "llm"
Requires-Dist: tiktoken>=0.11; extra == "llm"
Requires-Dist: torch<3,>=2.9; extra == "llm"
Dynamic: license-file

<p align="center">
  <img src="https://i.ibb.co/tM18v4Wn/snipe-icon.png" alt="SnipeX logo" width="160">
</p>

<h1 align="center">SnipeX</h1>

<p align="center"><strong>Single-GPU, oNe-week Idea-to-Prototype EXecution.</strong></p>

<p align="center">
  <a href="https://pypi.org/project/snipex/"><img src="https://img.shields.io/pypi/v/snipex" alt="PyPI version"></a>
  <a href="https://github.com/SeriaQ/SnipeX"><img src="https://img.shields.io/github/stars/SeriaQ/SnipeX" alt="GitHub stars"></a>
  <img src="https://img.shields.io/pypi/pyversions/snipex" alt="Python versions">
  <a href="https://github.com/SeriaQ/SnipeX/blob/main/LICENSE"><img src="https://img.shields.io/pypi/l/snipex" alt="MIT license"></a>
</p>

SnipeX is a collection of small AI algorithm prototypes. Algorithm-specific
dependencies are provided through optional package extras.
Every algorithm prototype in SnipeX is designed to run on a single consumer
GPU and reach a working result within one week.

---

## ⚡ Quick start

Installing the base package provides the `snipex` command, package version, and
help for the available algorithm families without installing PyTorch:

```bash
pip install snipex
snipex --help
```

---

## 🧠 LLM prototype (V0.1.0)

The first prototype is a NanoChat-inspired decoder-only language model with
RoPE, RMSNorm and QK norm, ReLU² MLPs, document-aware ClimbMix packing,
AdamW with linear warmup and cosine decay, MPS/CUDA support, and CUDA DDP.

Install the LLM prototype from PyPI:

```bash
pip install "snipex[llm]"
```

From a source checkout, create the development environment with:

```bash
uv sync --extra llm
```

FA3 is optional and CUDA-only. `auto` falls back to PyTorch SDPA when the
community kernel cannot be loaded. It is an implementation optimization, not
an algorithm extra, so source checkouts install it through a dependency group:

```bash
uv sync --extra llm --group flash
```

PyPI users who explicitly want FA3 can install `snipex[llm]` and `kernels`
separately. The normal `snipex[llm]` installation uses SDPA without `kernels`.

### 📚 Data and tokenizer

Keep every split in a separate directory. SnipeX never infers a split from a
filename and never reads the test directory during tokenizer or model training:

```text
/path/to/climbmix/
├── train/*.{parquet,jsonl}
├── val/*.{parquet,jsonl}
└── test/*.{parquet,jsonl}
```

Base-model JSONL uses one document per line:

```json
{"text": "One pretraining document."}
```

The equivalent Parquet schema has one string column named `text`. Mixing
Parquet and JSONL files inside one split is supported.

Train and inspect the 32K RustBPE tokenizer:

```bash
uv run snipex tok train \
  --train-data /path/to/climbmix/train \
  --out-dir artifacts/tokenizer

uv run snipex tok encode --tokenizer artifacts/tokenizer "Hello SnipeX"
uv run snipex tok decode --tokenizer artifacts/tokenizer 123 456
uv run snipex tok eval --tokenizer artifacts/tokenizer "Hello SnipeX"
uv run snipex tok eval \
  --tokenizer artifacts/tokenizer \
  --val-data /path/to/climbmix/val
```

### 🚀 Pretraining

Use a small model for an MPS smoke test:

```bash
uv run snipex llm pretrain \
  --train-data /path/to/climbmix/train \
  --val-data /path/to/climbmix/val \
  --tokenizer artifacts/tokenizer \
  --output runs/mps-smoke \
  --device mps --compile no --attention sdpa \
  --depth 2 --sequence-length 128 --device-batch-size 2 --steps 10
```

Use one CUDA GPU by running the command normally:

```bash
uv run snipex llm pretrain \
  --train-data /path/to/climbmix/train \
  --val-data /path/to/climbmix/val \
  --tokenizer artifacts/tokenizer \
  --output runs/base-single \
  --device cuda --depth 12 --sequence-length 512 \
  --device-batch-size 8 --total-batch-tokens 131072
```

Use two GPUs with DDP through `torchrun`:

```bash
uv run torchrun --standalone --nproc-per-node=2 -m snipex.cli llm pretrain \
  --train-data /path/to/climbmix/train \
  --val-data /path/to/climbmix/val \
  --tokenizer artifacts/tokenizer \
  --output runs/base-dual \
  --device cuda --depth 12 --sequence-length 512 \
  --device-batch-size 8 --total-batch-tokens 131072
```

`--device-batch-size` is sequences per GPU. `--total-batch-tokens` is the desired
global token count per optimizer step across every GPU and accumulation step;
SnipeX reports the derived accumulation count and rejects non-divisible values.
Omit both total batch and accumulation options to use one micro-batch per step.
Omit `--steps` to derive the run length from the parameter/data ratio.
Each data option names a directory; SnipeX recursively reads every supported
file in that directory. During training, `train` updates the model and `val`
provides periodic and final validation metrics. The `test` directory is not read.

Training writes `run.json`, TCurve CSV/plots under `metrics/`, periodic
checkpoints when requested, and an always-present `final.pt`. Resume with:

```bash
uv run snipex llm pretrain <same-data-and-output-options> \
  --resume runs/base-dual/step_001000.pt
```

Evaluate or generate from a base checkpoint:

```bash
uv run snipex llm eval \
  --checkpoint runs/base-dual/final.pt \
  --test-data /path/to/climbmix/test \
  --tokenizer artifacts/tokenizer

uv run snipex llm gen \
  --checkpoint runs/base-dual/final.pt \
  --tokenizer artifacts/tokenizer \
  --prompt "The purpose of a prototype is"
```

### 💬 SFT and chat

SFT accepts Parquet or JSONL with the standard `messages` structure used by
MS-SWIFT. JSONL contains one conversation per line:

```json
{"messages": [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi"}]}
```

Only assistant content and its end token contribute to the loss. V0.1.0
`data-prep` does not convert Alpaca, query/response, ShareGPT, tool, or
multimodal schemas.

Normalize one directory of Parquet files at a time. `data-prep` does not infer
repository layouts or splits; every source Parquet is converted to an output
Parquet with the same filename:

```bash
uv run snipex llm data-prep \
  --dataset mmlu \
  --input /path/to/raw-mmlu-train \
  --output sft-data/train \
  --tokenizer artifacts/tokenizer \
  --sequence-length 512

uv run snipex llm data-prep \
  --dataset mmlu \
  --input /path/to/raw-mmlu-val \
  --output sft-data/val \
  --tokenizer artifacts/tokenizer \
  --sequence-length 512
```

`--tokenizer` and `--sequence-length` must be supplied together. With both,
rows whose complete rendered conversation exceeds the model window are
filtered. With neither, `data-prep` only normalizes the schema and SFT will
reject any over-length row rather than truncate it.

The row converters are public Python functions. A custom converter with the
same `row -> {"messages": [...]}` contract can be passed directly to the loader:

```python
from snipex.llm.sft_data import iter_sft_batches
from snipex.llm.sft_prep import prep_mmlu

batches = iter_sft_batches(..., preprocess=prep_mmlu)
```

```bash
uv run snipex llm sft \
  --train-data smoltalk/train mmlu/train gsm8k/train \
  --val-data smoltalk/val mmlu/val gsm8k/val \
  --tokenizer artifacts/tokenizer \
  --checkpoint runs/base-dual/final.pt \
  --output runs/sft \
  --total-batch-tokens 32768 \
  --max-examples 10000

uv run snipex llm chat \
  --checkpoint runs/sft/final.pt \
  --tokenizer artifacts/tokenizer
```

Generation is deliberately simple in V0.1.0: temperature/top-k autoregressive
sampling without a KV cache.

## 🛠️ Development

Create the environment and run the local checks:

```bash
uv sync --extra llm
uv run ruff check .
uv run pytest -m "not gpu"
uv build
```

Use `algo/<name>` for an algorithm branch. Put its package under
`snipex/<name>/`, its tests under `tests/<name>/`, and declare only its required
dependencies in a matching extra in `pyproject.toml`:

```toml
[project.optional-dependencies]
example = ["example-dependency>=1"]
```

Install or test that algorithm with:

```bash
uv sync --extra example
```

Do not add an extra until a real algorithm needs it.

---

## ❤️ Support

If SnipeX helps you turn an idea into a working prototype, consider giving it
a star on [GitHub](https://github.com/SeriaQ/SnipeX).

SnipeX is released under the [MIT License](https://github.com/SeriaQ/SnipeX/blob/main/LICENSE).
