Metadata-Version: 2.4
Name: aquin
Version: 3.0.5
Summary: Aquin CLI. Run GPU inspection, steering, simulation, and evals locally with aquin load, aquin chat, and aquin inspect.
License-Expression: LicenseRef-Proprietary
Project-URL: Homepage, https://aquin.app
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests>=2.28
Requires-Dist: websockets>=12.0
Requires-Dist: pydantic>=2.0
Requires-Dist: httpx>=0.27
Requires-Dist: certifi>=2024.0
Requires-Dist: safetensors>=0.4
Requires-Dist: typer>=0.12
Requires-Dist: rich>=13.0
Requires-Dist: spinners>=0.0.24
Requires-Dist: textual>=0.80
Requires-Dist: numpy>=1.24
Requires-Dist: openai>=1.0
Requires-Dist: torch>=2.7
Requires-Dist: torchvision>=0.22
Requires-Dist: transformer-lens<3.5,>=3.4
Requires-Dist: transformers<6,>=4.57
Requires-Dist: accelerate>=0.30
Requires-Dist: datasets>=2.18
Requires-Dist: peft>=0.10
Requires-Dist: sentencepiece>=0.2
Requires-Dist: protobuf>=4.0
Requires-Dist: huggingface-hub>=0.23
Requires-Dist: sae-lens>=3.0
Requires-Dist: sentence-transformers>=3.0
Requires-Dist: pandas>=2.0
Requires-Dist: scipy>=1.12
Requires-Dist: scikit-learn>=1.4
Requires-Dist: matplotlib>=3.8
Requires-Dist: fastapi>=0.111
Requires-Dist: uvicorn>=0.29
Provides-Extra: viz
Requires-Dist: umap-learn>=0.5; extra == "viz"
Provides-Extra: compute
Provides-Extra: all
Requires-Dist: aquin[viz]; extra == "all"
Dynamic: license-file

# Aquin SDK

Record your training runs locally and observe them with [Aquin](https://aquin.app) — loss curves, learning rate, grad norm, epoch summaries, and web mirror sync via `aquin session start` + `aquin watch`.

## License

The Aquin CLI and Aquin Engine are proprietary software owned by Aquin Labs Private Limited. Licensor retains all right, title, and interest in the Software; no ownership is transferred to you.

Use requires a valid API key that is personal to you and must not be shared. Licensor grants a limited, non-exclusive, non-transferable, revocable license for your own internal use only. You may not use the Software without an API key, redistribute or sublicense it, reverse engineer it, modify or create derivative works, or use it to build a competing product. API keys may be revoked at any time; upon revocation you must cease all use immediately. You must keep the Software confidential and not disclose it without prior written consent from Aquin Labs. This license terminates automatically on breach, API key revocation, or any attempt to circumvent license enforcement; upon termination you must cease use and destroy all copies. The Software is provided "as is" without warranty; liability is limited to fees paid in the prior twelve months.

Licensing inquiries: aquin@aquin.app, ash@aquin.app — https://aquin.app, aquinlabs.com

See the full [AQUIN ENGINE — PROPRIETARY SOFTWARE LICENSE](LICENSE) or run:

```bash
aquin license
```

On first use, `aquin` displays the license and requires you to accept (`y`) before any other command runs. Non-interactive acceptance after review: `aquin --accept-license` or `AQUIN_LICENSE_ACCEPTED=1`.

## Install

```bash
pip install aquin
```

## Quickstart

```python
import aquin

run = aquin.init(
    base_model="meta-llama/Llama-3.2-1B-Instruct",
    run_name="my-lora-run",
    config={
        "lr": 2e-4, "epochs": 3, "rank": 16, "lora_alpha": 32,
        "method": "qlora", "per_device_train_batch_size": 2,
        "gradient_accumulation_steps": 8, "dataset": "data.jsonl",
    },
)

for epoch in range(3):
    for step, batch in enumerate(dataloader):
        loss = train_step(batch)
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0).item()
        run.log(
            step,
            loss=loss.item(),
            learning_rate=scheduler.get_last_lr()[0],
            grad_norm=grad_norm,
            epoch=epoch,
        )

run.checkpoint(model, step=step)
run.finish()
```

Then replay or sync to the web mirror:

```bash
aquin session start --id my-run --model llama-3.2-1b   # register engine + session
aquin watch <run_id>              # local replay / live tail
aquin watch ingest --run <run_id> --file aquin_run/metrics.jsonl
```

With an active session, watch events appear on the Aquin web mirror in real time.

## Post-training SAE (checkpoint diff, temp train, align)

After `run.checkpoint()` saves `aquin_run/checkpoints/checkpoint.pt`, run SAE tools on the real merged weights (not simulate's synthetic checkpoint):

```bash
aquin session start --id my-run --model llama-3.2-1b
aquin load sae llama-3.2-1b-l8

aquin sae diff --model llama-3.2-1b \
  --checkpoint aquin_run/checkpoints/checkpoint.pt \
  --prompts probes.jsonl --name my-run

aquin sae train --model llama-3.2-1b --layer 8 \
  --checkpoint aquin_run/checkpoints/checkpoint.pt \
  --quick --name my-run

aquin sae align \
  --sae-a ~/.aquin/sae/llama-3.2-1b/sae_layer8.pt \
  --sae-b ~/.aquin/sae/user/llama-3.2-1b/my-run/sae_layer8.pt
```

Docs: [Checkpoint SAE](https://aquin.app/docs/checkpoint-sae). **Watch** ingests metrics only; **simulate** diffs a synthetic NTK checkpoint at forecast end.

## API

### `aquin.init(base_model, run_name, config)`

Starts a new run. Creates `aquin_run/` in the current directory.

| Param | Description |
|---|---|
| `base_model` | HuggingFace model ID, e.g. `"meta-llama/Llama-3.2-1B-Instruct"` |
| `run_name` | Display name for the run |
| `config` | Dict of training hyperparameters (optional, can also pass to `finish()`) |

### `run.log(step, *, loss, ...)`

Record metrics for one training step. Call every step inside your loop.

| Param | Description |
|---|---|
| `step` | Global training step (required) |
| `loss` | Scalar training loss (required) |
| `learning_rate` | Current LR — enables LR chart |
| `grad_norm` | Gradient norm — enables grad norm chart |
| `epoch` | Current epoch — enables epoch summary table |
| `momentum_norm` | Optimizer momentum norm — enables momentum chart |
| `step_ms` | Wall-clock time for this step in ms |

### `run.checkpoint(model, step)`

Saves the model checkpoint locally. One checkpoint per run — always replaces the previous save. Call once at the end of training. Stored under `aquin_run/checkpoints/` for local analysis.

### `run.finish(config)`

Flushes all metrics to disk. Pass `config` here if you didn't pass it to `aquin.init()`.

## CLI

```bash
aquin login       # save your API key
aquin session start --id <name> --model <id>   # connect engine + load model
aquin watch list  # list observed training runs
aquin sae help    # checkpoint SAE diff / train / align
aquin help        # full command list (mode-filtered)
aquin status      # account, API key, and session state
```

## Using with HuggingFace Trainer / TRL

Use a `TrainerCallback` to hook into the training loop:

```python
import time
from transformers import TrainerCallback

class AquinCallback(TrainerCallback):
    def __init__(self, run):
        self.run = run
        self._step_start = 0.0

    def on_step_begin(self, args, state, control, **kwargs):
        self._step_start = time.time()

    def on_log(self, args, state, control, logs=None, **kwargs):
        if not logs or "loss" not in logs:
            return
        self.run.log(
            step=state.global_step,
            loss=float(logs["loss"]),
            learning_rate=float(logs["learning_rate"]) if "learning_rate" in logs else None,
            grad_norm=float(logs["grad_norm"]) if "grad_norm" in logs else None,
            epoch=int(state.epoch) if state.epoch is not None else None,
            step_ms=round((time.time() - self._step_start) * 1000),
        )

    def on_train_end(self, args, state, control, **kwargs):
        model = kwargs.get("model")
        if model:
            self.run.checkpoint(model, step=state.global_step)
```

---

## Building and publishing a new release

**Prerequisites:** Python 3.13, Nuitka, MSVC (Visual Studio Build Tools with Desktop C++ workload).

**1. Compile to native extensions**
```bash
cd cli
python scripts/build_nuitka.py
# Compiles engine/ + compute/ to .pyd, removes .py source, audits on finish
```

**2. Build the wheel**
```bash
python -m build --wheel
# Output: dist/aquin-<version>-py3-none-any.whl
```

**3. Audit — confirm no source leaked**
```bash
python scripts/build_nuitka.py --check
# Must print: Audit passed
```

**4. Bump version before releasing**
Edit `version` in `pyproject.toml`, then repeat steps 1–3.

**5. Distribute**
Send the wheel directly to users (`pip install aquin-*.whl`) or upload to R2 and share a signed link.

**Notes:**
- `.pyd` files and `dist/` are gitignored — never commit compiled artifacts
- After building, `engine/` and `compute/` have no `.py` source locally either — keep a clean git working tree by running builds in a separate branch or restoring source from git after building
- To rebuild from scratch: `git restore cli/aquin/engine cli/aquin/compute` then repeat from step 1
