diff --git a/.planning/PROJECT.md b/.planning/PROJECT.md new file mode 100644 index 0000000..eb66b9e --- /dev/null +++ b/.planning/PROJECT.md @@ -0,0 +1,118 @@ +# Needle + +## What This Is + +Needle is a lightweight Python package for running a compact language model and +agent workflows on constrained devices. It provides a native inference facade, +JAX/Flax reference and training code, LoRA fine-tuning, quantization/export to +`.cact`, a CLI, and a local playground. + +This milestone focuses on making the existing capabilities understandable and +usable by Python and machine-learning beginners through Chinese-first +documentation and reproducible end-to-end examples. + +## Core Value + +A beginner can install Needle and reliably go from a first inference request to +a fine-tuned, exported model without guessing which assets, commands, or +runtime constraints apply. + +## Requirements + +### Validated + +- [x] Native `Needle` API supports completion, tool execution, schema extraction, + reset, and weight loading — existing in `needle/__init__.py`. +- [x] CLI exposes fetch/download, run, fine-tune, data generation, build, + playground, and related workflows — existing in `needle/cli.py`. +- [x] JAX/Flax reference model defines the Simple Attention Network, decoding, + LoRA training, quantization, and `.cact` export paths — existing in + `needle/model/`. +- [x] Test suite covers model, inference, fine-tuning, export, packaging, + environments, and CLI behavior — existing in `tests/`. + +### Active + +- [ ] DOC-01: Provide a Chinese-first installation and environment guide that + explains CPU and GPU options, optional extras, model assets, caches, and + offline operation. +- [ ] DOC-02: Provide a copy-paste quickstart that runs a pre-trained inference + example and explains the public `Needle` API and CLI equivalents. +- [ ] DOC-03: Provide an end-to-end fine-tuning tutorial covering JSONL data + format, prompt rendering, LoRA training, checkpoint outputs, and common + resource/configuration choices. +- [ ] DOC-04: Provide an export/deployment tutorial covering merge, quantization, + `.cact` generation, engine compatibility, and running the exported artifact. +- [ ] DOC-05: Document the model structure and data flow, including the native + runtime path versus the JAX/Flax training path, attention/MLP components, + tokenizer contract, and process-global state constraints. +- [ ] DOC-06: Add troubleshooting and safety notes for download verification, + pickle checkpoints, playground exposure, concurrency, dependency drift, and + known training/export limitations; track code hardening as later work. +- [ ] DOC-07: Keep examples and commands testable on CPU and GPU where supported, + with a clear verification checklist for each tutorial. + +### Out of Scope + +- Rewriting the native inference engine or changing the `.cact` binary contract + — documentation should reflect the current implementation first. +- Adding authentication, hosted multi-user serving, or a production web service + — the current playground remains a local development tool. +- Solving every security or performance concern in this documentation milestone + — risks are recorded and prioritized for later implementation phases. + +## Context + +- The package targets CPython 3.9+ and uses setuptools with optional `train`, + `gpu`, `metal`, and `test` extras. +- Production inference loads a platform-specific native library through ctypes; + JAX/Flax/Optax/SentencePiece are used for reference inference, fine-tuning, + and export. +- Model and tokenizer artifacts are commonly fetched from Hugging Face and may + be cached locally; `HF_HUB_OFFLINE=1` supports air-gapped use. +- The native engine keeps process-global active state, so base and tuned agents + have ordering and process-isolation constraints. +- The codebase map in `.planning/codebase/` is the evidence source for current + structure, integrations, conventions, testing, and concerns. + +## Constraints + +- **Audience**: Write for Python/ML beginners, while linking to source paths for + readers who need implementation detail. +- **Language**: Chinese is the primary user-facing documentation language; keep + API names, commands, paths, and code identifiers exact. +- **Platforms**: Cover both CPU-first setup and supported NVIDIA CUDA/Apple + Metal acceleration without claiming unsupported combinations. +- **Compatibility**: Preserve public APIs, CLI behavior, checkpoint formats, and + `.cact` tensor ordering while improving documentation. +- **Verification**: Every tutorial must state prerequisites, expected output, + and a practical way to verify success. + +## Key Decisions + +| Decision | Rationale | Outcome | +|----------|-----------|---------| +| Chinese-first documentation | The intended onboarding audience asked for Chinese guidance | - Pending | +| Cover inference, LoRA, and deployment in one path | Users need a complete journey from install to usable tuned artifact | - Pending | +| Support CPU and GPU guidance | Hardware availability varies and beginner setup should not assume CUDA | - Pending | +| Record hardening risks before fixing them | Documentation can prevent misuse without expanding the first milestone into a security rewrite | - Pending | + +## Evolution + +This document evolves at phase transitions and milestone boundaries. + +**After each phase transition** (via `$gsd-transition`): +1. Requirements invalidated? -> Move to Out of Scope with reason +2. Requirements validated? -> Move to Validated with phase reference +3. New requirements emerged? -> Add to Active +4. Decisions to log? -> Add to Key Decisions +5. "What This Is" still accurate? -> Update if drifted + +**After each milestone** (via `$gsd-complete-milestone`): +1. Full review of all sections +2. Core Value check - still the right priority? +3. Audit Out of Scope - reasons still valid? +4. Update Context with current state + +--- +*Last updated: 2026-08-31 after initialization* diff --git a/.planning/REQUIREMENTS.md b/.planning/REQUIREMENTS.md new file mode 100644 index 0000000..446e0da --- /dev/null +++ b/.planning/REQUIREMENTS.md @@ -0,0 +1,106 @@ +# Requirements: Needle Documentation and Onboarding + +**Defined:** 2026-08-31 +**Core Value:** A beginner can install Needle and reliably go from a first inference request to a fine-tuned, exported model without guessing which assets, commands, or runtime constraints apply. + +## v1 Requirements + +### Installation and Environment + +- [x] **INST-01**: A beginner can create an isolated Python 3.9+ environment and install the correct base, training, test, GPU, or Metal extras for their platform. +- [x] **INST-02**: The guide explains native engine, checkpoint, tokenizer, Hugging Face cache locations, offline mode, and required environment variables. +- [x] **INST-03**: The guide states a supported CPU/GPU/Metal matrix and gives a CPU fallback when acceleration is unavailable. + +### Inference Quickstart + +- [x] **INFR-01**: A beginner can run a copy-paste pre-trained inference example and identify the expected successful output. +- [x] **INFR-02**: The guide demonstrates equivalent CLI and `Needle` Python API flows, including a minimal typed tool call or extraction example. +- [x] **INFR-03**: Each quickstart includes a verification command and links common errors to troubleshooting guidance. + +### Model and Runtime Concepts + +- [ ] **MODL-01**: The documentation explains the native ctypes runtime path versus the JAX/Flax reference and training path. +- [ ] **MODL-02**: The documentation describes the Simple Attention Network components, tokenizer contract, checkpoint formats, and `.cact` artifact roles with source links. +- [ ] **MODL-03**: The documentation calls out process-global native state, base/tuned agent ordering, and when to use a fresh process. + +### LoRA Fine-Tuning + +- [ ] **LORA-01**: A beginner can create a valid JSONL fine-tuning dataset with documented fields, chat/tool markers, and target masking behavior. +- [ ] **LORA-02**: A beginner can run a small local LoRA fine-tuning job, understand its prerequisites and resource knobs, and locate the adapter/checkpoint outputs. +- [ ] **LORA-03**: The guide explains how to inspect or validate a trained adapter and documents the current limitation around interrupted-run resume. + +### Export and Deployment + +- [ ] **DEPL-01**: A beginner can merge a LoRA adapter, quantize it, and produce a `.cact` archive using the supported CLI workflow. +- [ ] **DEPL-02**: The guide explains engine version, tensor order, model geometry, and tokenizer vocabulary compatibility requirements. +- [ ] **DEPL-03**: A beginner can load the exported archive in a fresh process and compare its output with the reference or base path using a verification checklist. + +### Troubleshooting and Safety + +- [ ] **SAFE-01**: Troubleshooting maps installation, asset, backend, checkpoint, tokenizer, and export failures to observable symptoms and fixes. +- [ ] **SAFE-02**: Documentation warns about untrusted pickle checkpoints, native downloads, API keys, and exposing the unauthenticated local playground. +- [ ] **SAFE-03**: Documentation records known concurrency, dependency drift, performance, and test-coverage limitations without implying they are solved. + +### Documentation Quality + +- [x] **DOCS-01**: The primary onboarding path is Chinese-first, preserves exact commands/API identifiers, and links every conceptual claim to a source path or official reference. +- [x] **DOCS-02**: Every tutorial states prerequisites, expected output, cleanup/cache behavior, and a practical verification step for CPU and supported accelerator branches. +- [ ] **DOCS-03**: Examples are organized by user goal and can be checked in automated tests or a documented manual verification pass. + +## v2 Requirements + +### Productization + +- **PROD-01**: Provide authenticated hosted or multi-user serving guidance. +- **PROD-02**: Provide a GUI for training job management and artifact browsing. +- **PROD-03**: Add automatic artifact signing, checksums, or registry provenance enforcement. + +### Advanced Learning Material + +- **LEARN-01**: Provide notebook-based walkthroughs and interactive visualizations. +- **LEARN-02**: Publish benchmark tables across model sizes, hardware backends, and quantization levels. + +## Out of Scope + +| Feature | Reason | +|---------|--------| +| Native engine rewrite | This milestone documents the current runtime contract rather than changing it. | +| New `.cact` format or tensor ordering | Format changes require coordinated engine/export work and are not documentation-only. | +| Production security hardening | Risks are documented and tracked, but code remediation is a later phase. | +| Universal OS/GPU support claims | Backend availability changes; docs will state tested/supportable combinations only. | + +## Traceability + +| Requirement | Phase | Status | +|-------------|-------|--------| +| INST-01 | Phase 1 | Complete | +| INST-02 | Phase 1 | Complete | +| INST-03 | Phase 1 | Complete | +| INFR-01 | Phase 1 | Complete | +| INFR-02 | Phase 1 | Complete | +| INFR-03 | Phase 1 | Complete | +| MODL-01 | Phase 2 | Pending | +| MODL-02 | Phase 2 | Pending | +| MODL-03 | Phase 2 | Pending | +| LORA-01 | Phase 3 | Pending | +| LORA-02 | Phase 3 | Pending | +| LORA-03 | Phase 3 | Pending | +| DEPL-01 | Phase 4 | Pending | +| DEPL-02 | Phase 4 | Pending | +| DEPL-03 | Phase 4 | Pending | +| SAFE-01 | Phase 2 | Pending | +| SAFE-02 | Phase 2 | Pending | +| SAFE-03 | Phase 2 | Pending | +| DOCS-01 | Phase 1 | Complete | +| DOCS-02 | Phase 1 | Complete | +| DOCS-03 | Phase 4 | Pending | + +**Coverage:** + +- v1 requirements: 21 total +- Mapped to phases: 21 +- Unmapped: 0 + +--- +*Requirements defined: 2026-08-31* +*Last updated: 2026-08-31 after initial definition* diff --git a/.planning/ROADMAP.md b/.planning/ROADMAP.md new file mode 100644 index 0000000..7b2d24c --- /dev/null +++ b/.planning/ROADMAP.md @@ -0,0 +1,111 @@ +# Roadmap: Needle Documentation and Onboarding + +## Overview + +Deliver a Chinese-first beginner journey through the existing Needle lifecycle: +install and fetch assets, run inference, understand the model/runtime boundary, +fine-tune with LoRA, and export/deploy a verified `.cact` artifact. Each phase +ships a usable documentation slice with executable examples and explicit limits. + +## Phases + +- [x] **Phase 1: Install and First Inference** - Make a clean environment and first response repeatable. (completed 2026-08-31) +- [ ] **Phase 2: Model and Runtime Concepts** - Explain architecture, artifacts, troubleshooting, and safety boundaries. +- [ ] **Phase 3: LoRA Fine-Tuning** - Document data preparation and a reproducible local adaptation run. +- [ ] **Phase 4: Export and Deployment Verification** - Teach merge, quantize, `.cact` loading, and end-to-end checks. + +## Phase Details + +### Phase 1: Install and First Inference + +**Goal**: A beginner can install Needle on CPU or a supported accelerator, obtain required assets, and run a first inference request. +**Mode**: mvp +**Depends on**: Nothing (first phase) +**Requirements**: [INST-01, INST-02, INST-03, INFR-01, INFR-02, INFR-03, DOCS-01, DOCS-02] +**Success Criteria** (what must be TRUE): + + 1. A clean Python 3.9+ environment can be installed using documented commands for CPU and supported accelerator branches. + 2. A reader can fetch or locate engine/checkpoint/tokenizer assets and knows the cache and offline controls. + 3. A reader can copy a CLI or `Needle` API example, observe the expected response, and run a verification check. + 4. Common install and first-run failures link to actionable troubleshooting entries. + +**Plans**: 1/2 plans executed + +Plans: + +- [x] 01-01-PLAN.md +- [x] 01-02-PLAN.md +- [x] 01-01: Write installation, asset/cache, and backend matrix guide. +- [x] 01-02: Write and verify CLI/API inference quickstart with expected output. + +### Phase 2: Model and Runtime Concepts + +**Goal**: A beginner can explain which runtime path and artifact applies to inference, training, and deployment, and can avoid known unsafe usage. +**Mode**: mvp +**Depends on**: Phase 1 +**Requirements**: [MODL-01, MODL-02, MODL-03, SAFE-01, SAFE-02, SAFE-03] +**Success Criteria** (what must be TRUE): + + 1. Architecture documentation traces a request through `Needle`/ctypes and separately through JAX/Flax reference code. + 2. The roles and compatibility constraints of `.pkl`, tokenizer, LoRA adapter, and `.cact` artifacts are clear with source links. + 3. A reader can identify process-global state, pickle/download risks, playground exposure risks, and documented performance/dependency limitations. + 4. Troubleshooting is organized by symptom and includes a safe next action. + +**Plans**: 2 plans + +Plans: + +- [ ] 02-01: Write model architecture, data-flow, and artifact glossary. +- [ ] 02-02: Write troubleshooting and safety reference from codebase concerns. + +### Phase 3: LoRA Fine-Tuning + +**Goal**: A beginner can prepare valid supervision data and complete a small, reproducible LoRA fine-tuning run. +**Mode**: mvp +**Depends on**: Phase 2 +**Requirements**: [LORA-01, LORA-02, LORA-03] +**Success Criteria** (what must be TRUE): + + 1. A documented JSONL example renders to the expected chat/tool markers and explains which tokens contribute to loss. + 2. A small local fine-tuning command runs with stated CPU/GPU prerequisites and resource knobs. + 3. Adapter/checkpoint outputs have documented paths and a validation or inspection step. + 4. The guide clearly states current interrupted-run resume limitations and safe checkpoint handling. + +**Plans**: 2 plans + +Plans: + +- [ ] 03-01: Document JSONL schema, rendering/masking, and dataset validation. +- [ ] 03-02: Document and run the LoRA CLI workflow with output verification. + +### Phase 4: Export and Deployment Verification + +**Goal**: A beginner can turn a trained adapter into a compatible `.cact` artifact and verify it in a fresh runtime process. +**Mode**: mvp +**Depends on**: Phase 3 +**Requirements**: [DEPL-01, DEPL-02, DEPL-03, DOCS-03] +**Success Criteria** (what must be TRUE): + + 1. Merge, quantize, and build commands produce a named `.cact` archive from the documented adapter output. + 2. Engine version, tensor order, geometry, and tokenizer vocabulary checks are explicit and testable. + 3. A fresh process loads the archive and compares output against a reference/base path using a checklist. + 4. Examples have a documented manual or automated verification route suitable for CI follow-up. + +**Plans**: 2 plans + +Plans: + +- [ ] 04-01: Write merge/quantize/export/deployment tutorial and compatibility checklist. +- [ ] 04-02: Add or document example verification commands and review all links/commands. + +## Progress + +**Execution Order:** +Phases execute in numeric order: 1 -> 2 -> 3 -> 4 + +| Phase | Plans Complete | Status | Completed | +|-------|----------------|--------|-----------| +| 1. Install and First Inference | 2/2 | Complete | 2026-08-31 | +| 2. Model and Runtime Concepts | 0/2 | Not started | - | +| 3. LoRA Fine-Tuning | 0/2 | Not started | - | +| 4. Export and Deployment Verification | 0/2 | Not started | - | diff --git a/.planning/STATE.md b/.planning/STATE.md new file mode 100644 index 0000000..aebc7c7 --- /dev/null +++ b/.planning/STATE.md @@ -0,0 +1,86 @@ +--- +gsd_state_version: 1.0 +current_phase: 2 +current_phase_name: Model and Runtime Concepts +status: planning +stopped_at: Phase 01 complete, ready to plan Phase 2 +last_updated: "2026-08-31T14:26:10.336Z" +last_activity: 2026-08-31 +last_activity_desc: Phase 01 complete, transitioned to Phase 2 +state_head: 7c5922f22eb091a4a7548dc97d217e766f3f189c +progress: + total_phases: 4 + completed_phases: 1 + total_plans: 2 + completed_plans: 2 + percent: 25 +--- + +# Project State + +## Project Reference + +See: `.planning/PROJECT.md` (updated 2026-08-31) + +**Core value:** A beginner can install Needle and reliably go from a first inference request to a fine-tuned, exported model without guessing which assets, commands, or runtime constraints apply. +**Current focus:** Phase 02 — Model and Runtime Concepts + +## Current Position + +Phase: 2 — Model and Runtime Concepts +Plan: Not started +Status: Ready to plan +Last activity: 2026-08-31 — Phase 01 complete, transitioned to Phase 2 + +Progress: [██░░░░░░░░] 25% + +## Performance Metrics + +**Velocity:** + +- Total plans completed: 2 +- Average duration: n/a +- Total execution time: 0 hours + +**By Phase:** + +| Phase | Plans | Total | Avg/Plan | +|-------|-------|-------|----------| +| 01 | 2 | - | - | +**Per-Plan Metrics:** + +| Plan | Duration | Tasks | Files | +|------|----------|-------|-------| +| Phase 01 P01 | 14m | 3 tasks | 2 files | + +## Accumulated Context + +### Decisions + +Decisions are logged in `.planning/PROJECT.md` Key Decisions table. + +- Chinese-first documentation with exact commands and source links. +- Vertical MVP roadmap from install through deployment. +- Research, plan checks, verification, and parallel execution enabled. +- [Phase 01]: Phase 1 plan 01 uses uv editable [train,test] installation and explicit needle fetch before CPU inference. +- [Phase 01]: Phase 1 plan 01 verifies online engine fetch followed by HF_HUB_OFFLINE=1 loading; CUDA and Metal remain deferred. + +### Pending Todos + +None yet. + +### Blockers/Concerns + +- CUDA and Metal installation/runtime verification remain intentionally deferred to a later phase. + +## Deferred Items + +| Category | Item | Status | Deferred At | Milestone | +|----------|------|--------|-------------|-----------| +| Productization | Hosted auth, GUI training, artifact signing | Deferred | 2026-08-31 | v1 docs | + +## Session Continuity + +Last session: 2026-08-31T14:04:38.688Z +Stopped at: Phase 01 complete, ready to plan Phase 2 +Resume file: None diff --git a/.planning/WINDOWS.md b/.planning/WINDOWS.md new file mode 100644 index 0000000..48598a9 --- /dev/null +++ b/.planning/WINDOWS.md @@ -0,0 +1,35 @@ +--- +schema_version: 1 +open_count: 1 +waived_count: 0 +fixed_count: 0 +total_count: 1 +last_updated: 2026-08-31T14:03:01.105Z +--- + +# Broken Windows Ledger + +> Cross-phase defect register. With `workflow.windows_enforce` enabled, `/gsd-ship` blocks while `open_count > 0`. +> Waive with `gsd-tools windows waive ""` (reason required). +> Mark fixed with `gsd-tools windows fixed `. + +| id | phase | kind | file | line | description | status | reason | recorded_at | resolved_at | +|----|-------|------|------|------|-------------|--------|--------|-------------|-------------| +| 1 | 01 | todo | doc/installation.md | 5 | CUDA and Metal installation and verification are intentionally deferred to a later phase. | open | | 2026-08-31T14:03:01.105Z | | + +````json +[ + { + "id": 1, + "kind": "todo", + "phase": "01", + "file": "doc/installation.md", + "line": 5, + "description": "CUDA and Metal installation and verification are intentionally deferred to a later phase.", + "status": "open", + "reason": "", + "recorded_at": "2026-08-31T14:03:01.105Z", + "resolved_at": null + } +] +```` diff --git a/.planning/codebase/ARCHITECTURE.md b/.planning/codebase/ARCHITECTURE.md new file mode 100644 index 0000000..39834d6 --- /dev/null +++ b/.planning/codebase/ARCHITECTURE.md @@ -0,0 +1,182 @@ + +# Architecture + +**Analysis Date:** 2026-08-31 + +## System Overview + +```text + Python public API / CLI + `needle/__init__.py` `needle/cli.py` + | | + v v + Native C engine (ctypes) JAX/Flax reference + training + `needle/agent/fetch.py` `needle/model/*.py` + | | + v v + `.cact` weights + grammar `.pkl` checkpoints / LoRA + | | + +------------+------------+ + v + Device inference/runtime + `needle.Needle`, playground, + ready-made environments +``` + +The package has two intentionally separate execution paths. Production inference uses a platform-specific shared library (`libneedle.so`, `.dylib`, or `.dll`) loaded by `ctypes` and a self-contained `.cact` archive. Training, checkpoint inspection, reference generation, and export use Python, JAX, Flax, and NumPy under `needle/model`. + +## Component Responsibilities + +| Component | Responsibility | File | +|-----------|----------------|------| +| Public agent API | Resolve Python callables/Pydantic models to JSON schemas, bind the engine, complete requests, execute tool loops | `needle/__init__.py` | +| Tool schema builder | Convert annotations, `Literal`, enums, `Field`, docstrings, and Pydantic models into schemas | `needle/agent/tools.py` | +| Engine fetcher | Select platform tag, download/cache the native engine from Hugging Face | `needle/agent/fetch.py` | +| CLI router | Parse `run`, `finetune`, `generate-data`, `build`, `download`, `fetch`, and `playground` commands | `needle/cli.py` | +| Reference model | Define TransformerConfig, Simple Attention Network modules, masks, KV-window sizing | `needle/model/architecture.py` | +| Reference decode | Load `.pkl` checkpoints and run greedy/temperature generation in JAX | `needle/model/run.py`, `needle/model/decode.py` | +| Fine-tuning pipeline | Render JSONL examples, synthesize data, train LoRA, merge adapter, call export | `needle/model/finetune.py` | +| Quantization | Fake-QAT and Cactus Quants codebooks, mixed bit maps, deployment quantization | `needle/model/quantize.py` | +| Export format | Pack tensors, metadata, codebooks, and SentencePiece tokenizer into `.cact`; read it back | `needle/model/export.py` | +| Playground | Threaded HTTP server around `Needle`, model loading, completion, and background fine-tune | `needle/playground/server.py` | +| Environments | Curated tool schemas and frozen acceptance cases | `needle/environments/*.py`, `needle/environments/_harness.py` | + +## Pattern Overview + +**Overall:** Thin public facade over a native runtime, with a Flax reference/training implementation and an offline binary export boundary. + +**Key Characteristics:** +- The native engine is process-global: `needle/__init__.py` keeps one loaded library, active agent, and active weight blob. A tuned archive cannot be unloaded in-process. +- Tool declarations are data at the engine boundary. Python callables are retained locally for `Needle.run()` execution, while schemas are serialized to JSON for constrained decoding. +- The model output head is tied to `embedding.embedding`; `.cact` export transposes and quantizes runtime matrices in a fixed positional tensor order. +- The reference stack is scanned over `num_layers` with per-layer parameter axes, optional rematerialization, and configurable flash attention. + +## Layers + +**Public API and orchestration:** +- Purpose: User-facing completion, extraction, and agentic tool execution. +- Location: `needle/__init__.py`, `needle/agent/tools.py`. +- Contains: `Needle`, `extract`, `tool`, `Field`, schema conversion and result normalization. +- Depends on: `ctypes`, native engine symbols, Hugging Face fetch fallback. +- Used by: README examples, environments, playground, downstream applications. + +**Native engine adapter:** +- Purpose: Load the platform binary and expose `needle_init`, `needle_complete`, `needle_reset`, and `needle_load`. +- Location: `needle/__init__.py`, `needle/agent/fetch.py`. +- Contains: cache lookup, download, symbol signatures, output buffer handling, process-global weight state. +- Depends on: platform detection and Hugging Face model artifacts. +- Used by: every production `Needle` call and the playground. + +**Model definition:** +- Purpose: Compute logits and auxiliary heads in JAX/Flax. +- Location: `needle/model/architecture.py`. +- Contains: `ZCRMSNorm`, `MultiHeadAttention`, `HadamardMLP`, `Block`, `Stack`, `Engram`, mHC routing, contrastive/confidence heads, masks. +- Depends on: JAX, Flax, quantization helpers. +- Used by: fine-tuning, checkpoint generation, reference decode, export metadata. + +**Training/export:** +- Purpose: Convert JSONL supervision into a LoRA adapter and deployable archive. +- Location: `needle/model/finetune.py`, `needle/model/quantize.py`, `needle/model/export.py`. +- Contains: prompt rendering, masked causal loss, LoRA over attention projections, CQ packing, tokenizer embedding. +- Depends on: base `.pkl` checkpoint, SentencePiece tokenizer, optional OpenRouter API for data generation. +- Used by: CLI `finetune`, `generate-data`, and `build`; playground background fine-tune. + +## Data Flow + +### Production Completion Path + +1. `Needle.__init__` resolves tools to schemas and records callable implementations (`needle/__init__.py:56-114`). +2. `_bind` locates/downloads the native library, optionally loads `.cact` bytes, and calls `needle_init` with system text, tools JSON, and optional tool index (`needle/__init__.py:75-99`). +3. `Needle.complete` sends UTF-8 text and a bounded output buffer to `needle_complete` (`needle/__init__.py:119-137`). +4. The native engine returns a JSON envelope containing `type`, text/function calls, and confidence; tuned weights force confidence to `None` because the confidence head is not adapted. +5. `Needle.run` invokes local Python functions for returned calls, serializes results, feeds them back through `_complete`, and attaches all executed results (`needle/__init__.py:139-160`). + +### Reference Model Path + +1. `load_checkpoint` reads a format-v2 pickle and reconstructs `TransformerConfig` (`needle/model/run.py:37-72`). +2. `SimpleAttentionNetwork.__call__` embeds token IDs, computes RoPE and engram key/value vectors, runs `Stack`, applies final normalization, and projects against the tied embedding table (`needle/model/architecture.py:478-540`). +3. Each scanned `Block` performs ZCRMS pre-normalization, GQA attention with RoPE and sigmoid output gating, a residual attention gate, then a sandwich-normalized Hadamard MLP (`needle/model/architecture.py:305-375`). +4. `generate`/`batch_generate` repeatedly call the JIT model, append tokens until EOS/max length, and decode with SentencePiece (`needle/model/run.py:89-183`). + +### Fine-Tune and Deployment Path + +1. JSONL examples are rendered with chat/tool markers by `render_example`; only target tokens contribute to masked cross-entropy (`needle/model/finetune.py:194-243`). +2. `finetune_local` loads the frozen base checkpoint, initializes LoRA A/B matrices only for `q_proj`, `k_proj`, `v_proj`, `gate_proj`, and `out_proj`, then optimizes them with clipped AdamW and a warmup/cosine schedule (`needle/model/finetune.py:254-390`). +3. `build_main` merges the adapter into the base parameters (`merge_lora`) and calls `write_export` (`needle/model/finetune.py:404-440`). +4. `write_export` emits a fixed-header `.cact` containing quantized layer-major tensors, codebooks, optional probe heads, and a raw tokenizer blob (`needle/model/export.py:340-393`). +5. `Needle(weights=...)` loads that archive into the same native engine without recompilation. + +**State Management:** JAX model parameters are immutable pytrees; optimizer state is local to `finetune_local`. Native runtime state is mutable and process-global (`_active`, `_active_weights`, `_active_blob`), and `Needle.reset()` resets the native conversation/KV state. Playground serializes engine operations with a lock. + +## Key Abstractions + +**`TransformerConfig`:** Dataclass carrying model geometry, dtype, RoPE, engram, mHC, KV-window, and quantization settings (`needle/model/architecture.py:58-94`). Keep checkpoint config and export header geometry aligned. + +**`SimpleAttentionNetwork`:** Flax top-level model. It owns embedding, `Stack`, auxiliary heads, engram sites, and optional MTP block (`needle/model/architecture.py:478-579`). + +**`Needle`:** Runtime facade that owns schemas and Python callables while delegating decoding to native C (`needle/__init__.py:56-169`). Use `complete` for one response, `run` for tool execution, and `extract` for one-shot schema output. + +**`.cact` positional archive:** Binary contract between Python export and the native engine. Tensor order and header geometry are defined in `needle/model/export.py:1-80`; do not reorder tensors without a matching engine change. + +## Entry Points + +**Python package:** `needle/__init__.py` exports `Needle`, `tool`, `Field`, `extract`. + +**CLI:** `needle/cli.py:main` is registered as the `needle` console script in `pyproject.toml`; command handlers delegate to model/runtime modules. + +**Reference generation:** `needle/model/run.py:main` backs `needle run` and requires a `.pkl` checkpoint. + +**Playground:** `needle/playground/server.py:main` serves static UI and HTTP endpoints on `127.0.0.1:7860` by default. + +**Environment suites:** `python -m needle.environments.smart_home` (and sibling modules) exercises the native agent through `_harness.run_tests`. + +## Architectural Constraints + +- **Process-global native state:** One native engine/library and one active weight archive are shared per process; construct base agents before tuned agents or isolate processes. +- **Format compatibility:** `.cact` archives are tied to engine version (`ENGINE_VERSION` in `needle/agent/fetch.py`); rebuild archives after package/engine upgrades. +- **Shape constraints:** Export currently requires equal query/key and value head dimensions and rejects unsupported lexicon or local/global sliding-window configurations (`needle/model/export.py:89-111`). +- **Memory constraint:** KV cache sizing is computed from an ~11.5 MiB budget and aligned to `KV_GROUP`; `effective_kv_window` caps it to the configured maximum (`needle/model/architecture.py:603-620`). +- **Tokenizer contract:** Export embeds the SentencePiece vocabulary and special-token IDs; tokenizer vocabulary must equal `config.vocab_size` (`needle/model/export.py:340-346`). +- **Single-threaded model update:** JAX training mutates no shared model state, while playground native calls are explicitly locked (`needle/playground/server.py:17-54`). + +## Anti-Patterns + +### Loading Base Weights After Tuned Weights + +**What happens:** Constructing a base `Needle` after a tuned agent raises instead of unloading the tuned archive (`needle/__init__.py:79-84`). +**Why it's wrong:** The native engine cannot unload weights and would silently answer with the wrong model. +**Do this instead:** Instantiate base agents first, use one tuned archive per process, or run separate processes. + +### Bypassing Schema Generation + +**What happens:** Passing arbitrary callable objects without usable annotations/docstrings produces weak schemas (`needle/agent/tools.py:100-145`). +**Why it's wrong:** Constrained decoding depends on accurate JSON types, enums, bounds, and descriptions. +**Do this instead:** Annotate every argument, use `Literal`/`Field` for constraints, and document callable behavior. + +### Editing `.cact` Tensor Order Independently + +**What happens:** A custom export that changes positional tensor order can still produce a file but native loading interprets weights incorrectly. +**Why it's wrong:** The runtime directory is intentionally nameless and positional (`needle/model/export.py:367-379`). +**Do this instead:** Extend both export and engine format together, with round-trip tests in `tests/test_build.py`. + +## Error Handling + +**Strategy:** Raise explicit Python exceptions at boundaries; return structured tool errors inside `Needle.run`; return JSON error bodies from playground handlers. + +**Patterns:** +- Missing/invalid checkpoints raise `ValueError` with format-version details (`needle/model/run.py:54-72`). +- Native negative return codes become `RuntimeError` (`needle/__init__.py:123-134`). +- Tool lookup/execution errors are appended as `{"error": ...}` results rather than aborting the loop (`needle/__init__.py:147-158`). +- Playground catches request exceptions and responds with an error envelope (`needle/playground/server.py:133-170`). + +## Cross-Cutting Concerns + +**Logging:** CLI/model paths print aligned progress lines; native XLA noise is filtered at startup in `needle/cli.py:20-111`; playground suppresses HTTP access logs. + +**Validation:** Tool schema constraints are compiled into the engine's constrained decoder; environment suites compare exact function-call JSON and optionally apply confidence gates (`needle/environments/_harness.py:11-46`). + +**Authentication:** The runtime itself has no user auth. Optional OpenRouter data synthesis uses `OPENROUTER_API_KEY` in `needle/model/finetune.py`; Hugging Face access uses the ambient `huggingface_hub` configuration. + +--- + +*Architecture analysis: 2026-08-31* diff --git a/.planning/codebase/CONCERNS.md b/.planning/codebase/CONCERNS.md new file mode 100644 index 0000000..bed1af1 --- /dev/null +++ b/.planning/codebase/CONCERNS.md @@ -0,0 +1,123 @@ +# Codebase Concerns + +**Analysis Date:** 2026-08-31 + +## Tech Debt + +**Fine-tune checkpointing and resume:** +- Issue: `finetune_local` accepts `checkpoint_dir` and creates the directory only when training finishes; it never writes step checkpoints or supports resuming an interrupted run. +- Files: `needle/model/finetune.py:294-400`, `needle/cli.py:111-132` +- Impact: Long JAX/accelerator runs lose all progress on interruption, and the advertised checkpoint directory can appear healthy while remaining empty until completion. +- Fix approach: Persist adapter parameters, optimizer state, config, and step at a configurable interval using an atomic rename; add an explicit `--resume` path and a recovery test. + +**Silent configuration drops:** +- Issue: `TransformerConfig.__init__` ignores unknown checkpoint keys instead of rejecting or recording them. +- Files: `needle/model/architecture.py:58-90` +- Impact: A misspelled or newer config field silently falls back to a default, potentially constructing a model with incompatible dimensions or attention/engram behavior before export. +- Fix approach: Validate keys and dimension invariants (head divisibility, engram layers, dtype) and fail with the offending field; keep a migration path for intentionally deprecated fields. + +**Unpinned training stack and release metadata drift:** +- Issue: JAX, jaxlib, optax, numpy, and sentencepiece are broadly specified without a lockfile; package metadata reports `2.0.8` while the checked-out repository is tagged `v2.0.11`, and the native engine is independently pinned to `2.0.3`. +- Files: `pyproject.toml:3-31`, `requirements-train.txt:1-8`, `needle/__init__.py:12`, `needle/agent/fetch.py:8` +- Impact: Reproducing a fine-tune/export environment is difficult, and package/engine/CACT compatibility can be misdiagnosed from inconsistent version identifiers. +- Fix approach: Generate a supported lock/test matrix, make the package version single-sourced, and publish an explicit package-to-engine/format compatibility table. + +## Known Bugs + +**Empty batch generation is not handled:** +- Symptoms: Calling `batch_generate` with an empty `prompts` sequence raises `ValueError` from `max(plens)` instead of returning an empty result. +- Files: `needle/model/run.py:119-126` +- Trigger: Any caller builds a batch dynamically and passes no requests. +- Workaround: Guard at the API boundary and return `[]` (or reject with a clear `ValueError`) before computing `max()`. + +**Engine calls race across agents/threads:** +- Symptoms: Concurrent calls can overwrite process-global active agent/weights and the per-agent response buffer, producing responses for the wrong tool set or corrupted JSON. +- Files: `needle/__init__.py:34-99`, `needle/__init__.py:119-168` +- Trigger: Two threads call `Needle.complete`, `run`, or `reset`, or bind agents with different tool/weight sets in one process; `_bind` has no lock. +- Workaround: Serialize all calls in the application (the playground does this only around its own `Engine.complete`); isolate incompatible workloads in separate processes. + +## Security Considerations + +**Unsafe pickle deserialization:** +- Risk: Checkpoints and LoRA adapters are loaded with `pickle.load`, which can execute arbitrary code when an attacker controls a file or a downloaded artifact. +- Files: `needle/model/run.py:31-35`, `needle/model/run.py:84-85`, `needle/model/finetune.py:412-417` +- Current mitigation: A format-version field is checked only after deserialization; no signature or trusted-source check is performed. +- Recommendations: Replace the interchange format with a safe tensor container (for example, NPZ/safetensors plus JSON metadata), or require signatures and an explicit trusted-file policy before loading pickle. + +**Unverified native artifact and tokenizer downloads:** +- Risk: The package downloads a wheel/engine and tokenizer from Hugging Face, extracts a shared library, and writes it to the cache/package without a project-level digest or signature verification. +- Files: `needle/agent/fetch.py:57-100`, `needle/model/tokenizer.py:82-112`, `needle/__init__.py:16-31` +- Current mitigation: Hugging Face client cache/ETag handling is delegated to the SDK; the archive member path is fixed. +- Recommendations: Pin and verify SHA-256/signatures for every native and tokenizer artifact, reject unexpected archive metadata, and expose an offline verification failure. + +**Unauthenticated playground control plane:** +- Risk: When bound beyond localhost, any caller can upload arbitrary model bytes, invoke OpenRouter-backed fine-tuning with their own or a supplied API key, and consume unbounded memory/disk/network resources. +- Files: `needle/playground/server.py:99-170`, `needle/playground/server.py:176-186` +- Current mitigation: CLI defaults to `127.0.0.1`; uploaded filenames are reduced with `basename`. +- Recommendations: Keep localhost as the only default, require an auth token for non-loopback hosts, enforce request/body/sample limits, rate-limit `/complete` and `/finetune`, and never accept API keys from unauthenticated remote clients. + +## Performance Bottlenecks + +**Per-event telemetry thread creation:** +- Problem: Every tracked operation creates a new daemon thread and performs a network request, including frequent `complete` calls. +- Files: `needle/_telemetry.py:59-84`, `needle/__init__.py:119-121`, `needle/__init__.py:139-141` +- Cause: There is no bounded queue or shared sender; high-QPS applications can accumulate threads and connection overhead. +- Improvement path: Use a bounded queue and one sender thread (or sampled/batched events), with explicit shutdown behavior and back-pressure/drop metrics. + +**Dense Walsh-Hadamard materialization during model application:** +- Problem: `HadamardMLP` constructs a dense `n x n` NumPy matrix (next power of two of `d_model`) in the module call path. +- Files: `needle/model/architecture.py:287-302` +- Cause: The matrix is represented explicitly even though the model description calls for an O(n log n) transform; repeated tracing/compilation multiplies memory and compile cost. +- Improvement path: Cache the matrix per dimension/dtype or implement a structured Walsh-Hadamard transform; benchmark compile and peak memory on the 45M-parameter preset. + +## Fragile Areas + +**CACT exporter supports only a subset of model configurations:** +- Files: `needle/model/export.py:102-118`, `needle/model/export.py:420-433` +- Why fragile: Export rejects split qk/v head dimensions, lexicon layers, and sliding-window layer patterns, while `TransformerConfig` can represent related fields. A checkpoint can train/load but fail only at deployment build time. +- Safe modification: Run `_geometry` validation before expensive quantization, and add fixture checkpoints for every supported/unsupported configuration with actionable migration messages. +- Test coverage: Existing build tests cover basic tiny exports only (`tests/test_build.py:13-37`); no negative geometry cases are exercised. + +**Shared temporary filenames in playground fine-tuning:** +- Files: `needle/playground/server.py:14-15`, `needle/playground/server.py:66-93` +- Why fragile: All sessions use fixed files in the process temp directory (`needle_playground_data.jsonl`, adapter, and output), so retries or another process can overwrite active artifacts. +- Safe modification: Allocate a per-job temporary directory, retain an explicit job id in status, and clean up only after download/expiry. +- Test coverage: No playground HTTP or concurrent fine-tune tests exist. + +## Scaling Limits + +**Single-process engine/session state:** +- Current capacity: One active native engine state and one `_FT` fine-tune job per Python process. +- Limit: Distinct weight sets cannot be unloaded/rebound safely, and simultaneous playground users share one conversation/engine lock and global status. +- Scaling path: Move model sessions/jobs behind worker processes with explicit job/session identifiers, or document a strict single-session deployment contract. + +## Dependencies at Risk + +**JAX/Flax/Metal compatibility:** +- Risk: The training extras leave most versions unconstrained while the Metal path relies on an exact JAX/JAXLIB pair and environment mutation before import. +- Impact: Backend upgrades can break compilation or silently change numerics; CPU/GPU/Metal behavior is not reproducible from the repository alone. +- Migration plan: Maintain tested constraints per backend, publish lockfiles, and run a small model smoke test for each supported Python/platform combination. + +## Missing Critical Features + +**Production request hardening:** +- Problem: The playground has no authentication, quotas, request size limits, cancellation, or durable job storage. +- Blocks: Safely exposing the demo server to a team/network and operating fine-tuning as a shared service. + +## Test Coverage Gaps + +**Native fetch, security, and concurrency paths:** +- What's not tested: Artifact digest/signature validation, malicious pickle rejection, tokenizer/engine cache failures, concurrent `Needle` calls, and playground endpoint authorization/resource limits. +- Files: `needle/agent/fetch.py`, `needle/model/tokenizer.py`, `needle/model/run.py`, `needle/playground/server.py`, `needle/__init__.py` +- Risk: Supply-chain, data-corruption, and denial-of-service regressions can ship unnoticed; most engine tests are skipped when the native library is absent (`tests/conftest.py:7-19`). +- Priority: High + +**Configuration/export edge cases:** +- What's not tested: Unknown config keys, invalid dimension combinations, empty `batch_generate`, and unsupported CACT geometry. +- Files: `needle/model/architecture.py`, `needle/model/run.py`, `needle/model/export.py`, `tests/test_build.py`, `tests/test_run.py` +- Risk: Failures occur late during training/export or only in production deployment. +- Priority: Medium + +--- + +*Concerns audit: 2026-08-31* diff --git a/.planning/codebase/CONVENTIONS.md b/.planning/codebase/CONVENTIONS.md new file mode 100644 index 0000000..21a3da8 --- /dev/null +++ b/.planning/codebase/CONVENTIONS.md @@ -0,0 +1,99 @@ +# Coding Conventions + +**Analysis Date:** 2026-08-31 + +## Naming Patterns + +**Files:** +- Python modules use lowercase `snake_case.py`, grouped by package responsibility, for example `needle/model/finetune.py` and `needle/agent/fetch.py`. +- Tests use `test_.py` and functions use `test_()`, for example `tests/test_render.py::test_encode_loss_mask_targets_only`. + +**Functions:** +- Public and private functions use `snake_case`; private implementation helpers begin with `_`, such as `_parse_array` in `needle/model/finetune.py` and `_library_path` in `needle/__init__.py`. +- Boolean/configuration helpers use descriptive predicates or accessors (`_engine_available`, `has_default`, `effective_kv_window`). + +**Variables:** +- Local variables and parameters use `snake_case`; short mathematical names are used in tensor code where shape context is clear (`B`, `T`, `D`, `q`, `k`, `v` in `needle/model/architecture.py`). +- Module constants use uppercase with underscores (`PAD_ID`, `EOS_ID`, `DEFAULT_BASE`, `LORA_TARGETS`). + +**Types:** +- Classes use `PascalCase` (`TransformerConfig`, `SimpleAttentionNetwork`, `SANTokenizer`, `Needle`). +- Type annotations use built-in generics where practical (`list`, `dict`) and `typing`/`Annotated` for Python 3.9-compatible unions and schema metadata, as shown in `needle/agent/tools.py` and `needle/environments/smart_home.py`. +- Dataclass configuration belongs in `@dataclass` classes; model hyperparameters are centralized in `needle/model/architecture.py::TransformerConfig`. + +## Code Style + +**Formatting:** +- No Black, Ruff, Flake8, isort, or formatter configuration is present in `pyproject.toml` or the repository root. Preserve the existing four-space indentation, blank-line grouping, and manually wrapped calls. +- Keep imports at module scope and group standard-library imports before third-party imports and local relative imports, following `needle/model/architecture.py` and `needle/__init__.py`. +- Use trailing commas in multiline calls/collections where the surrounding file does; keep lines readable rather than introducing a new formatter dependency. + +**Linting:** +- No lint command or enforced lint rules are configured. New code should still avoid unused imports, wildcard imports, mutable default arguments, and broad exception handling except at explicit process/network boundaries. +- The release workflow validates behavior with `pytest -q -m "not slow"` in `.github/workflows/release.yaml`; it does not run a linter. + +## Import Organization + +**Order:** +1. Standard library (`os`, `json`, `typing`, `dataclasses`, etc.). +2. Third-party dependencies (`numpy`, `jax`, `flax`, `pytest`, `pydantic`). +3. Local package imports, using relative imports inside `needle` (`from .tokenizer import ...`, `from . import quantize`). + +**Path Aliases:** +- No import aliases or package path aliases are configured. Use package imports such as `from needle.model...` in tests and relative imports within package modules. + +## Error Handling + +**Patterns:** +- Raise `ValueError` for invalid user/configuration data and incompatible checkpoint/export formats, for example `needle/model/run.py::load_checkpoint` and `needle/model/export.py::_geometry`. +- Raise `RuntimeError` when an external engine, tokenizer download, native call, or response envelope fails; preserve the original exception with `raise ... from e` where useful (`needle/model/tokenizer.py::get_tokenizer`, `needle/__init__.py::Needle._complete`). +- Catch narrowly when the failure is expected (`JSONDecodeError`, `OSError`, `EntryNotFoundError`). Broad `except Exception` is reserved for isolation boundaries such as tool execution in `Needle.run`, telemetry, and optional platform probing. +- Tool execution errors are converted into structured `{"error": ...}` results so one failing tool does not abort the agent loop (`needle/__init__.py::Needle.run`). +- Validate external schemas and arguments before acting; environment definitions encode bounds/enums using `needle.Field` and `typing.Literal` (`needle/environments/smart_home.py`). + +## Logging + +**Framework:** `print()` for CLI/progress output; no logging framework is configured. + +**Patterns:** +- CLI and download/training progress uses aligned `print` messages with labels such as `fetch`, `file`, `weights`, and `next` (`needle/cli.py`, `needle/model/run.py`, `needle/model/tokenizer.py`). +- Streaming generation writes incremental text directly to stdout and flushes (`needle/model/run.py::generate`). +- User-facing warnings use `warnings.warn` for tuned-weight confidence limitations (`needle/__init__.py::Needle.__init__`). +- Anonymous telemetry is isolated in `needle/_telemetry.py`; failures are swallowed there so instrumentation cannot break inference. + +## Comments + +**When to Comment:** +- Comment non-obvious runtime constraints, binary formats, backend workarounds, or algorithmic invariants. Examples include the Metal PJRT compatibility note in `needle/model/finetune.py` and quantization format documentation in `needle/model/export.py`. +- Keep comments close to the implementation and avoid narrating straightforward assignments. + +**JSDoc/TSDoc:** +- Python docstrings document public behavior and tool schemas. Function docstrings in environment modules include an overview and an `Args:` section consumed by `needle/agent/tools.py::build_schema`. +- Public helpers such as `needle.extract` and `needle.environments._harness.run_tests` have concise behavioral docstrings; private tensor helpers generally rely on names and nearby comments. + +## Function Design + +**Size:** +- Keep orchestration in small helpers and isolate serialization, tokenization, model math, and CLI dispatch in their existing modules. Large model routines may be compact tensor pipelines, but avoid mixing CLI parsing with numerical implementation. + +**Parameters:** +- Prefer explicit keyword arguments for configuration-heavy APIs and defaults that preserve current behavior (`Needle(..., max_new_tokens=256)`, `TransformerConfig`). +- Use annotations on public/tool-facing parameters; use `Annotated[..., needle.Field(...)]` for validation constraints and `Literal` for closed sets. + +**Return Values:** +- Return plain dictionaries/lists for JSON/native boundaries (`Needle.complete`, tool results, generated examples). +- Return typed Pydantic instances only when the caller supplies a Pydantic schema (`needle.extract`). +- Preserve array dtypes/shapes at numerical boundaries and use NumPy/JAX conversion explicitly rather than implicit Python coercion. + +## Module Design + +**Exports:** +- `needle/__init__.py` defines the public surface through `__all__` (`Needle`, `tool`, `Field`, `extract`, `__version__`). +- Model internals are imported from their focused modules; avoid adding engine, training, or quantization implementation to the package root. + +**Barrel Files:** +- There are no broad barrel modules. `needle/model/__init__.py`, `needle/agent/__init__.py`, and `needle/environments/__init__.py` provide lightweight package entry points/registries only. + +--- + +*Convention analysis: 2026-08-31* diff --git a/.planning/codebase/INTEGRATIONS.md b/.planning/codebase/INTEGRATIONS.md new file mode 100644 index 0000000..96ab8e4 --- /dev/null +++ b/.planning/codebase/INTEGRATIONS.md @@ -0,0 +1,84 @@ +# External Integrations + +**Analysis Date:** 2026-08-31 + +## APIs & External Services + +**Model and artifact distribution:** +- Hugging Face Hub (`Cactus-Compute/needle2`) - Publishes native engine wheels, base checkpoints, tokenizer files and optional tuned `.cact` archives. Used by `needle/agent/fetch.py`, `needle/model/run.py`, `needle/model/tokenizer.py`, and `needle/cli.py`. + - SDK/Client: `huggingface_hub` (`hf_hub_download`, `list_repo_files`, `HfApi`) + - Auth: public downloads require no configured secret; upload uses the Hugging Face client/token environment handled by `HfApi`. + +**Synthetic fine-tuning data:** +- OpenRouter (default `https://openrouter.ai/api/v1/chat/completions`) - Optional generation/augmentation of tool-calling JSONL examples in `needle/model/finetune.py` and the playground worker in `needle/playground/server.py`. + - SDK/Client: Python `urllib.request` with an OpenAI-compatible JSON request. + - Auth: `OPENROUTER_API_KEY` (Bearer token); endpoint can be changed with `OPENROUTER_URL`. + - Default model: `deepseek/deepseek-v4-flash` (`DEFAULT_MODEL` in `needle/model/finetune.py`). + +**Native inference runtime:** +- Cactus Needle engine - A platform-specific shared library fetched from Hugging Face and loaded with `ctypes` in `needle/__init__.py`. + - SDK/Client: C ABI functions `needle_init`, `needle_complete`, `needle_reset`, and `needle_load`. + - Auth: none; local binary and `.cact` bytes are supplied by the caller. + +## Data Storage + +**Databases:** +- None detected. Runtime state is in-memory Python objects and native engine state (`needle/__init__.py`, `needle/playground/server.py`). + +**File Storage:** +- Local filesystem only for checkpoints, LoRA adapters, `.cact` exports, JSONL datasets and playground downloads. +- Hugging Face cache is used for remote artifacts; native engine defaults to `~/.cache/cactus-needle//` (`needle/agent/fetch.py`, `needle/__init__.py`). +- Telemetry anonymous ID is stored at `~/.cactus_needle/telemetry_id` (`needle/_telemetry.py`). + +**Caching:** +- Hugging Face's local cache backs downloads; `tool_index_path` optionally persists tool embeddings keyed by schema/model fingerprint (documented in `doc/apis.md`). +- No Redis, Memcached or remote cache integration is present. + +## Authentication & Identity + +**Auth Provider:** +- No user authentication or identity provider is implemented. +- External credentials are limited to `OPENROUTER_API_KEY` for data synthesis and Hugging Face credentials used implicitly by `HfApi` when `--upload` is requested. +- Telemetry uses a locally generated random anonymous ID, not an account identity (`needle/_telemetry.py`). + +## Monitoring & Observability + +**Error Tracking:** +- None detected (no Sentry or hosted error tracker). + +**Logs:** +- CLI and training progress use `print()`/stderr; `needle/cli.py` installs an XLA stderr noise filter. +- Anonymous usage counts are sent asynchronously to a Supabase Edge Function endpoint (`needle/_telemetry.py`), containing function name, package/engine versions, OS/arch/Python and optional non-content properties; prompts and outputs are not sent. + +## CI/CD & Deployment + +**Hosting:** +- The package is installable from Python packaging infrastructure; model/engine artifacts are hosted on Hugging Face. No application hosting configuration is included. + +**CI Pipeline:** +- No CI workflow files were detected at repository root. Tests are run locally with Pytest (`pyproject.toml`, `tests/`). + +## Environment Configuration + +**Required env vars:** +- Normal inference: none if engine/checkpoint assets are cached or downloadable. +- Synthetic generation: `OPENROUTER_API_KEY`; optional `OPENROUTER_URL`. +- Telemetry opt-out: `NEEDLE_TELEMETRY=0` or `DO_NOT_TRACK=1`. +- Air-gapped deployment: optionally `HF_HUB_OFFLINE=1` and/or `NEEDLE_LIB_PATH`. + +**Secrets location:** +- Supplied through process environment or Hugging Face's standard local credential configuration; no secret files are read by the package. Never commit API keys to datasets or source. + +## Webhooks & Callbacks + +**Incoming:** +- None. The local playground exposes HTTP endpoints (`/complete`, `/reset`, `/finetune`, `/load-model`) but they are local server routes, not third-party webhooks (`needle/playground/server.py`). + +**Outgoing:** +- OpenRouter HTTPS POST requests for synthetic examples (`needle/model/finetune.py`). +- Supabase telemetry HTTPS POST requests to the configured/default endpoint (`needle/_telemetry.py`). +- Hugging Face Hub download/upload API requests for model artifacts (`needle/agent/fetch.py`, `needle/model/run.py`, `needle/model/tokenizer.py`, `needle/cli.py`). + +--- + +*Integration audit: 2026-08-31* diff --git a/.planning/codebase/STACK.md b/.planning/codebase/STACK.md new file mode 100644 index 0000000..dc5dd24 --- /dev/null +++ b/.planning/codebase/STACK.md @@ -0,0 +1,81 @@ +# Technology Stack + +**Analysis Date:** 2026-08-31 + +## Languages + +**Primary:** +- Python (>=3.9) - Public package API, CLI, inference orchestration, training and export in `needle/`. +- JAX/Python numerical code - Neural-network definition, decoding, quantization and LoRA updates in `needle/model/architecture.py`, `needle/model/decode.py`, `needle/model/quantize.py`, and `needle/model/finetune.py`. + +**Secondary:** +- C/C++ shared library (prebuilt, outside this repository) - Native inference engine loaded through `ctypes` in `needle/__init__.py`; the Python package does not compile it locally. +- HTML/CSS/JavaScript - Browser playground assets in `needle/playground/index.html`, `needle/playground/app.js`, and `needle/playground/style.css`. + +## Runtime + +**Environment:** +- CPython 3.9 or newer (declared by `requires-python` in `pyproject.toml`). +- Native engine selected by platform/architecture and loaded with `ctypes.CDLL` from `needle/__init__.py`. + +**Package Manager:** +- `pip`/PEP 517 setuptools build (`setuptools>=68.0`, `pyproject.toml`). +- Lockfile: missing; dependency versions are specified as unconstrained or minimum versions in `pyproject.toml` and `requirements*.txt`. + +## Frameworks + +**Core:** +- Flax Linen (`flax>=0.10.2`, training extra) - Transformer/Simple Attention Network modules in `needle/model/architecture.py`. +- JAX (`jax`, `jaxlib`, training extra) - Array operations, autodiff, JIT and accelerator execution in `needle/model/`. +- SentencePiece (`sentencepiece`, training extra) - Tokenizer model loading/encoding in `needle/model/tokenizer.py`. +- Python standard library `argparse` and `http.server` - CLI in `needle/cli.py` and local playground server in `needle/playground/server.py`. + +**Testing:** +- Pytest (`pytest`, test extra) - Tests under `tests/`, configured by `tool.pytest.ini_options` in `pyproject.toml`. +- Pydantic (`pydantic`, test extra/runtime-optional) - Typed extraction schemas and test fixtures; integrated dynamically in `needle/agent/tools.py` and `needle/__init__.py`. + +**Build/Dev:** +- Setuptools package discovery and package data, configured in `pyproject.toml`. +- Cactus Quants export format (`.cact`) implemented in `needle/model/export.py` and quantization utilities in `needle/model/quantize.py`. + +## Key Dependencies + +**Critical:** +- `huggingface_hub` - Downloads the base checkpoint, tokenizer, native engine wheels and published `.cact` archives (`needle/model/run.py`, `needle/model/tokenizer.py`, `needle/agent/fetch.py`, `needle/cli.py`). +- `jax`/`jaxlib` - Required for checkpoint inference utilities and all fine-tuning/export paths (`needle/model/run.py`, `needle/model/finetune.py`). +- `flax` - Parameterized neural-network modules and tree traversal for LoRA (`needle/model/architecture.py`, `needle/model/finetune.py`). +- `optax` - AdamW, warmup/cosine schedule, gradient clipping and loss helpers in `needle/model/finetune.py`. +- `sentencepiece` - Required to train or load the model tokenizer (`needle/model/tokenizer.py`). + +**Infrastructure:** +- `numpy` - Checkpoint conversion, array serialization and export packing across `needle/model/`. +- `pydantic` (optional) - Converts `BaseModel` schemas into tool contracts and typed extraction results (`needle/agent/tools.py`, `needle/__init__.py`). + +## Configuration + +**Environment:** +- `NEEDLE_LIB_PATH` overrides native engine lookup; otherwise the package directory and `~/.cache/cactus-needle//` are searched (`needle/__init__.py`). +- `HF_HUB_OFFLINE=1` prevents Hugging Face network access for air-gapped operation (documented in `doc/apis.md`). +- `OPENROUTER_API_KEY` authorizes optional synthetic data generation; `OPENROUTER_URL` overrides the OpenAI-compatible endpoint (`needle/model/finetune.py`). +- `NEEDLE_HF_REPO` selects the Hugging Face destination for `needle build --upload` (`needle/model/finetune.py`). +- `NEEDLE_TELEMETRY=0` or `DO_NOT_TRACK=1` disables anonymous telemetry; `CI` also disables it (`needle/_telemetry.py`). +- `ENABLE_PJRT_COMPATIBILITY` is set automatically on macOS before JAX initialization for the Metal plugin (`needle/model/finetune.py`). + +**Build:** +- `pyproject.toml` defines package metadata, optional extras (`train`, `gpu`, `metal`, `test`), console script `needle = needle.cli:main`, package data and pytest paths. +- `requirements.txt` contains runtime installation; `requirements-train.txt` extends it with JAX/Flax/Optax/SentencePiece training dependencies. +- Model/tokenizer assets (`*.model`, `*.vocab`) are included as package data for `needle.model`; playground static assets are included for `needle.playground`. + +## Platform Requirements + +**Development:** +- Python 3.9+ and a platform-supported JAX backend. Install `cactus-needle[train,gpu]` for NVIDIA CUDA 12 or `cactus-needle[train,metal]` for Apple Silicon Metal (pins JAX 0.4.38). +- Network access is needed once to fetch the native engine/checkpoint/tokenizer from Hugging Face unless assets are pre-populated in cache. + +**Production:** +- A supported native engine binary for the target platform (downloaded by `needle fetch` or `needle download `). +- Runtime RAM target is approximately 28 MB for the bundled 14 MB Needle 2 engine/weights, as described in `README.md`; tuned `.cact` archives use the same engine. + +--- + +*Stack analysis: 2026-08-31* diff --git a/.planning/codebase/STRUCTURE.md b/.planning/codebase/STRUCTURE.md new file mode 100644 index 0000000..0af8c1e --- /dev/null +++ b/.planning/codebase/STRUCTURE.md @@ -0,0 +1,149 @@ +# Codebase Structure + +**Analysis Date:** 2026-08-31 + +## Directory Layout + +```text +needle/ +├── __init__.py # Public API and ctypes native-engine facade +├── cli.py # Console command parser/router +├── _telemetry.py # Anonymous usage tracking/opt-out +├── agent/ +│ ├── tools.py # Callable/Pydantic -> JSON schema conversion +│ └── fetch.py # Platform engine download/cache +├── model/ +│ ├── architecture.py # Flax Simple Attention Network and config +│ ├── decode.py # Cached JAX decode implementation +│ ├── run.py # Checkpoint loading and reference generation +│ ├── finetune.py # Data synthesis, JSONL loading, LoRA, build CLI +│ ├── quantize.py # Fake-QAT and CQ quantization utilities +│ ├── export.py # .cact binary writer/reader and tokenizer blob +│ └── tokenizer.py # SentencePiece tokenizer and special IDs +├── environments/ # Tool surfaces and acceptance suites +└── playground/ # Browser assets and threaded HTTP server +tests/ # Pytest unit, integration, and slow training tests +doc/ # API, environment, and fine-tuning guides +assets/ # README architecture/frontier images +``` + +## Directory Purposes + +**`needle/`:** Installable package. Keep public API changes in `needle/__init__.py` and CLI changes in `needle/cli.py`; avoid importing JAX from the runtime facade so production installs remain lightweight. + +**`needle/agent/`:** Runtime integration boundary. `tools.py` owns schema semantics and `fetch.py` owns platform/Hugging Face artifact retrieval. Native shared libraries are downloaded into the user cache, not committed here. + +**`needle/model/`:** Training/reference implementation. `architecture.py` is the source of truth for parameter names and geometry; `finetune.py` assumes those names when selecting LoRA targets; `export.py` assumes the same tree when serializing. + +**`needle/environments/`:** Product-like examples. Each module defines `SYSTEM`, `TOOLS`, and `TEST_CASES`; `_harness.py` supplies lazy agent construction and exact acceptance scoring. + +**`needle/playground/`:** Self-contained static frontend plus `server.py`. It stores uploaded/tuned temporary files in the system temp directory and runs fine-tuning in a daemon thread. + +**`tests/`:** Pytest suite configured by `pyproject.toml`. `tests/conftest.py` supplies shared fixtures such as tiny checkpoints; engine-dependent tests use the `requires_engine` marker. + +**`doc/`:** User-facing operational guidance. Read `doc/apis.md` for runtime contracts, `doc/finetuning.md` for dataset/training details, and `doc/environments.md` for tool-surface design. + +## Key File Locations + +**Entry Points:** +- `needle/__init__.py`: `Needle`, `extract`, `tool`, and `Field` public API. +- `needle/cli.py`: `main()` and all console subcommands. +- `needle/playground/server.py`: local HTTP server entry point. +- `needle/model/run.py`: reference checkpoint generation entry point. + +**Configuration:** +- `pyproject.toml`: package metadata, Python >=3.9, dependencies/extras, console script, package data, pytest config. +- `needle/model/architecture.py`: `TransformerConfig` and `PRESETS` model geometry. +- `needle/agent/fetch.py`: `HF_REPO`, `ENGINE_VERSION`, platform tags. +- `needle/model/tokenizer.py`: tokenizer path and special token IDs. + +**Core Logic:** +- `needle/model/architecture.py`: model block, engram, mHC, masks, KV budgeting. +- `needle/model/decode.py`: KV-cached forward pass and generation internals. +- `needle/__init__.py`: native completion and agent loop. +- `needle/agent/tools.py`: schema reflection and validation metadata. + +**Training/Deployment:** +- `needle/model/finetune.py`: JSONL encoding, LoRA optimizer, adapter serialization, `.cact` build orchestration. +- `needle/model/export.py`: binary archive contract and round-trip reader. +- `needle/model/quantize.py`: CQ/fake quantization algorithms. + +**Testing:** +- `tests/test_inference.py`: native completion, extraction, loops, and multiple-agent behavior. +- `tests/test_finetune.py`: adapter contents and merge/build integration. +- `tests/test_build.py`: archive creation, bit widths, and projection round trips. +- `tests/test_run.py`: prompt rendering and reference CLI behavior. +- `tests/test_tools.py`: schema reflection and constraints. +- `tests/test_environments.py`: environment acceptance contracts. + +## Naming Conventions + +**Files:** +- Lowercase snake_case for Python modules (`finetune.py`, `data_capture.py`). +- Leading underscore for private helpers/modules (`_harness.py`, `_telemetry.py`). +- Tests use `test_.py`; fixtures and test data are kept under `tests/` or temporary paths. + +**Directories:** +- Lowercase package names (`agent`, `model`, `environments`, `playground`). +- No source-generated build directory is required; checkpoints and `.cact` outputs are user-selected paths (README defaults to `checkpoints/`). + +**Python symbols:** +- Classes use PascalCase (`SimpleAttentionNetwork`, `TransformerConfig`, `Needle`). +- Functions and variables use snake_case (`load_checkpoint`, `merge_lora`). +- Constants use uppercase (`BOS_ID`, `ENGINE_VERSION`, `LORA_TARGETS`). + +## Where to Add New Code + +**New Public Runtime Feature:** +- Primary code: `needle/__init__.py` for API behavior; add schema support in `needle/agent/tools.py` when required. +- Tests: `tests/test_inference.py` or `tests/test_tools.py`. +- Documentation: `doc/apis.md` and the relevant README section. + +**New Model Layer/Head:** +- Implementation: `needle/model/architecture.py`, with config fields in `TransformerConfig`. +- Reference execution: update `needle/model/decode.py` if cached/native-equivalent inference needs the layer. +- Export: update tensor order/header handling in `needle/model/export.py` and add round-trip coverage to `tests/test_build.py`. + +**New Fine-Tune Behavior:** +- Data/rendering/optimizer: `needle/model/finetune.py`. +- Quantization changes: `needle/model/quantize.py`. +- Tests: `tests/test_finetune.py`; use the `slow` marker for accelerator-dependent work. + +**New CLI Command:** +- Parser and dispatch: `needle/cli.py`. +- Implementation: keep command-specific logic in `needle/model/` or `needle/agent/`, not in the parser function. +- Test: add argument/behavior coverage under `tests/test_*.py`. + +**New Environment:** +- Add `needle/environments/.py` defining `SYSTEM`, `TOOLS`, and `TEST_CASES`. +- Reuse `needle.environments._harness.agent_for` and `run_tests`; add a focused module test in `tests/test_environments.py`. + +**Utilities:** +- Shared schema helpers belong in `needle/agent/tools.py`; model math belongs in `needle/model/architecture.py` or `needle/model/quantize.py` according to ownership. + +## Special Directories + +**`checkpoints/` (runtime-created):** +- Purpose: default destination for downloaded/base checkpoints and LoRA adapters. +- Generated: Yes. +- Committed: No by convention; use explicit paths for reproducible artifacts. + +**User cache `~/.cache/cactus-needle//`:** +- Purpose: downloaded native engine library. +- Generated: Yes. +- Committed: No. + +**`needle/model` package data:** +- `*.model` and `*.vocab` tokenizer files are declared in `pyproject.toml` and may be downloaded from Hugging Face when absent. + +**System temp directory:** +- Playground uploads, generated JSONL, adapters, and tuned `.cact` files are placed in `tempfile.gettempdir()` by `needle/playground/server.py`; they are ephemeral and not source-controlled. + +**`assets/`:** +- Purpose: README visual assets (`banner.png`, `architecture.png`, frontier images). +- Generated: No for normal development. +- Committed: Yes. + +--- + +*Structure analysis: 2026-08-31* diff --git a/.planning/codebase/TESTING.md b/.planning/codebase/TESTING.md new file mode 100644 index 0000000..3040185 --- /dev/null +++ b/.planning/codebase/TESTING.md @@ -0,0 +1,137 @@ +# Testing Patterns + +**Analysis Date:** 2026-08-31 + +## Test Framework + +**Runner:** +- `pytest` is configured in `pyproject.toml` with `testpaths = ["tests"]`. +- The only registered marker is `slow`, documented as end-to-end JAX build/finetune tests that initialize a tiny model. +- CI runs `pytest -q -m "not slow"` in `.github/workflows/release.yaml` after installing `.[test,train]`. + +**Assertion Library:** +- Use plain Python `assert` for scalar/structure checks and `numpy.testing.assert_allclose` for numerical tensors (`tests/test_lora.py`, `tests/test_build.py`). +- `pytest.raises(..., match=...)` verifies exception type and message; `pytest.mark.skipif` gates optional runtime features. + +**Run Commands:** +```bash +pytest -q # Run the complete suite +pytest -q -m "not slow" # Fast/CI suite without JAX fine-tune and build tests +pytest -q -m slow # Explicitly run slow model-training/export tests +pytest -q tests/test_tools.py # Run one focused module +pytest -q tests/test_lora.py::test_merge_lora_adds_scaled_delta # Run one test +``` + +## Test File Organization + +**Location:** +- Tests are in the top-level `tests/` directory, separate from implementation modules. There are no co-located `needle/**/test_*.py` files. +- `tests/conftest.py` contains shared engine detection and the session-scoped `tiny_checkpoint` fixture. + +**Naming:** +- Files are `test_.py`; tests are descriptive `test_()` functions. Test classes are not used. + +**Structure:** +``` +tests/ +├── conftest.py # shared fixtures and optional-engine marker +├── test_tools.py # schema/decorator unit tests +├── test_lora.py # LoRA tensor math unit tests +├── test_render.py # prompt/token/loss-mask unit tests +├── test_generate.py # data synthesis and JSON parsing tests +├── test_run.py, test_weights.py # checkpoint/CLI/runtime behavior +├── test_build.py, test_finetune.py # slow JAX/export integration tests +└── test_inference.py, test_environments.py # native-engine and environment suites +``` + +## Test Structure + +**Suite Organization:** +```python +def test_merge_lora_adds_scaled_delta(): + import jax.numpy as jnp + from needle.model.finetune import merge_lora + + params = {"w": {"kernel": jnp.zeros((3, 4))}} + lora = {("w", "kernel"): {"A": jnp.ones((3, 2)), "B": jnp.ones((2, 4))}} + merged = merge_lora(params, lora, scale=0.5) + np.testing.assert_allclose(np.asarray(merged["w"]["kernel"]), np.ones((3, 4))) +``` + +**Patterns:** +- Keep each test focused on one behavior and construct minimal local inputs. +- Import implementation modules inside tests when optional dependencies or monkeypatch setup must occur first (`tests/test_finetune.py`, `tests/test_tools.py`). +- Use fixtures for reusable state (`tok`, `tiny_checkpoint`, `engine`, `tuned`), with `scope="session"` only for expensive immutable setup. +- Assert externally visible contracts: exact schema fields, output types, file existence, checkpoint metadata, array shape/dtype, or allowed response types. + +## Mocking + +**Framework:** +- Use pytest's `monkeypatch` fixture; no standalone mock library is configured. + +**Patterns:** +```python +monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") +monkeypatch.setattr(finetune, "generate_examples", fake_generator) +monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) +``` + +- Native engine behavior is replaced with a small `_Stub` object in `tests/test_weights.py`; its methods record calls and populate a JSON envelope in the provided ctypes buffer. +- Use `warnings.catch_warnings()` to suppress expected tuned-weight warnings while testing behavior. + +**What to Mock:** +- Mock network/API calls (`needle.model.finetune._openrouter`), environment variables, platform/native library discovery, and ctypes engine bindings when testing control flow. +- Mock expensive generated-data calls with deterministic rows and test deduplication/termination separately (`tests/test_generate.py`). + +**What NOT to Mock:** +- Keep pure schema, rendering, LoRA math, quantization, and serialization logic real. Slow tests intentionally initialize a tiny JAX model rather than mocking model parameters (`tests/conftest.py::tiny_checkpoint`). + +## Fixtures and Factories + +**Test Data:** +- Use `tmp_path`/`tmp_path_factory` for JSONL, checkpoint, adapter, and `.cact` files; tests write only ephemeral files. +- Build minimal dictionaries inline for tool schemas and calls. `tests/test_finetune.py::TOOLS` and `_write_data` are representative fixtures/factories. +- `tiny_checkpoint` creates a 2-layer, 64-dimension `SimpleAttentionNetwork`, converts leaves to NumPy, and serializes a format-v2 pickle for training/export tests. + +**Location:** +- Shared fixtures belong in `tests/conftest.py`; feature-specific helpers stay at the top of each test module (`_finetune_args`, `_build_args`, `_parse_array` inputs). + +## Coverage + +**Requirements:** +- No coverage target or coverage configuration is enforced. CI only runs the non-slow pytest selection. + +**View Coverage:** +```bash +pytest --cov=needle --cov-report=term-missing +``` +This command is available only when a coverage plugin is installed; it is not declared in `pyproject.toml`. + +## Test Types + +**Unit Tests:** +- Most tests cover deterministic helpers without the native engine: tool schema generation (`tests/test_tools.py`), prompt rendering/token masks (`tests/test_render.py`), LoRA operations (`tests/test_lora.py`), quantized export round trips (`tests/test_build.py`), and CLI/checkpoint parsing (`tests/test_run.py`, `tests/test_fetch.py`). + +**Integration Tests:** +- `tests/test_finetune.py` runs one-epoch JAX LoRA training and then exports/reads a `.cact`; marked `slow`. +- `tests/test_environments.py` validates every declared environment's tool surface, frozen cases, and direct tool execution. The smoke inference test is engine-gated. + +**E2E Tests:** +- `tests/test_inference.py` is marked with `requires_engine`, which skips the module when the platform-specific native engine is absent. It exercises `Needle.complete`, `run`, extraction, and multiple agents against the installed engine. + +## Common Patterns + +**Async Testing:** +- No async test functions or async fixtures are present. Concurrent generation is tested synchronously by monkeypatching worker-facing functions (`tests/test_generate.py`). + +**Error Testing:** +```python +with pytest.raises(RuntimeError, match="cannot unload"): + needle.Needle(tools="[]") +``` +- Assert invalid/missing input behavior with `pytest.raises`, and use exact/regex message fragments for important operational diagnostics (`tests/test_weights.py`, `tests/test_generate.py`). +- Engine-optional tests should use the shared `requires_engine` marker rather than failing on machines without the downloaded native library. + +--- + +*Testing analysis: 2026-08-31* diff --git a/.planning/config.json b/.planning/config.json new file mode 100644 index 0000000..46cd8ae --- /dev/null +++ b/.planning/config.json @@ -0,0 +1,92 @@ +{ + "model_profile": "inherit", + "commit_docs": true, + "parallelization": true, + "search_gitignored": false, + "brave_search": false, + "firecrawl": false, + "exa_search": false, + "tavily_search": false, + "ref_search": false, + "perplexity": false, + "jina": false, + "git": { + "branching_strategy": "none", + "create_tag": true, + "phase_branch_template": "gsd/phase-{phase}-{slug}", + "milestone_branch_template": "gsd/{milestone}-{slug}", + "quick_branch_template": null + }, + "workflow": { + "research": true, + "plan_check": true, + "verifier": true, + "nyquist_validation": true, + "auto_advance": false, + "node_repair": true, + "node_repair_budget": 2, + "ui_phase": true, + "ui_safety_gate": true, + "ai_integration_phase": true, + "api_coverage_gate": true, + "human_verify_mode": "end-of-phase", + "context_guard_mode": "warn", + "text_mode": false, + "research_before_questions": false, + "discuss_mode": "discuss", + "skip_discuss": false, + "code_review": true, + "code_review_depth": "standard", + "code_review_command": null, + "pattern_mapper": true, + "plan_bounce": false, + "plan_bounce_script": null, + "plan_bounce_passes": 2, + "auto_prune_state": false, + "post_planning_gaps": true, + "security_enforcement": true, + "security_asvs_level": 1, + "security_block_on": "high" + }, + "ship": { + "pr_body_sections": [ + { + "heading": "User Stories & Acceptance Criteria", + "enabled": false, + "source": "REQUIREMENTS.md ## User Stories || REQUIREMENTS.md ## Acceptance Criteria", + "fallback": "- Acceptance criteria are covered by the linked requirements and verification evidence." + }, + { + "heading": "Risks & Dependencies", + "enabled": true, + "source": "PLAN.md ## Risks || PLAN.md ## Dependencies", + "fallback": "- No known high-risk rollout dependencies." + }, + { + "heading": "Success Metrics & Release Criteria", + "enabled": false, + "source": "REQUIREMENTS.md ## Definition of Done || VERIFICATION.md ## Release Criteria", + "fallback": "- Release when automated verification and required manual checks pass." + }, + { + "heading": "Stakeholder Review & Approval", + "enabled": false, + "template": "- Product owner approval pending for {phase_name}." + } + ] + }, + "hooks": { + "context_warnings": true + }, + "project_code": null, + "phase_naming": "sequential", + "agent_skills": {}, + "claude_md_path": "./.claude/CLAUDE.md", + "plan_review": { + "source_grounding": false, + "source_grounding_authority": "grep" + }, + "resolve_model_ids": "omit", + "mode": "yolo", + "granularity": "standard" +} diff --git a/.planning/onboarding/SUMMARY.md b/.planning/onboarding/SUMMARY.md new file mode 100644 index 0000000..b82a1fd --- /dev/null +++ b/.planning/onboarding/SUMMARY.md @@ -0,0 +1,61 @@ +# Onboarding Summary + +## Project State +- PROJECT.md: present +- REQUIREMENTS.md: present +- ROADMAP.md: present +- STATE.md: present + +## Codebase Context +- Brownfield repo: yes +- Map readiness: complete +- Codebase map: `.planning/codebase/` (complete codebase map) +- Fast map available: yes + +## Docs Context +- Existing ADR/PRD/SPEC/RFC candidates: 0 + +## What This Project Contains + +Needle is a Python package with two main execution paths: + +- Production inference: `needle.Needle` loads a platform native engine through + ctypes and runs against a `.cact` archive. +- Reference/training: JAX/Flax code defines the Simple Attention Network, + decoding, LoRA fine-tuning, quantization, and export. + +The CLI in `needle/cli.py` covers fetching assets, inference, data generation, +fine-tuning, building exports, downloading engines, and the playground. + +## How To Start + +1. Read `.planning/codebase/STACK.md` for Python extras, backends, assets, and + environment variables. +2. Read `.planning/codebase/ARCHITECTURE.md` for model structure and native vs + JAX data flow. +3. Use the first roadmap phase to turn the current README/docs into a tested + Chinese-first install and inference quickstart. +4. Continue through LoRA and `.cact` deployment only after the first inference + path is verified. + +## Fine-Tuning Entry Points + +- `needle/model/finetune.py`: JSONL rendering, masked causal loss, LoRA update, + merge, and build orchestration. +- `doc/finetuning.md`: existing fine-tuning notes to reconcile with the new + beginner tutorial. +- `tests/test_finetune.py` and `tests/test_build.py`: executable behavior and + export checks to reuse for documentation verification. + +## Important Constraints + +- Native state is process-global; isolate base and tuned agents when needed. +- `.cact` tensor order, geometry, tokenizer vocabulary, and engine version must + remain aligned. +- Treat downloaded native artifacts and pickle checkpoints as trusted inputs + only; the local playground is not an authenticated production server. +- Git commits require a writable `.git` index; this environment now has a + successful initialization commit. + +## Recommended Next Step +- `$gsd-manager` (then `$gsd-discuss-phase 1` or `$gsd-plan-phase 1`) diff --git a/.planning/phases/01-install-and-first-inference/01-01-PLAN.md b/.planning/phases/01-install-and-first-inference/01-01-PLAN.md new file mode 100644 index 0000000..bd1074e --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-01-PLAN.md @@ -0,0 +1,85 @@ +--- +phase: 01-install-and-first-inference +plan: 01 +type: execute +wave: 1 +depends_on: [] +files_modified: + - README.md + - doc/installation.md +autonomous: true +user_setup: [] +must_haves: + truths: + - "A beginner can create a Python 3.9+ uv environment and install the editable package with train/test extras." + - "A reader knows what the native engine, checkpoint, tokenizer, and .cact files are before downloading assets." + - "A reader can explicitly fetch assets with needle fetch, inspect the default cache, and perform an offline load check." + artifacts: + - path: "doc/installation.md" + provides: "Chinese-first CPU installation, asset/cache, offline, and backend matrix guide" + contains: "uv venv" + - path: "README.md" + provides: "Short linked entry path to the installation guide" + contains: "needle fetch" + key_links: + - from: "doc/installation.md" + to: "pyproject.toml" + via: "uv editable install command uses declared extras" + - from: "doc/installation.md" + to: "needle/agent/fetch.py" + via: "fetch/cache behavior and platform selection are documented from implementation" +--- + +# Phase 1 Plan 01: Install, Assets, and CPU Baseline + +## Objective + +- **What:** A Chinese-first installation and asset guide for a clean CPU-first environment. +- **Why:** Every later inference and training tutorial depends on a reproducible environment and known artifact locations. +- **Output:** `doc/installation.md` plus a concise README entry link. + +## Context + +@.planning/PROJECT.md +@.planning/ROADMAP.md +@.planning/STATE.md +@.planning/phases/01-install-and-first-inference/01-CONTEXT.md +@.planning/codebase/STACK.md +@.planning/codebase/CONVENTIONS.md + +## Tasks + + + Write the uv CPU installation guide + doc/installation.md + Document prerequisites, `uv venv`, editable `uv pip install -e ".[train,test]"`, Python 3.9+ support, CPU baseline, optional GPU/Metal TODO note, telemetry opt-out, and expected command output. Keep commands exact and explain each dependency group. + Run the documented environment commands in a clean temporary environment, then run `python -c "import needle"` and `uv pip check`. Record the actual Python and package versions used. + A beginner can follow the guide without guessing which uv command or extras to choose, and the guide states the CPU-only Phase 1 acceptance boundary. + + + + Document explicit asset fetch, cache, and offline verification + doc/installation.md + Explain native engine, checkpoint, tokenizer, and `.cact` roles before commands. Use `needle fetch` as the primary command, describe `needle download ` as an advanced alternative, document the default cache inspection/cleanup approach, and add an online-first then `HF_HUB_OFFLINE=1` verification flow. State that missing assets should be fixed by running fetch rather than relying on implicit download. + Run `needle fetch` on the CPU environment, confirm expected cached files/version metadata, set `HF_HUB_OFFLINE=1`, and execute the documented load/inference check. Do not add a checksum command unless the implementation provides one. + The asset lifecycle and offline check are explicit, reproducible, and linked to the current fetch implementation. + + + + Add the installation entry link to README + README.md + Add a short Chinese-first “从安装开始” entry near the existing setup material that links to `doc/installation.md` and preserves the existing project positioning and API examples. + Check links from the repository root and render Markdown locally or with the project’s existing documentation check if available. + README readers reach the canonical installation guide without duplicate conflicting commands. + + +## Verification + +1. In a clean CPU environment, complete the documented uv install and `uv pip check`. +2. Run `needle fetch`, inspect the default cache, enable `HF_HUB_OFFLINE=1`, and complete the documented load check. +3. Confirm README links to the guide and no command claims GPU/Metal support was verified in this phase. + +## Success Criteria + +- [ ] Clean CPU setup and explicit asset/offline verification are reproducible from the guide. +- [ ] README and installation docs agree on package extras, commands, cache behavior, and Phase 1 scope. diff --git a/.planning/phases/01-install-and-first-inference/01-01-SUMMARY.md b/.planning/phases/01-install-and-first-inference/01-01-SUMMARY.md new file mode 100644 index 0000000..d184afd --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-01-SUMMARY.md @@ -0,0 +1,140 @@ +--- +phase: 01-install-and-first-inference +plan: 01 +subsystem: documentation +tags: [uv, cpu, installation, huggingface, offline] + +requires: [] +provides: + - Chinese-first uv installation guide for Python 3.9+ CPU environments + - Explicit native engine asset fetch, cache inspection, and offline verification flow + - README entry point linking source users to the canonical installation guide +affects: [01-02-inference, fine-tuning, onboarding] + +tech-stack: + added: [] + patterns: [uv editable installs, explicit fetch-before-inference, HF_HUB_OFFLINE verification] + +key-files: + created: [doc/installation.md] + modified: [README.md] + +key-decisions: + - "Use uv venv and editable uv pip install -e \".[train,test]\" as the source checkout baseline." + - "Fetch only the platform native engine in Phase 1; explain checkpoint, tokenizer, and .cact roles for later workflows." + - "Verify online fetch followed by HF_HUB_OFFLINE=1 loading on CPU, without asserting exact generated text." + +requirements-completed: [] + +coverage: + - id: D1 + description: "Chinese-first CPU installation guide with uv environment creation and train/test extras" + verification: + - kind: integration + ref: "uv venv + uv pip install -e .\"[train,test]\" in /tmp/needle-uv-ddSeBK" + status: pass + - kind: unit + ref: "uv pip check --python /tmp/needle-uv-ddSeBK/.venv/bin/python" + status: pass + human_judgment: false + - id: D2 + description: "Explicit engine fetch, cache inspection, and HF_HUB_OFFLINE inference check" + verification: + - kind: e2e + ref: "HOME=/tmp/needle-home-03UaCo /tmp/needle-uv-ddSeBK/.venv/bin/needle fetch" + status: pass + - kind: e2e + ref: "HF_HUB_OFFLINE=1 Needle.complete('hello', max_new_tokens=16)" + status: pass + human_judgment: false + - id: D3 + description: "README Chinese installation entry links to doc/installation.md and preserves API quickstart" + verification: + - kind: other + ref: "git diff --check and README relative link inspection" + status: pass + human_judgment: false + +duration: 14min +completed: 2026-08-31 +status: complete +--- + +# Phase 1 Plan 1: Install, Assets, and CPU Baseline Summary + +**uv-based CPU onboarding with explicit Needle engine fetching, cache/offline checks, and a Chinese README entry path** + +## Performance + +- **Duration:** 14 min +- **Started:** 2026-08-31T13:48:00Z +- **Completed:** 2026-08-31T14:04:00Z +- **Tasks:** 3 +- **Files modified:** 2 (plus this summary and the Windows ledger) + +## Accomplishments + +- Added `doc/installation.md` covering Python 3.9+, `uv venv`, editable `.[train,test]` installation, CPU scope, telemetry opt-out, and expected checks. +- Documented the distinction between native engine, JAX checkpoint, SentencePiece tokenizer, and `.cact` archive, including default cache inspection and cleanup. +- Added an online-first `needle fetch` then `HF_HUB_OFFLINE=1` verification flow and a Chinese-first README entry that points to the guide. + +## Task Commits + +1. **Write the uv CPU installation guide** - `fd0202a` (`docs`) +2. **Document explicit asset fetch, cache, and offline verification** - `b174059` (`docs`) +3. **Add the installation entry link to README** - `786d75d` (`docs`) +4. **Auto-fix temporary README断链** - `4819483` (`fix`) + +## Files Created/Modified + +- `doc/installation.md` - Chinese-first setup, artifact lifecycle, cache, fetch, and offline CPU verification instructions. +- `README.md` - Source-checkout installation entry with the shortest uv/fetch path and published-package distinction. +- `.planning/WINDOWS.md` - Broken-windows ledger entry for intentionally deferred CUDA/Metal verification. + +## Decisions Made + +- Keep the repository's declared Python 3.9+ range and install the complete training/test extras for future phases. +- Treat `needle fetch` as an explicit prerequisite; do not rely on implicit downloads during the quickstart. +- Keep CUDA and Metal out of Phase 1 acceptance while naming them as later work. + +## Deviations from Plan + +### Auto-fixed Issues + +**1. [Rule 1 - Bug] Avoided a README link to a not-yet-created inference guide** + +- **Found during:** Task 3 (Add the installation entry link to README) +- **Issue:** The initial entry linked `doc/inference.md`, which belongs to dependent plan 01-02 and did not yet exist. +- **Fix:** Removed that link while retaining the instruction to continue to the later inference guide; plan 01-02 can add its link when the file is created. +- **Files modified:** `README.md` +- **Verification:** `git diff --check`; installation link resolves from repository root. +- **Committed in:** `4819483` + +**Total deviations:** 1 auto-fixed (Rule 1) +**Impact on plan:** No scope creep; the plan's required installation link remains valid and no broken link is left behind. + +## Issues Encountered + +- The first isolated fetch command selected an empty virtual-environment path due to an overly shallow `find`; rerunning with the known environment path succeeded. No repository change was needed. +- Hugging Face emitted an unauthenticated-rate-limit warning during fetch; public assets downloaded successfully, so no authentication gate was required. + +## Known Stubs + +- `doc/installation.md:5` - CUDA and Metal installation/runtime verification are intentionally deferred to a later phase, matching the Phase 1 CPU-only boundary. + +## User Setup Required + +None - no external service configuration is required for this CPU baseline. + +## Next Phase Readiness + +The source checkout can be installed reproducibly with uv, the native CPU engine was fetched into the documented default cache, and offline loading was verified. Plan 01-02 can add the first fixed CLI inference example and its documentation smoke tests; it should restore a README link to `doc/inference.md` after creating that file. + +--- +*Phase: 01-install-and-first-inference* +*Completed: 2026-08-31* + +## Self-Check: PASSED + +- `01-01-SUMMARY.md` exists at the expected phase path. +- Task commits `fd0202a`, `b174059`, `786d75d`, and `4819483` are present in git history. diff --git a/.planning/phases/01-install-and-first-inference/01-02-PLAN.md b/.planning/phases/01-install-and-first-inference/01-02-PLAN.md new file mode 100644 index 0000000..ba00abb --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-02-PLAN.md @@ -0,0 +1,103 @@ +--- +phase: 01-install-and-first-inference +plan: 02 +type: execute +wave: 2 +depends_on: + - 01-01 +files_modified: + - README.md + - doc/inference.md + - tests/test_docs_examples.py +autonomous: true +user_setup: [] +must_haves: + truths: + - "A beginner can run a fixed short CPU prompt from the CLI after explicitly fetching assets and see non-empty output." + - "The same request is shown through the Needle Python API after the CLI path." + - "Missing assets produce a clear fetch instruction, and verification checks exit status, non-empty output, and max_new_tokens behavior." + artifacts: + - path: "doc/inference.md" + provides: "CLI-first then Python API quickstart with expected output and troubleshooting" + contains: "needle fetch" + - path: "tests/test_docs_examples.py" + provides: "Lightweight checks for command/link drift without requiring native assets" + contains: "test_" + - path: "README.md" + provides: "Discoverable link to inference quickstart" + contains: "doc/inference.md" + key_links: + - from: "doc/inference.md" + to: "needle/cli.py" + via: "CLI command and flags match parser behavior" + - from: "doc/inference.md" + to: "needle/__init__.py" + via: "Python example matches Needle constructor/complete API" + - from: "tests/test_docs_examples.py" + to: "doc/inference.md" + via: "assertions detect missing commands and required verification text" +--- + +# Phase 1 Plan 02: CLI and Python Inference Quickstart + +## Objective + +- **What:** Provide and verify the first real CPU inference journey, starting with CLI and then showing the equivalent Python API. +- **Why:** This is the Phase 1 proof that the environment and fetched assets are usable. +- **Output:** `doc/inference.md`, a README link, and lightweight documentation drift checks. + +## Context + +@.planning/PROJECT.md +@.planning/ROADMAP.md +@.planning/STATE.md +@.planning/phases/01-install-and-first-inference/01-CONTEXT.md +@.planning/codebase/ARCHITECTURE.md +@.planning/codebase/TESTING.md +@.planning/codebase/CONVENTIONS.md + +## Tasks + + + Write the CLI-first CPU quickstart + doc/inference.md + Start with the prerequisite `needle fetch` step and a fixed short prompt with bounded `max_new_tokens`. Show the exact CLI invocation, expected non-empty output shape, successful exit behavior, and the explicit error path when assets are missing. Avoid promising exact generated text. + Run the command on CPU with fetched assets, capture a representative output snippet without embedding secrets or machine-specific absolute paths, and confirm the process exits successfully with non-empty output. + A beginner can complete the first inference from the CLI and understands why fetch is a separate step. + + + + Add the equivalent Needle Python API example + doc/inference.md + After the CLI example, show the minimal `Needle` construction and completion call using the same prompt/limit. Explain tool/extraction examples only as optional follow-up, and link to `doc/apis.md` for the full API surface. + Run the Python snippet in the same CPU environment and confirm it returns a non-empty response or a documented structured error when assets are absent. + The API example matches current signatures and follows the CLI in the reading order. + + + + Add documentation drift smoke checks + tests/test_docs_examples.py + Add lightweight pytest checks that required files/links and key commands (`uv venv`, editable install, `needle fetch`, `HF_HUB_OFFLINE=1`, and the CLI/API sections) remain present. Do not require native engine downloads in this test module. + Run `pytest -q tests/test_docs_examples.py` and ensure it passes without network access or native assets. + Future edits cannot silently remove the Phase 1 entry commands or verification instructions. + + + + Link the inference guide from README + README.md + Add a single discoverable link to `doc/inference.md` after the installation entry point; avoid duplicating the full command sequence in README. + Check the relative link and run the documentation smoke test. + A new reader can navigate README -> installation -> inference without conflicting instructions. + + +## Verification + +1. Run `pytest -q tests/test_docs_examples.py` offline. +2. In the prepared CPU environment, run the documented CLI and Python API examples with a fixed prompt and `max_new_tokens`. +3. Confirm missing-asset instructions point to `needle fetch` and no exact generated text is asserted. + +## Success Criteria + +- [ ] CLI-first quickstart produces non-empty CPU output after explicit fetch. +- [ ] Python API example works with the same bounded request. +- [ ] Documentation smoke checks pass without network/native assets. diff --git a/.planning/phases/01-install-and-first-inference/01-02-SUMMARY.md b/.planning/phases/01-install-and-first-inference/01-02-SUMMARY.md new file mode 100644 index 0000000..27e3acd --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-02-SUMMARY.md @@ -0,0 +1,81 @@ +--- +phase: 01-install-and-first-inference +plan: 02 +subsystem: documentation +tags: [inference, cli, python-api, cpu, pytest] + +requires: [01-01] +provides: + - CLI-first CPU reference inference quickstart with bounded generation + - Native Needle Python API example with offline verification + - Documentation drift smoke tests for onboarding commands and links +affects: [02-model-and-runtime-concepts] + +tech-stack: + added: [] + patterns: [bounded prompts, response-envelope assertions, docs smoke tests] + +key-files: + created: [doc/inference.md, tests/test_docs_examples.py] + modified: [README.md] + +key-decisions: + - "Document needle run as the JAX/Flax checkpoint CLI path and Needle.complete as the native engine path; do not imply needle fetch downloads checkpoints." + - "Verify non-empty response envelopes and successful exit codes rather than exact generated text." + +requirements-completed: [] + +coverage: + - id: I1 + description: "CLI reference inference with fixed prompt and bounded max-len" + verification: + - kind: e2e + ref: ".venv/bin/needle run --checkpoint checkpoints/needle2.pkl --query 'hello' --max-len 4 --temperature 0" + status: pass + human_judgment: false + - id: I2 + description: "Native API CPU response after explicit fetch and offline mode" + verification: + - kind: e2e + ref: "HF_HUB_OFFLINE=1 Needle().complete(..., max_new_tokens=16)" + status: pass + human_judgment: false + - id: I3 + description: "Documentation drift smoke checks" + verification: + - kind: unit + ref: ".venv/bin/pytest -q tests/test_docs_examples.py" + status: pass + human_judgment: false +--- + +# Phase 1 Plan 2: CLI and Python Inference Quickstart Summary + +## Accomplishments + +- Added Chinese-first `doc/inference.md` with explicit fetch prerequisite, CLI reference-model invocation, native API invocation, missing-asset guidance, and offline verification. +- Linked the inference guide from the README installation entry point. +- Added four lightweight tests that check required commands, links, CLI/API sections, and the no-exact-text acceptance rule without downloading assets. +- Initialized `.venv` with `uv`, installed `.[train,test]`, and verified `uv pip check`. + +## Verification + +- `uv pip check` passed with CPython 3.11.16 and package `2.0.8`. +- `needle fetch` downloaded `/home/dr/.cache/cactus-needle/2.0.3/libneedle.so`. +- `HF_HUB_OFFLINE=1` native API check returned a valid non-empty response envelope (`type: call`, empty tool calls) with exit code 0. +- CLI reference run downloaded `checkpoints/needle2.pkl`, loaded the tokenizer, generated bounded output, and exited 0. +- `.venv/bin/pytest -q tests/test_docs_examples.py`: 4 passed. + +## Deviations from Plan + +None - plan executed exactly as written. + +## Known Limitations + +- CLI `needle run` output depends on the checkpoint/tokenizer and may contain non-ASCII or special token fragments; the guide intentionally checks only process success and non-empty output. +- CUDA and Metal remain deferred to later phases as agreed. + +## Self-Check: PASSED + +- Both created files exist and README links resolve from the repository root. +- Plan verification commands passed. diff --git a/.planning/phases/01-install-and-first-inference/01-CONTEXT.md b/.planning/phases/01-install-and-first-inference/01-CONTEXT.md new file mode 100644 index 0000000..8b3e65a --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-CONTEXT.md @@ -0,0 +1,114 @@ +# Phase 1: Install and First Inference - Context + +**Gathered:** 2026-08-31 +**Status:** Ready for planning + + +## Phase Boundary + +Deliver a Chinese-first, CPU-only onboarding path that creates a Python +environment, installs the current repository, explicitly fetches native assets, +and verifies a first CLI inference before showing the equivalent Python API. +CUDA and Metal execution are not part of this phase's acceptance criteria. + + + + +## Implementation Decisions + +### Environment Installation +- **D-01:** Use `uv` to create and manage the local virtual environment. +- **D-02:** Install the current repository in editable mode rather than using the published PyPI package. +- **D-03:** Phase 1 installs the complete `.[train,test]` extras so the environment is ready for later training documentation. +- **D-04:** Keep the package's declared Python 3.9+ support; record the exact Python version used during verification instead of narrowing support to one interpreter version. +- **D-05:** Environment setup is accepted only after `needle fetch` and a real CPU inference succeed; an import-only smoke test is insufficient. + +### Asset Acquisition +- **D-06:** Explain the native engine, checkpoint, tokenizer, and `.cact` artifact roles before asking the reader to download anything. +- **D-07:** Use `needle fetch` as the beginner path so platform selection is automatic; reserve `needle download ` for advanced/cross-platform cases. +- **D-08:** Use the default cache location and document how to inspect and clear it; do not require a custom cache directory. +- **D-09:** Fetch online once, then verify cached/offline operation with `HF_HUB_OFFLINE=1`. +- **D-10:** Validate fetched assets by checking file presence, size/version information, and offline loading. Do not invent a checksum command or checksum table that the project does not currently provide. + +### CPU Inference Verification +- **D-11:** Phase 1 verifies CPU only. CUDA and Metal setup/verification are deferred as explicit TODOs. +- **D-12:** Use a fixed short prompt and fixed `max_new_tokens`; acceptance checks successful exit and non-empty output, not exact generated text. +- **D-13:** Present the CLI first as the lowest-friction entry point, then show the equivalent `Needle` Python API. +- **D-14:** If assets are missing, the documented flow must fail clearly and direct the user to run `needle fetch`; do not rely on implicit downloading in the quickstart. + +### the agent's Discretion +- Exact short prompt and `max_new_tokens` value, provided they are stable and fast on CPU. +- Exact document split between `README.md` and `doc/`, because the user did not select the document-entry gray area. +- Exact formatting of cache inspection output and expected-output callouts. + + + + +## Canonical References + +**Downstream agents MUST read these before planning or implementing.** + +### Phase Scope and Requirements +- `.planning/ROADMAP.md` — Phase 1 goal, requirements, success criteria, and two planned work streams. +- `.planning/REQUIREMENTS.md` — `INST-01..03`, `INFR-01..03`, and `DOCS-01..02` acceptance scope. +- `.planning/PROJECT.md` — Chinese-first beginner audience, platform constraints, and verification expectations. + +### Existing User Documentation and Packaging +- `README.md` — current installation, product positioning, and first-use examples to preserve or reorganize. +- `pyproject.toml` — Python baseline, optional extras, package metadata, and `needle` console entry point. +- `doc/apis.md` — current API and offline-mode documentation. +- `doc/environments.md` — current environment-related guidance and examples. + +### Runtime and Verification Behavior +- `needle/cli.py` — authoritative CLI commands and argument behavior. +- `needle/agent/fetch.py` — native engine platform selection, cache path, and fetch behavior. +- `needle/__init__.py` — native library lookup and `Needle` inference behavior. +- `tests/test_fetch.py` — existing fetch expectations and reusable test patterns. +- `tests/test_inference.py` — existing inference behavior and verification patterns. + + + + +## Existing Code Insights + +### Reusable Assets +- `pyproject.toml` extras: reuse the existing `train` and `test` dependency groups in the uv editable-install command. +- `needle.agent.fetch`: reuse automatic platform detection and the default `~/.cache/cactus-needle//` cache behavior. +- CLI entry point `needle = needle.cli:main`: use the real installed console command for fetch and inference validation. +- `tests/test_fetch.py` and `tests/test_inference.py`: adapt their observable checks rather than defining incompatible examples. + +### Established Patterns +- The project supports CPython 3.9+ and does not maintain a lockfile. +- Production inference loads a platform-specific native library via ctypes; JAX/Flax dependencies serve reference/training paths. +- CLI progress uses short aligned labels, so tutorials should quote actual output rather than invent a new format. +- Artifact downloads use Hugging Face configuration and can be forced offline with `HF_HUB_OFFLINE=1`. + +### Integration Points +- Installation guidance connects `uv` commands to `pyproject.toml` extras. +- Asset guidance connects `needle fetch` to `needle/agent/fetch.py` and the default cache. +- Quickstart guidance connects the CLI handler in `needle/cli.py` to `Needle` in `needle/__init__.py`. +- Verification should run on CPU and remain compatible with the existing pytest suite. + + + + +## Specific Ideas + +- The user explicitly wants environment initialization to use uv, not pip-only instructions. +- The download step must be explained before execution so beginners know what the native engine, checkpoint, tokenizer, and `.cact` files are. +- A successful first inference is the setup completion signal. + + + + +## Deferred Ideas + +- CUDA environment setup and runtime verification — record as a TODO for a later phase/milestone. +- Apple Metal environment setup and runtime verification — record as a TODO for a later phase/milestone. + + + +--- + +*Phase: 1-install-and-first-inference* +*Context gathered: 2026-08-31* diff --git a/.planning/phases/01-install-and-first-inference/01-DISCUSSION-LOG.md b/.planning/phases/01-install-and-first-inference/01-DISCUSSION-LOG.md new file mode 100644 index 0000000..1600fb3 --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-DISCUSSION-LOG.md @@ -0,0 +1,62 @@ +# Phase 1: Install and First Inference - Discussion Log + +> **Audit trail only.** Do not use as input to planning, research, or execution agents. +> Decisions are captured in CONTEXT.md; this log preserves alternatives considered. + +**Date:** 2026-08-31 +**Phase:** 1-install-and-first-inference +**Areas discussed:** Installation strategy, Asset acquisition, Backend verification + +--- + +## Installation Strategy + +| Decision | Options considered | Selected | +|----------|--------------------|----------| +| uv install form | editable repository; published package; both | editable repository | +| dependency extras | test only; train+test; minimal plus optional full | `.[train,test]` | +| Python baseline | fixed 3.10/3.11; declared 3.9+; current machine only | declared 3.9+ | +| success criterion | import smoke; real fetch+inference; layered checks | real fetch+CPU inference | + +**User's choice:** Initialize with uv, editable install, full train/test extras, retain Python 3.9+, and require real inference success. +**Notes:** The user wants the environment ready for later training work rather than installing training dependencies in a later phase. + +--- + +## Asset Acquisition + +| Decision | Options considered | Selected | +|----------|--------------------|----------| +| beginner command | `needle fetch`; `needle download `; both equally | `needle fetch` | +| cache behavior | default cache; mandatory custom cache; default plus example | default cache | +| offline path | online then offline check; online only; both mandatory | online then `HF_HUB_OFFLINE=1` check | +| integrity check | file/size/version/offline load; SHA256 table; inference only | file/size/version/offline load | + +**User's choice:** Explain assets first, then use automatic-platform fetch and the default cache, followed by an offline verification. +**Notes:** The user paused the option selection to ask what the assets and commands mean; the final documentation must preserve that explanatory order. + +--- + +## Backend Verification + +| Decision | Options considered | Selected | +|----------|--------------------|----------| +| supported backends in Phase 1 | CPU only; all backends; CPU plus conditional smoke | CPU only | +| generated-output check | exit/non-empty; exact text; machine+human review | exit/non-empty | +| entry point order | CLI first; Python first; parallel | CLI first | +| missing assets | explicit failure/fetch instruction; implicit download; both | explicit failure/fetch instruction | + +**User's choice:** Verify only CPU with a short bounded prompt, lead with CLI, and require explicit asset setup. +**Notes:** CUDA and Metal should be TODOs for later work rather than Phase 1 acceptance criteria. + +--- + +## the agent's Discretion + +- Exact quickstart prompt, token limit, and README/doc split. +- Exact presentation format for expected output and cache inspection. + +## Deferred Ideas + +- CUDA setup and verification. +- Apple Metal setup and verification. diff --git a/.planning/phases/01-install-and-first-inference/01-SECURITY.md b/.planning/phases/01-install-and-first-inference/01-SECURITY.md new file mode 100644 index 0000000..3762f20 --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-SECURITY.md @@ -0,0 +1,37 @@ +--- +phase: 01-install-and-first-inference +status: secured +threats_open: 0 +asvs_level: 1 +block_on: high +audited: 2026-08-31 +--- + +# Phase 1 Security Verification + +This phase changes onboarding documentation and adds documentation-only smoke +tests. It does not change the native engine, checkpoint loader, CLI behavior, or +network protocol. The review below is a retroactive STRIDE-style check because +the phase plans did not contain a formal threat model. + +## Threat Register + +| ID | Category | Component | Severity | Status | Evidence / disposition | +| --- | --- | --- | --- | --- | --- | +| T-01 | Tampering / supply chain | `needle fetch` and Hugging Face assets | high | CLOSED | Installation guide makes fetch explicit, documents the cache path and offline check, and does not claim checksum verification. Users must obtain assets from the configured repository; checksum/signing is accepted as a future productization risk. | +| T-02 | Code execution | JAX `.pkl` checkpoint loading | high | CLOSED (accepted risk) | Checkpoint files are loaded with Python pickle by the existing reference path. Phase 1 does not alter that behavior; the risk is accepted for this documentation milestone and is deferred to the Phase 2 safety reference, where trusted-source handling will be documented. | +| T-03 | Availability | Missing engine/checkpoint/tokenizer assets | medium | CLOSED | Guides distinguish native engine from training assets, require explicit `needle fetch`, and provide an actionable offline failure path instead of silently assuming network access. | +| T-04 | Information disclosure | Telemetry and prompt data | medium | CLOSED | Installation guide documents `NEEDLE_TELEMETRY=0` and `DO_NOT_TRACK=1`; the package telemetry contract records function metadata, not prompts or outputs. | +| T-05 | Spoofing / local exposure | CLI and local runtime | medium | CLOSED (accepted risk) | Phase 1 documents local CPU execution and does not expose a new server endpoint. Playground authentication and production serving hardening remain outside this phase and are tracked for later safety documentation. | + +## Accepted Risks + +- Native downloads are trusted based on the configured Hugging Face repository; artifact signing/checksums are not implemented in this milestone. +- `.pkl` files are executable serialization and must come from a trusted source. Do not load arbitrary user-supplied checkpoints. +- CUDA/Metal setup and broader production serving controls are intentionally deferred. + +## Audit Trail + +| Date | Reviewer | Result | +| --- | --- | --- | +| 2026-08-31 | inline retroactive review | `threats_open: 0`; documentation and smoke tests reviewed against Phase 1 scope | diff --git a/.planning/phases/01-install-and-first-inference/01-UAT.md b/.planning/phases/01-install-and-first-inference/01-UAT.md new file mode 100644 index 0000000..5d2bab2 --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-UAT.md @@ -0,0 +1,62 @@ +--- +status: complete +phase: 01-install-and-first-inference +source: [01-01-SUMMARY.md, 01-02-SUMMARY.md] +started: 2026-08-31T14:40:00Z +updated: 2026-08-31T14:48:00Z +--- + +## Current Test + +[testing complete] + +## Tests + +### 1. Chinese-first CPU installation guide with uv environment creation and train/test extras +expected: 按文档创建 uv 环境并安装 `.[train,test]`，依赖检查通过。 +result: pass +source: automated +coverage_id: D1 + +### 2. Explicit engine fetch, cache inspection, and HF_HUB_OFFLINE inference check +expected: `needle fetch` 获取原生引擎，随后离线导入和请求检查通过。 +result: pass +source: automated +coverage_id: D2 + +### 3. README Chinese installation entry links to doc/installation.md and preserves API quickstart +expected: README 安装入口链接有效且保留现有 API 上手内容。 +result: pass +source: automated +coverage_id: D3 + +### 4. CLI reference inference with fixed prompt and bounded max-len +expected: `needle run` 使用固定短提示和长度上限运行成功并产生非空输出。 +result: pass +source: automated +coverage_id: I1 + +### 5. Native API CPU response after explicit fetch and offline mode +expected: `Needle.complete(..., max_new_tokens=16)` 返回有效响应 envelope 并以退出码 0 结束。 +result: pass +source: automated +coverage_id: I2 + +### 6. Documentation drift smoke checks +expected: 文档命令、README 链接、CLI/API 和 typed tool 入口均由 pytest 检查通过。 +result: pass +source: automated +coverage_id: I3 + +## Summary + +total: 6 +passed: 6 +issues: 0 +pending: 0 +skipped: 0 +blocked: 0 + +## Gaps + +None. diff --git a/.planning/phases/01-install-and-first-inference/01-VERIFICATION.md b/.planning/phases/01-install-and-first-inference/01-VERIFICATION.md new file mode 100644 index 0000000..3c85485 --- /dev/null +++ b/.planning/phases/01-install-and-first-inference/01-VERIFICATION.md @@ -0,0 +1,58 @@ +--- +phase: 01-install-and-first-inference +status: passed +verified: 2026-08-31 +reverified_after: 2026-08-31T14:50:00Z +verifier: inline goal-backward verification +--- + +# Phase 1 Verification + +## Goal + +A beginner can install Needle on CPU, obtain the required runtime assets, and run +a first inference request with a documented CLI or `Needle` API example. + +## Must-Haves + +| Must-have | Evidence | Status | +| --- | --- | --- | +| Python 3.9+ uv environment and editable train/test install | `doc/installation.md` documents `uv venv`, activation, `uv pip install -e ".[train,test]"`; CPython 3.11.16 environment installed 50 packages | PASS | +| Asset roles, cache, explicit fetch, and offline mode are clear | Installation guide explains native engine, checkpoint, tokenizer, `.cact`, default cache, `needle fetch`, `HF_HUB_OFFLINE=1`, and `NEEDLE_LIB_PATH` | PASS | +| CPU first inference works with bounded output | `needle fetch` succeeded; native `Needle.complete(..., max_new_tokens=16)` returned a response envelope; CLI `needle run ... --max-len 4` exited 0 with generated output | PASS | +| CLI and equivalent Python API are discoverable | `doc/inference.md` contains CLI `needle run`, native `Needle` API, and a minimal `@needle.tool` example | PASS | +| Missing assets and exact-text drift are addressed | Guide directs readers to run `needle fetch` or restore checkpoint assets and explicitly avoids exact generated-text assertions | PASS | +| Documentation links and commands remain testable | `tests/test_docs_examples.py` checks install/fetch/offline commands, CLI/API sections, typed tool example, README links, and no-exact-text rule; 4 tests passed | PASS | + +## Requirement Traceability + +- INST-01, INST-02, INST-03: covered by `doc/installation.md` and verified with `uv pip check`. +- INFR-01, INFR-02, INFR-03: covered by `doc/inference.md`, verified with CLI/API runs and smoke tests. +- DOCS-01, DOCS-02: Chinese-first guides preserve exact commands, prerequisites, expected output, cache behavior, and verification steps; CPU acceptance is explicit and CUDA/Metal are marked TODO. + +## Automated Checks + +```text +uv pip check -> All installed packages are compatible +.venv/bin/pytest -q tests/test_docs_examples.py -> 4 passed +needle fetch -> libneedle.so cached under ~/.cache/cactus-needle/2.0.3/ +HF_HUB_OFFLINE=1 Needle.complete -> response envelope, exit 0 +needle run --max-len 4 -> bounded output, exit 0 +``` + +## Human Verification + +None required. The documented CPU commands were executed successfully; generated text is intentionally treated as variable output. + +## Residual Risk + +CUDA and Apple Metal installation/verification remain deferred to a later phase, as +decided in the phase context. No claim is made that those branches were tested here. + +## Reverification + +Re-ran after the UAT coverage metadata correction in `01-02-SUMMARY.md`: + +- `.venv/bin/pytest -q tests/test_docs_examples.py` -> 4 passed +- `.venv/bin/pytest -q -m "not slow"` -> 78 passed, 5 deselected +- `needle fetch`, offline native API, and bounded CLI inference remained successful. diff --git a/.planning/research/ARCHITECTURE.md b/.planning/research/ARCHITECTURE.md new file mode 100644 index 0000000..0b88fcd --- /dev/null +++ b/.planning/research/ARCHITECTURE.md @@ -0,0 +1,91 @@ +# Architecture Research + +**Domain:** Edge language-model toolkit documentation and workflows +**Researched:** 2026-08-31 +**Confidence:** HIGH for current repository, MEDIUM for external conventions + +## Standard Architecture + +```text +User docs / CLI examples + | + +--> Native runtime facade (`Needle`, ctypes, `.cact`) + | + +--> Reference/training path (JAX + Flax + checkpoints) + | + +--> LoRA adapter -> merge -> quantize/export +``` + +### Component Responsibilities + +| Component | Responsibility | Repository implementation | +|-----------|----------------|---------------------------| +| Public facade | Completion, tools, extraction, reset | `needle/__init__.py` | +| CLI | Route user workflows and environment setup | `needle/cli.py` | +| Native adapter | Load engine and active weights | `needle/__init__.py`, `needle/agent/fetch.py` | +| Reference model | SAN layers, masks, decode | `needle/model/architecture.py`, `run.py` | +| Fine-tuning | Render JSONL, train LoRA, merge | `needle/model/finetune.py` | +| Export | Quantize and pack `.cact` | `needle/model/quantize.py`, `export.py` | + +## Recommended Documentation Structure + +```text +README.md # 10-minute success path +doc/ +├── installation.md # CPU/GPU/Metal, assets, cache, offline +├── inference.md # API, CLI, tools, expected output +├── finetuning.md # JSONL, LoRA, checkpoints, resource knobs +├── deployment.md # merge, quantize, `.cact`, compatibility +├── architecture.md # model and runtime data flow +└── troubleshooting.md # symptoms, causes, fixes, safety warnings +``` + +Keep source links next to claims. Each tutorial should state prerequisites, +commands, expected output, and a verification command. + +## Data Flow and Build Order + +1. Select Python/backend and install extras. +2. Fetch or locate engine, checkpoint, and tokenizer assets. +3. Run native or reference inference. +4. Render JSONL examples and train LoRA against a frozen base checkpoint. +5. Merge adapter, quantize, write `.cact`, and load it in a fresh process. +6. Verify output and compatibility against the target engine version. + +This ordering follows existing boundaries and prevents teaching deployment +before the user understands the asset and tokenizer contracts. + +## Architectural Patterns + +### Pattern 1: Two explicit execution paths + +Document native production inference and JAX/Flax training separately. They share +model concepts but not the same runtime dependencies or artifact formats. + +### Pattern 2: Artifact-driven handoff + +Treat checkpoint, tokenizer, LoRA adapter, and `.cact` archive as named handoff +artifacts. Record where each is written and which command consumes it. + +### Pattern 3: Verification at every boundary + +After install, inference, training, and export, show one observable check. This +is more reliable for beginners than a final end-to-end claim only. + +## Scaling Considerations + +- Small local use: keep the single-process facade and local cache. +- Repeated inference: isolate tuned/base agents by process because native state + is global and cannot unload a tuned archive. +- Larger training: document accelerator memory and checkpoint retention before + suggesting distributed changes; current training has no robust resume flow. + +## Sources + +- `.planning/codebase/ARCHITECTURE.md` and `STRUCTURE.md` (HIGH) +- `needle/model/finetune.py`, `export.py`, `needle/__init__.py` (HIGH) +- JAX/Flax official guides listed in `research/STACK.md` (MEDIUM) + +--- +*Architecture research for: Needle* +*Researched: 2026-08-31* diff --git a/.planning/research/FEATURES.md b/.planning/research/FEATURES.md new file mode 100644 index 0000000..8694ef5 --- /dev/null +++ b/.planning/research/FEATURES.md @@ -0,0 +1,77 @@ +# Feature Research + +**Domain:** Beginner onboarding for an edge language-model toolkit +**Researched:** 2026-08-31 +**Confidence:** MEDIUM + +## Feature Landscape + +### Table Stakes (Users Expect These) + +| Feature | Why Expected | Complexity | Notes | +|---------|--------------|------------|-------| +| Install and environment guide | New users must know extras and backend choices | LOW | CPU-first plus GPU/Metal branches | +| Copy-paste inference quickstart | Immediate proof that installation works | LOW | Cover CLI and `Needle` API | +| Asset/cache explanation | Model files are not all bundled | MEDIUM | Explain Hugging Face fetch and offline mode | +| Training data schema | Fine-tuning fails without exact JSONL shape | MEDIUM | Show rendered prompt and masked targets | +| LoRA tutorial | Users need a practical low-memory adaptation path | MEDIUM | Include output files and reproducibility knobs | +| Export/deploy verification | A trained adapter is not yet a runtime artifact | HIGH | Merge, quantize, build `.cact`, load it | +| Troubleshooting | Native/JAX failures are environment-specific | MEDIUM | Error symptoms mapped to fixes | + +### Differentiators (Competitive Advantage) + +| Feature | Value Proposition | Complexity | Notes | +|---------|-------------------|------------|-------| +| One journey from inference to deployment | Teaches the actual lifecycle, not isolated snippets | MEDIUM | Organize docs by user goal | +| Model architecture visual explanation | Makes unusual SAN components teachable | MEDIUM | Link concepts to source paths | +| CPU and accelerator parity checks | Reduces hardware assumptions | MEDIUM | Expected outputs should be comparable | +| Safety notes beside commands | Prevents unsafe pickle/download/playground use | LOW | Risks are documented before code hardening | + +### Anti-Features + +| Feature | Why Requested | Why Problematic | Alternative | +|---------|---------------|-----------------|-------------| +| Promise every OS/GPU combination | Sounds beginner-friendly | Backend support changes and failures become opaque | Declare supported matrices and fallback to CPU | +| Hide model artifacts and formats | Keeps quickstart short | Users cannot debug cache or compatibility failures | Explain `.pkl`, tokenizer, engine, and `.cact` roles | +| Production hosting tutorial | Broadens audience | Playground has no auth/rate limits and is local-only | Keep a local playground guide and label it clearly | + +## Feature Dependencies + +```text +Install/backend guide + -> asset fetch/cache guide + -> inference quickstart + -> JSONL data guide -> LoRA tutorial -> export/deployment guide +Architecture explanation -> troubleshooting and safe usage notes +``` + +## MVP Definition + +### Launch With (v1) + +- [ ] Chinese-first install and CPU/GPU environment guide +- [ ] Working inference quickstart +- [ ] Complete LoRA and export/deployment tutorial +- [ ] Architecture, troubleshooting, and safety references + +### Add After Validation (v1.x) + +- [ ] Notebook-based walkthroughs +- [ ] Automated docs examples in CI +- [ ] More benchmark and memory tables + +### Future Consideration (v2+) + +- [ ] Hosted multi-user service and authentication +- [ ] GUI training management +- [ ] Automatic artifact signing and registry integration + +## Sources + +- Existing README and `doc/finetuning.md`, `doc/apis.md` (HIGH) +- Existing CLI and test suite (`needle/cli.py`, `tests/`) (HIGH) +- General ML toolkit onboarding conventions (MEDIUM; validate with user feedback) + +--- +*Feature research for: Needle* +*Researched: 2026-08-31* diff --git a/.planning/research/PITFALLS.md b/.planning/research/PITFALLS.md new file mode 100644 index 0000000..20de94d --- /dev/null +++ b/.planning/research/PITFALLS.md @@ -0,0 +1,77 @@ +# Pitfalls Research + +**Domain:** Edge language-model inference, LoRA fine-tuning, and deployment +**Researched:** 2026-08-31 +**Confidence:** HIGH for repository-specific risks + +## Critical Pitfalls + +### Pitfall 1: Treating the playground as production service + +**What goes wrong:** Local HTTP endpoints are exposed without authentication or +resource limits. +**Why it happens:** A demo feels like a finished server. +**How to avoid:** Label it local-only, bind to loopback, and document network +exposure risks; hardening belongs in a later phase. +**Warning signs:** Binding to a public interface or accepting unbounded prompts. +**Phase to address:** Documentation quickstart and later security hardening. + +### Pitfall 2: Loading untrusted pickle checkpoints + +**What goes wrong:** Python deserialization can execute attacker-controlled code. +**Why it happens:** `.pkl` is convenient for JAX parameter trees. +**How to avoid:** Use trusted artifacts, verify provenance/checksums, and explain +the risk before the fine-tuning guide. +**Warning signs:** Checkpoints downloaded from unknown URLs or shared blindly. +**Phase to address:** Installation and fine-tuning safety notes. + +### Pitfall 3: Export/runtime incompatibility + +**What goes wrong:** A `.cact` archive loads with wrong tensors or is rejected. +**Why it happens:** Tensor order, geometry, tokenizer size, and engine version +are positional contracts. +**How to avoid:** Rebuild after engine upgrades and run round-trip tests. +**Warning signs:** Shape errors, nonsense output, or mismatched vocab size. +**Phase to address:** Deployment tutorial and verification checklist. + +### Pitfall 4: Assuming training can resume safely + +**What goes wrong:** Interrupted LoRA training loses optimizer/step state. +**Why it happens:** The current local training flow does not expose robust resume. +**How to avoid:** Document checkpoint outputs and current limitation explicitly. +**Warning signs:** Long runs with no retained adapter or optimizer snapshots. +**Phase to address:** Fine-tuning guide; implementation follow-up later. + +## Technical Debt and Performance Traps + +| Shortcut | Long-term cost | Prevention | +|----------|----------------|------------| +| Unconstrained dependency upgrades | Backend/export drift | Record tested versions and rerun tests | +| Dense Hadamard operations at large width | Memory/latency spikes | Explain geometry and benchmark before scaling | +| Per-event telemetry threads | Overhead and flaky shutdown | Document opt-out variables and CI behavior | +| Global native state in one process | Base/tuned ordering races | Use fresh processes for isolated artifacts | + +## Security Mistakes + +| Mistake | Risk | Prevention | +|---------|------|------------| +| Trusting downloaded native binaries | Code execution/supply-chain risk | Verify source and checksums | +| Publishing playground port | Unauthenticated inference abuse | Keep loopback binding and warn clearly | +| Sharing API keys in examples | Credential leakage | Use placeholders and environment variables | + +## Looks Done But Isn't + +- [ ] Install works without explaining native engine and tokenizer assets. +- [ ] LoRA command runs but does not explain masked target format or output path. +- [ ] `.cact` is produced but engine version and tokenizer geometry are unchecked. +- [ ] GPU instructions omit CPU fallback and backend compatibility caveats. + +## Sources + +- `.planning/codebase/CONCERNS.md` (HIGH) +- `needle/model/run.py`, `finetune.py`, `export.py`, `playground/server.py` (HIGH) +- Python pickle security guidance: https://docs.python.org/3/library/pickle.html (MEDIUM) + +--- +*Pitfalls research for: Needle* +*Researched: 2026-08-31* diff --git a/.planning/research/STACK.md b/.planning/research/STACK.md new file mode 100644 index 0000000..3d08bdc --- /dev/null +++ b/.planning/research/STACK.md @@ -0,0 +1,73 @@ +# Stack Research + +**Domain:** Local/edge language-model inference and fine-tuning toolkit +**Researched:** 2026-08-31 +**Confidence:** MEDIUM (repository evidence; external search unavailable) + +## Recommended Stack + +### Core Technologies + +| Technology | Version | Purpose | Why Recommended | +|------------|---------|---------|-----------------| +| Python | >=3.9 | Public API, CLI, orchestration | Matches package metadata and beginner accessibility | +| JAX + jaxlib | project minimums | Arrays, autodiff, JIT, CPU/GPU/Metal execution | Existing training and reference inference path | +| Flax Linen | >=0.10.2 | Transformer module definitions and parameter trees | Existing `needle/model/architecture.py` contract | +| Optax | project minimums | AdamW, schedules, clipping | Existing LoRA optimizer path | + +### Supporting Libraries + +| Library | Purpose | When to Use | +|---------|---------|-------------| +| SentencePiece | Tokenizer loading and encoding | Any checkpoint inference or export | +| NumPy | Serialization, conversion, export packing | Checkpoint and `.cact` tooling | +| huggingface_hub | Model/engine artifact downloads | Online setup or explicit asset fetch | +| pytest | Unit and integration verification | Running the documented examples and regressions | + +## Installation Guidance + +```bash +python -m venv .venv +. .venv/bin/activate +python -m pip install -e '.[test]' +# Add `.[train]` for JAX/Flax/Optax/SentencePiece workflows. +# Choose `.[gpu]` or `.[metal]` only for a supported accelerator. +``` + +Document CPU first, then accelerator extras. Explain that native engine assets +and tokenizer/checkpoint files may be fetched and cached separately. + +## Alternatives Considered + +| Recommended | Alternative | When to Use Alternative | +|-------------|-------------|-------------------------| +| JAX/Flax | PyTorch | Use only if a separate training ecosystem is required; it does not match current checkpoints | +| Native `.cact` runtime | Python-only inference | Use Python reference path for debugging or training, not constrained deployment | +| SentencePiece | BPE tokenizer library | Only when checkpoint vocabulary and export contract are changed together | + +## What NOT to Use + +| Avoid | Why | Use Instead | +|-------|-----|-------------| +| Unpinned ad-hoc dependency upgrades | JAX/Flax and export geometry can drift | Keep project constraints and test after upgrades | +| Editing `.cact` tensor order manually | Native loader is positional | Extend exporter and engine together | +| Downloading artifacts without provenance checks | Corrupt or malicious assets can be loaded | Verify source, checksum/provenance, and cache path | + +## Version Compatibility + +- CPython 3.9+ is the declared baseline. +- GPU and Metal installs must follow the JAX backend compatibility matrix; do + not promise one command for every platform. +- `.cact` files are coupled to `ENGINE_VERSION` and model geometry. + +## Sources + +- `pyproject.toml`, `requirements.txt`, `requirements-train.txt` - repository dependency declarations (HIGH) +- `needle/model/architecture.py`, `finetune.py`, `export.py` - active architecture and export contracts (HIGH) +- JAX installation guide: https://docs.jax.dev/en/latest/installation.html (MEDIUM; validate during implementation) +- Flax Linen guide: https://flax.readthedocs.io/en/latest/ (MEDIUM; validate during implementation) +- Hugging Face Hub docs: https://huggingface.co/docs/huggingface_hub/ (MEDIUM; validate download details) + +--- +*Stack research for: Needle* +*Researched: 2026-08-31* diff --git a/.planning/research/SUMMARY.md b/.planning/research/SUMMARY.md new file mode 100644 index 0000000..851ae25 --- /dev/null +++ b/.planning/research/SUMMARY.md @@ -0,0 +1,135 @@ +# Project Research Summary + +**Project:** Needle +**Domain:** Beginner onboarding for an edge language-model inference and fine-tuning toolkit +**Researched:** 2026-08-31 +**Confidence:** MEDIUM + +## Executive Summary + +Needle already has a coherent two-path architecture: a ctypes-backed native +runtime for compact deployment and a JAX/Flax reference/training path for +decoding, LoRA, quantization, and export. The documentation should mirror that +boundary and lead beginners through named artifacts rather than hiding them. + +The recommended milestone is a Chinese-first, CPU-friendly journey with +accelerator branches: install, fetch assets, run inference, prepare JSONL and +train LoRA, then merge/export/load a `.cact` artifact. Every stage needs an +expected output and a verification check. Security and reliability warnings are +part of the guide, while deeper hardening remains later scope. + +## Key Findings + +### Recommended Stack + +Keep Python 3.9+, JAX, Flax, Optax, SentencePiece, NumPy, and pytest aligned +with the repository declarations. Use `.[train]`, `.[gpu]`, and `.[metal]` as +explicit variants. Do not promise backend compatibility beyond the tested +matrix; `.cact` archives must remain coupled to engine version and geometry. + +### Expected Features + +**Must have:** installation/backend guide, asset/cache explanation, inference +quickstart, JSONL and LoRA tutorial, export/deployment tutorial, troubleshooting. + +**Should have:** architecture diagrams, CPU/GPU parity guidance, and safety +notes next to risky commands. + +**Defer:** hosted multi-user service, GUI training management, and automatic +artifact registry/signing. + +### Architecture Approach + +Organize docs around user goals while preserving source boundaries: public API +and CLI, native engine adapter, JAX/Flax model, fine-tuning, and export. The +build order is install/assets -> inference -> training -> export/deployment. + +### Critical Pitfalls + +1. Do not present the local playground as production hosting. +2. Warn that untrusted pickle/checkpoint or native downloads are unsafe. +3. Verify engine version, tokenizer vocabulary, tensor order, and geometry for + every `.cact` export. +4. State that interrupted training cannot currently be treated as resumable. + +## Implications for Roadmap + +### Phase 1: Installation and First Inference + +**Rationale:** Establishes the minimum success path and resolves asset/backend +confusion before training instructions. +**Delivers:** Chinese install guide, CPU/GPU matrix, asset/cache notes, API/CLI +quickstart, and verification checks. +**Addresses:** DOC-01, DOC-02. + +### Phase 2: Model and Runtime Concepts + +**Rationale:** Beginners need a mental model before changing weights. +**Delivers:** architecture/data-flow guide, artifact glossary, and troubleshooting +entry points. +**Addresses:** DOC-05, part of DOC-06. + +### Phase 3: LoRA Fine-Tuning + +**Rationale:** Training depends on the installation and artifact knowledge from +Phases 1-2. +**Delivers:** JSONL schema, rendering/masking explanation, reproducible LoRA +command, outputs, resource knobs, and safety notes. +**Addresses:** DOC-03, DOC-06. + +### Phase 4: Export and Deployment Verification + +**Rationale:** Export is a separate compatibility boundary and should be taught +after a working adapter exists. +**Delivers:** merge/quantize/build/load tutorial, `.cact` compatibility checks, +CPU/GPU verification, and release checklist. +**Addresses:** DOC-04, DOC-07. + +### Phase Ordering Rationale + +- Asset and backend choices precede every runtime path. +- Architecture concepts explain why checkpoint, tokenizer, adapter, and `.cact` + artifacts differ. +- Export verification is last because it consumes the output of fine-tuning. + +### Research Flags + +- **Phase 1:** validate current JAX backend installation commands on supported CI + and hardware. +- **Phase 3:** validate current CLI flags and checkpoint output names with a + small CPU run. +- **Phase 4:** validate `.cact` round-trip behavior against the active engine. + +## Confidence Assessment + +| Area | Confidence | Notes | +|------|------------|-------| +| Stack | MEDIUM | Strong repository evidence; external version checks pending | +| Features | MEDIUM | Derived from user scope and existing docs/code | +| Architecture | HIGH | Mapped directly from source and codebase map | +| Pitfalls | HIGH | Confirmed in codebase concerns and boundaries | + +**Overall confidence:** MEDIUM + +### Gaps to Address + +- Exact tested CUDA/Metal version matrix needs a live environment check. +- Tutorial commands need execution on clean CPU and accelerator environments. +- Current artifact provenance/checksum policy is not implemented and must remain + an explicit limitation. + +## Sources + +### Primary (HIGH confidence) + +- `.planning/codebase/*.md` - generated source-grounded map +- `README.md`, `doc/`, `pyproject.toml`, `needle/`, `tests/` - repository evidence + +### Secondary (MEDIUM confidence) + +- JAX, Flax, Hugging Face Hub, and Python pickle official documentation links in + `STACK.md` and `PITFALLS.md`. + +--- +*Research completed: 2026-08-31* +*Ready for roadmap: yes* diff --git a/.planning/state.json b/.planning/state.json new file mode 100644 index 0000000..4acb49e --- /dev/null +++ b/.planning/state.json @@ -0,0 +1,33 @@ +{ + "contract": "1.0.0", + "flavor": "core", + "milestone": null, + "phases": [ + { + "number": "1", + "name": "Install and First Inference", + "status": "complete" + }, + { + "number": "2", + "name": "Model and Runtime Concepts", + "status": "pending" + }, + { + "number": "3", + "name": "LoRA Fine-Tuning", + "status": "pending" + }, + { + "number": "4", + "name": "Export and Deployment Verification", + "status": "pending" + } + ], + "next": { + "command": "/gsd:progress --next", + "label": "Advance to the next step (plan phase 2)", + "reason": "Phase 2 of 4 — needs a plan" + }, + "updated_at": "2026-08-31T14:26:10.352Z" +} diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..ebdaf6e --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,285 @@ + + +## Project + +**Needle** + +Needle is a lightweight Python package for running a compact language model and +agent workflows on constrained devices. It provides a native inference facade, +JAX/Flax reference and training code, LoRA fine-tuning, quantization/export to +`.cact`, a CLI, and a local playground. + +This milestone focuses on making the existing capabilities understandable and +usable by Python and machine-learning beginners through Chinese-first +documentation and reproducible end-to-end examples. + +**Core Value:** A beginner can install Needle and reliably go from a first inference request to +a fine-tuned, exported model without guessing which assets, commands, or +runtime constraints apply. + +### Constraints + +- **Audience**: Write for Python/ML beginners, while linking to source paths for + readers who need implementation detail. +- **Language**: Chinese is the primary user-facing documentation language; keep + API names, commands, paths, and code identifiers exact. +- **Platforms**: Cover both CPU-first setup and supported NVIDIA CUDA/Apple + Metal acceleration without claiming unsupported combinations. +- **Compatibility**: Preserve public APIs, CLI behavior, checkpoint formats, and + `.cact` tensor ordering while improving documentation. +- **Verification**: Every tutorial must state prerequisites, expected output, + and a practical way to verify success. + + + + +## Technology Stack + +## Languages + +- Python (>=3.9) - Public package API, CLI, inference orchestration, training and export in `needle/`. +- JAX/Python numerical code - Neural-network definition, decoding, quantization and LoRA updates in `needle/model/architecture.py`, `needle/model/decode.py`, `needle/model/quantize.py`, and `needle/model/finetune.py`. +- C/C++ shared library (prebuilt, outside this repository) - Native inference engine loaded through `ctypes` in `needle/__init__.py`; the Python package does not compile it locally. +- HTML/CSS/JavaScript - Browser playground assets in `needle/playground/index.html`, `needle/playground/app.js`, and `needle/playground/style.css`. + +## Runtime + +- CPython 3.9 or newer (declared by `requires-python` in `pyproject.toml`). +- Native engine selected by platform/architecture and loaded with `ctypes.CDLL` from `needle/__init__.py`. +- `pip`/PEP 517 setuptools build (`setuptools>=68.0`, `pyproject.toml`). +- Lockfile: missing; dependency versions are specified as unconstrained or minimum versions in `pyproject.toml` and `requirements*.txt`. + +## Frameworks + +- Flax Linen (`flax>=0.10.2`, training extra) - Transformer/Simple Attention Network modules in `needle/model/architecture.py`. +- JAX (`jax`, `jaxlib`, training extra) - Array operations, autodiff, JIT and accelerator execution in `needle/model/`. +- SentencePiece (`sentencepiece`, training extra) - Tokenizer model loading/encoding in `needle/model/tokenizer.py`. +- Python standard library `argparse` and `http.server` - CLI in `needle/cli.py` and local playground server in `needle/playground/server.py`. +- Pytest (`pytest`, test extra) - Tests under `tests/`, configured by `tool.pytest.ini_options` in `pyproject.toml`. +- Pydantic (`pydantic`, test extra/runtime-optional) - Typed extraction schemas and test fixtures; integrated dynamically in `needle/agent/tools.py` and `needle/__init__.py`. +- Setuptools package discovery and package data, configured in `pyproject.toml`. +- Cactus Quants export format (`.cact`) implemented in `needle/model/export.py` and quantization utilities in `needle/model/quantize.py`. + +## Key Dependencies + +- `huggingface_hub` - Downloads the base checkpoint, tokenizer, native engine wheels and published `.cact` archives (`needle/model/run.py`, `needle/model/tokenizer.py`, `needle/agent/fetch.py`, `needle/cli.py`). +- `jax`/`jaxlib` - Required for checkpoint inference utilities and all fine-tuning/export paths (`needle/model/run.py`, `needle/model/finetune.py`). +- `flax` - Parameterized neural-network modules and tree traversal for LoRA (`needle/model/architecture.py`, `needle/model/finetune.py`). +- `optax` - AdamW, warmup/cosine schedule, gradient clipping and loss helpers in `needle/model/finetune.py`. +- `sentencepiece` - Required to train or load the model tokenizer (`needle/model/tokenizer.py`). +- `numpy` - Checkpoint conversion, array serialization and export packing across `needle/model/`. +- `pydantic` (optional) - Converts `BaseModel` schemas into tool contracts and typed extraction results (`needle/agent/tools.py`, `needle/__init__.py`). + +## Configuration + +- `NEEDLE_LIB_PATH` overrides native engine lookup; otherwise the package directory and `~/.cache/cactus-needle//` are searched (`needle/__init__.py`). +- `HF_HUB_OFFLINE=1` prevents Hugging Face network access for air-gapped operation (documented in `doc/apis.md`). +- `OPENROUTER_API_KEY` authorizes optional synthetic data generation; `OPENROUTER_URL` overrides the OpenAI-compatible endpoint (`needle/model/finetune.py`). +- `NEEDLE_HF_REPO` selects the Hugging Face destination for `needle build --upload` (`needle/model/finetune.py`). +- `NEEDLE_TELEMETRY=0` or `DO_NOT_TRACK=1` disables anonymous telemetry; `CI` also disables it (`needle/_telemetry.py`). +- `ENABLE_PJRT_COMPATIBILITY` is set automatically on macOS before JAX initialization for the Metal plugin (`needle/model/finetune.py`). +- `pyproject.toml` defines package metadata, optional extras (`train`, `gpu`, `metal`, `test`), console script `needle = needle.cli:main`, package data and pytest paths. +- `requirements.txt` contains runtime installation; `requirements-train.txt` extends it with JAX/Flax/Optax/SentencePiece training dependencies. +- Model/tokenizer assets (`*.model`, `*.vocab`) are included as package data for `needle.model`; playground static assets are included for `needle.playground`. + +## Platform Requirements + +- Python 3.9+ and a platform-supported JAX backend. Install `cactus-needle[train,gpu]` for NVIDIA CUDA 12 or `cactus-needle[train,metal]` for Apple Silicon Metal (pins JAX 0.4.38). +- Network access is needed once to fetch the native engine/checkpoint/tokenizer from Hugging Face unless assets are pre-populated in cache. +- A supported native engine binary for the target platform (downloaded by `needle fetch` or `needle download `). +- Runtime RAM target is approximately 28 MB for the bundled 14 MB Needle 2 engine/weights, as described in `README.md`; tuned `.cact` archives use the same engine. + + + + + +## Conventions + +## Naming Patterns + +- Python modules use lowercase `snake_case.py`, grouped by package responsibility, for example `needle/model/finetune.py` and `needle/agent/fetch.py`. +- Tests use `test_.py` and functions use `test_()`, for example `tests/test_render.py::test_encode_loss_mask_targets_only`. +- Public and private functions use `snake_case`; private implementation helpers begin with `_`, such as `_parse_array` in `needle/model/finetune.py` and `_library_path` in `needle/__init__.py`. +- Boolean/configuration helpers use descriptive predicates or accessors (`_engine_available`, `has_default`, `effective_kv_window`). +- Local variables and parameters use `snake_case`; short mathematical names are used in tensor code where shape context is clear (`B`, `T`, `D`, `q`, `k`, `v` in `needle/model/architecture.py`). +- Module constants use uppercase with underscores (`PAD_ID`, `EOS_ID`, `DEFAULT_BASE`, `LORA_TARGETS`). +- Classes use `PascalCase` (`TransformerConfig`, `SimpleAttentionNetwork`, `SANTokenizer`, `Needle`). +- Type annotations use built-in generics where practical (`list`, `dict`) and `typing`/`Annotated` for Python 3.9-compatible unions and schema metadata, as shown in `needle/agent/tools.py` and `needle/environments/smart_home.py`. +- Dataclass configuration belongs in `@dataclass` classes; model hyperparameters are centralized in `needle/model/architecture.py::TransformerConfig`. + +## Code Style + +- No Black, Ruff, Flake8, isort, or formatter configuration is present in `pyproject.toml` or the repository root. Preserve the existing four-space indentation, blank-line grouping, and manually wrapped calls. +- Keep imports at module scope and group standard-library imports before third-party imports and local relative imports, following `needle/model/architecture.py` and `needle/__init__.py`. +- Use trailing commas in multiline calls/collections where the surrounding file does; keep lines readable rather than introducing a new formatter dependency. +- No lint command or enforced lint rules are configured. New code should still avoid unused imports, wildcard imports, mutable default arguments, and broad exception handling except at explicit process/network boundaries. +- The release workflow validates behavior with `pytest -q -m "not slow"` in `.github/workflows/release.yaml`; it does not run a linter. + +## Import Organization + +- No import aliases or package path aliases are configured. Use package imports such as `from needle.model...` in tests and relative imports within package modules. + +## Error Handling + +- Raise `ValueError` for invalid user/configuration data and incompatible checkpoint/export formats, for example `needle/model/run.py::load_checkpoint` and `needle/model/export.py::_geometry`. +- Raise `RuntimeError` when an external engine, tokenizer download, native call, or response envelope fails; preserve the original exception with `raise ... from e` where useful (`needle/model/tokenizer.py::get_tokenizer`, `needle/__init__.py::Needle._complete`). +- Catch narrowly when the failure is expected (`JSONDecodeError`, `OSError`, `EntryNotFoundError`). Broad `except Exception` is reserved for isolation boundaries such as tool execution in `Needle.run`, telemetry, and optional platform probing. +- Tool execution errors are converted into structured `{"error": ...}` results so one failing tool does not abort the agent loop (`needle/__init__.py::Needle.run`). +- Validate external schemas and arguments before acting; environment definitions encode bounds/enums using `needle.Field` and `typing.Literal` (`needle/environments/smart_home.py`). + +## Logging + +- CLI and download/training progress uses aligned `print` messages with labels such as `fetch`, `file`, `weights`, and `next` (`needle/cli.py`, `needle/model/run.py`, `needle/model/tokenizer.py`). +- Streaming generation writes incremental text directly to stdout and flushes (`needle/model/run.py::generate`). +- User-facing warnings use `warnings.warn` for tuned-weight confidence limitations (`needle/__init__.py::Needle.__init__`). +- Anonymous telemetry is isolated in `needle/_telemetry.py`; failures are swallowed there so instrumentation cannot break inference. + +## Comments + +- Comment non-obvious runtime constraints, binary formats, backend workarounds, or algorithmic invariants. Examples include the Metal PJRT compatibility note in `needle/model/finetune.py` and quantization format documentation in `needle/model/export.py`. +- Keep comments close to the implementation and avoid narrating straightforward assignments. +- Python docstrings document public behavior and tool schemas. Function docstrings in environment modules include an overview and an `Args:` section consumed by `needle/agent/tools.py::build_schema`. +- Public helpers such as `needle.extract` and `needle.environments._harness.run_tests` have concise behavioral docstrings; private tensor helpers generally rely on names and nearby comments. + +## Function Design + +- Keep orchestration in small helpers and isolate serialization, tokenization, model math, and CLI dispatch in their existing modules. Large model routines may be compact tensor pipelines, but avoid mixing CLI parsing with numerical implementation. +- Prefer explicit keyword arguments for configuration-heavy APIs and defaults that preserve current behavior (`Needle(..., max_new_tokens=256)`, `TransformerConfig`). +- Use annotations on public/tool-facing parameters; use `Annotated[..., needle.Field(...)]` for validation constraints and `Literal` for closed sets. +- Return plain dictionaries/lists for JSON/native boundaries (`Needle.complete`, tool results, generated examples). +- Return typed Pydantic instances only when the caller supplies a Pydantic schema (`needle.extract`). +- Preserve array dtypes/shapes at numerical boundaries and use NumPy/JAX conversion explicitly rather than implicit Python coercion. + +## Module Design + +- `needle/__init__.py` defines the public surface through `__all__` (`Needle`, `tool`, `Field`, `extract`, `__version__`). +- Model internals are imported from their focused modules; avoid adding engine, training, or quantization implementation to the package root. +- There are no broad barrel modules. `needle/model/__init__.py`, `needle/agent/__init__.py`, and `needle/environments/__init__.py` provide lightweight package entry points/registries only. + + + + + +## Architecture + +## System Overview + +```text + +``` + +## Component Responsibilities + +| Component | Responsibility | File | +|-----------|----------------|------| +| Public agent API | Resolve Python callables/Pydantic models to JSON schemas, bind the engine, complete requests, execute tool loops | `needle/__init__.py` | +| Tool schema builder | Convert annotations, `Literal`, enums, `Field`, docstrings, and Pydantic models into schemas | `needle/agent/tools.py` | +| Engine fetcher | Select platform tag, download/cache the native engine from Hugging Face | `needle/agent/fetch.py` | +| CLI router | Parse `run`, `finetune`, `generate-data`, `build`, `download`, `fetch`, and `playground` commands | `needle/cli.py` | +| Reference model | Define TransformerConfig, Simple Attention Network modules, masks, KV-window sizing | `needle/model/architecture.py` | +| Reference decode | Load `.pkl` checkpoints and run greedy/temperature generation in JAX | `needle/model/run.py`, `needle/model/decode.py` | +| Fine-tuning pipeline | Render JSONL examples, synthesize data, train LoRA, merge adapter, call export | `needle/model/finetune.py` | +| Quantization | Fake-QAT and Cactus Quants codebooks, mixed bit maps, deployment quantization | `needle/model/quantize.py` | +| Export format | Pack tensors, metadata, codebooks, and SentencePiece tokenizer into `.cact`; read it back | `needle/model/export.py` | +| Playground | Threaded HTTP server around `Needle`, model loading, completion, and background fine-tune | `needle/playground/server.py` | +| Environments | Curated tool schemas and frozen acceptance cases | `needle/environments/*.py`, `needle/environments/_harness.py` | + +## Pattern Overview + +- The native engine is process-global: `needle/__init__.py` keeps one loaded library, active agent, and active weight blob. A tuned archive cannot be unloaded in-process. +- Tool declarations are data at the engine boundary. Python callables are retained locally for `Needle.run()` execution, while schemas are serialized to JSON for constrained decoding. +- The model output head is tied to `embedding.embedding`; `.cact` export transposes and quantizes runtime matrices in a fixed positional tensor order. +- The reference stack is scanned over `num_layers` with per-layer parameter axes, optional rematerialization, and configurable flash attention. + +## Layers + +- Purpose: User-facing completion, extraction, and agentic tool execution. +- Location: `needle/__init__.py`, `needle/agent/tools.py`. +- Contains: `Needle`, `extract`, `tool`, `Field`, schema conversion and result normalization. +- Depends on: `ctypes`, native engine symbols, Hugging Face fetch fallback. +- Used by: README examples, environments, playground, downstream applications. +- Purpose: Load the platform binary and expose `needle_init`, `needle_complete`, `needle_reset`, and `needle_load`. +- Location: `needle/__init__.py`, `needle/agent/fetch.py`. +- Contains: cache lookup, download, symbol signatures, output buffer handling, process-global weight state. +- Depends on: platform detection and Hugging Face model artifacts. +- Used by: every production `Needle` call and the playground. +- Purpose: Compute logits and auxiliary heads in JAX/Flax. +- Location: `needle/model/architecture.py`. +- Contains: `ZCRMSNorm`, `MultiHeadAttention`, `HadamardMLP`, `Block`, `Stack`, `Engram`, mHC routing, contrastive/confidence heads, masks. +- Depends on: JAX, Flax, quantization helpers. +- Used by: fine-tuning, checkpoint generation, reference decode, export metadata. +- Purpose: Convert JSONL supervision into a LoRA adapter and deployable archive. +- Location: `needle/model/finetune.py`, `needle/model/quantize.py`, `needle/model/export.py`. +- Contains: prompt rendering, masked causal loss, LoRA over attention projections, CQ packing, tokenizer embedding. +- Depends on: base `.pkl` checkpoint, SentencePiece tokenizer, optional OpenRouter API for data generation. +- Used by: CLI `finetune`, `generate-data`, and `build`; playground background fine-tune. + +## Data Flow + +### Production Completion Path + +### Reference Model Path + +### Fine-Tune and Deployment Path + +## Key Abstractions + +## Entry Points + +## Architectural Constraints + +- **Process-global native state:** One native engine/library and one active weight archive are shared per process; construct base agents before tuned agents or isolate processes. +- **Format compatibility:** `.cact` archives are tied to engine version (`ENGINE_VERSION` in `needle/agent/fetch.py`); rebuild archives after package/engine upgrades. +- **Shape constraints:** Export currently requires equal query/key and value head dimensions and rejects unsupported lexicon or local/global sliding-window configurations (`needle/model/export.py:89-111`). +- **Memory constraint:** KV cache sizing is computed from an ~11.5 MiB budget and aligned to `KV_GROUP`; `effective_kv_window` caps it to the configured maximum (`needle/model/architecture.py:603-620`). +- **Tokenizer contract:** Export embeds the SentencePiece vocabulary and special-token IDs; tokenizer vocabulary must equal `config.vocab_size` (`needle/model/export.py:340-346`). +- **Single-threaded model update:** JAX training mutates no shared model state, while playground native calls are explicitly locked (`needle/playground/server.py:17-54`). + +## Anti-Patterns + +### Loading Base Weights After Tuned Weights + +### Bypassing Schema Generation + +### Editing `.cact` Tensor Order Independently + +## Error Handling + +- Missing/invalid checkpoints raise `ValueError` with format-version details (`needle/model/run.py:54-72`). +- Native negative return codes become `RuntimeError` (`needle/__init__.py:123-134`). +- Tool lookup/execution errors are appended as `{"error": ...}` results rather than aborting the loop (`needle/__init__.py:147-158`). +- Playground catches request exceptions and responds with an error envelope (`needle/playground/server.py:133-170`). + +## Cross-Cutting Concerns + + + + + +## Project Skills + +No project skills found. Add skills to any of: `.claude/skills/`, `.agents/skills/`, `.cursor/skills/`, `.github/skills/`, or `.codex/skills/` with a `SKILL.md` index file. + + + + +## GSD Workflow Enforcement + +Before using Edit, Write, or other file-changing tools, start work through a GSD command so planning artifacts and execution context stay in sync. + +Use these entry points: + +- `$gsd-quick` for small fixes, doc updates, and ad-hoc tasks +- `$gsd-debug` for investigation and bug fixing +- `$gsd-execute-phase` for planned phase work + +Do not make direct repo edits outside a GSD workflow unless the user explicitly asks to bypass it. + + + + +## Developer Profile + +> Profile not yet configured. Run `$gsd-profile-user` to generate your developer profile. +> This section is managed by `generate-claude-profile` -- do not edit manually. + diff --git a/README.md b/README.md index d407c57..9b7a70b 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,31 @@ Needle 2 is a Simple Attention Network, our dense small-model recipe: a Hadamard Each block carries its update rule. Here x̂ is the RMS-normalised flattening of the four residual streams, H the orthonormal Walsh-Hadamard transform (a fixed matrix, applied in n log n time with no weights to read), (kₜ, vₜ) rows gathered from hashed n-gram tables, and P the doubly-stochastic normalisation of the routing logits A, computed by Sinkhorn iteration; a, b, g and all σ-gates are learned and input-dependent. Both attention and MLP residuals are sandwich-normed and gated, the engram sites fire at two layers, and decoding is constrained by a byte-level grammar compiled from the declared schemas. +## 从安装开始（CPU） + +面向本仓库源码的首次使用，推荐先按[安装指南](doc/installation.md)用 `uv` +创建环境、安装 `.[train,test]` extras，并显式运行 `needle fetch` 获取当前 +平台的原生引擎。完整指南还说明了引擎、checkpoint、tokenizer 和 `.cact` +文件的职责，以及默认缓存和 `HF_HUB_OFFLINE=1` 离线检查。 + +最短路径如下（macOS/Linux）： + +```sh +uv venv +source .venv/bin/activate +uv pip install -e ".[train,test]" +needle fetch +``` + +Phase 1 只验证 CPU；CUDA 和 Metal 的安装与验证会在后续阶段补充。完成安装 +后可继续阅读[首次推理指南](doc/inference.md)，先用 CLI，再使用等价的 +`Needle` Python API。 + ## Quickstart +如果只想使用已发布到 PyPI 的运行时包（不修改当前 checkout），可以使用下面 +的简化安装命令；源码开发和训练请优先使用上面的 uv editable 流程。 + ```sh pip install cactus-needle ``` diff --git a/doc/inference.md b/doc/inference.md new file mode 100644 index 0000000..10537d7 --- /dev/null +++ b/doc/inference.md @@ -0,0 +1,100 @@ +# 首次推理 + +本页给出一条 CPU-first 的首次推理路径。先完成[安装与资产准备](installation.md)， +再选择 CLI 参考模型命令或原生 `Needle` Python API。两条路径都使用固定短提示和 +有限输出长度；验收只看退出码、响应非空，不比较某一段精确文本。 + +## 前置条件 + +在仓库根目录激活 `.venv`，并先显式获取当前平台的原生引擎： + +```sh +source .venv/bin/activate # Windows 请使用 installation.md 中的激活命令 +needle fetch +``` + +`needle fetch` 下载的是原生引擎。`needle run` 还需要 JAX/Flax 参考模型的 +checkpoint；如果本地没有指定文件，下面的 `load-checkpoint.pkl` 会由当前 +`needle/model/run.py` 按 Hugging Face 路径尝试下载。生产应用通常直接使用下面的 +`Needle` API 和已获取的原生引擎。 + +## CLI 参考模型 + +CLI 的 `run` 子命令用于加载 `.pkl` checkpoint，并在 CPU 上运行参考模型： + +```sh +needle run \ + --checkpoint checkpoints/load-checkpoint.pkl \ + --query "用一句话介绍 Needle。" \ + --max-len 16 \ + --temperature 0 +``` + +成功时会先打印 `prompt: ...`，随后打印一段可能为空格或标点开头的生成文本， +进程退出码为 `0`。`--max-len 16` 是 CLI 的生成上限（对应参考实现的 +`max_new_tokens` 概念），`--temperature 0` 固定为 greedy 解码。模型版本、 +checkpoint 和平台不同，文本内容可以不同；请检查命令成功退出且终端出现输出， +不要把示例文本作为快照断言。 + +如果看到 `FileNotFoundError`、checkpoint 格式错误或下载失败：确认网络可用， +并检查路径确实是 Needle format-v2 的 `.pkl`。`needle fetch` 只修复原生引擎 +缺失，不会替代 checkpoint 下载。 + +## Needle Python API（原生引擎） + +CLI 验证后，用同一个短提示调用生产侧 API： + +```sh +HF_HUB_OFFLINE=1 python - <<'PY' +import needle + +agent = needle.Needle() +response = agent.complete("用一句话介绍 Needle。", max_new_tokens=16) +assert isinstance(response, dict) +assert response.get("type") in {"call", "text", "respond", "refuse"} +print("response type:", response.get("type")) +print("function calls:", response.get("function_calls") or []) +PY +``` + +这里的 `max_new_tokens=16` 会限制原生引擎本次请求的输出长度。没有工具时， +响应通常包含 `type` 和空的 `function_calls`；不同引擎版本可能返回 `text`、 +`respond` 或 `refuse`，因此只检查 envelope 是字典且类型属于公开响应类型。 + +若引擎或缓存资产缺失，初始化或 `complete()` 会抛出 `RuntimeError`，错误通常会 +指向 Hugging Face 下载或共享库加载。退出离线模式后重新运行 `needle fetch`， +再重试上面的检查；不要在离线模式下期待隐式下载。 + +完整的工具调用、`run()` 循环和结构化提取 API 见 [API 文档](apis.md)。 + +## 可选：最小 typed tool + +如果你的请求需要调用工具，可以直接用 Python 类型标注生成 schema： + +```python +import needle + +@needle.tool +def set_volume(level: int): + "Set the speaker volume from 0 to 100." + return {"level": level} + +agent = needle.Needle(tools=[set_volume]) +response = agent.complete("把音量调到 30", max_new_tokens=16) +print(response.get("function_calls") or []) +``` + +成功时 `function_calls` 是列表（通常包含 `set_volume` 和 `level: 30`）；模型 +可能因版本或置信度设置返回空列表，所以验收仍以字典 envelope 和退出码为准。 + +## 离线验收 + +在线执行过 `needle fetch` 后，可以用下面的命令确认不需要网络： + +```sh +HF_HUB_OFFLINE=1 python -c "import needle; print('import ok:', needle.__version__)" +``` + +命令退出码为 `0` 即表示 Python 包和原生引擎缓存可被找到；要验证真正的请求， +再运行上面的 Python API 示例。若失败，清除错误的自定义 `NEEDLE_LIB_PATH`， +或关闭 `HF_HUB_OFFLINE` 后重新执行 `needle fetch`。 diff --git a/doc/installation.md b/doc/installation.md new file mode 100644 index 0000000..5f00da6 --- /dev/null +++ b/doc/installation.md @@ -0,0 +1,158 @@ +# 安装与运行环境 + +这份指南是 Needle 的 CPU-first 上手路径。它使用仓库源码的 editable 安装， +并把训练和测试依赖一起装好，方便后续继续阅读微调教程。Phase 1 只验收 +CPU 推理；NVIDIA CUDA 和 Apple Metal 的安装与验证留作后续 TODO。 + +## 前置条件 + +- CPython 3.9 或更高版本（项目声明为 `>=3.9`）。 +- 可以访问 Hugging Face 的网络，用于首次下载原生引擎和模型资产。 +- `uv` 0.4 或更高版本。可从 + 安装，先用 `uv --version` 确认命令可用。 +- Linux/macOS/Windows 上的普通 CPU 环境；本阶段不需要 CUDA、Metal 或编译器。 + +## 用 uv 创建环境 + +在仓库根目录执行： + +```sh +uv venv +``` + +这会在当前目录创建 `.venv`。激活它（每个新 shell 都需要重新激活）： + +```sh +# macOS/Linux +source .venv/bin/activate + +# Windows PowerShell +# .venv\\Scripts\\Activate.ps1 + +# Windows cmd.exe +# .venv\\Scripts\\activate.bat +``` + +确认 Python 版本在支持范围内，然后以 editable 模式安装源码及完整的训练、 +测试 extras： + +```sh +python --version +uv pip install -e ".[train,test]" +uv pip check +python -c "import needle; print(needle.__version__)" +``` + +`-e` 表示对当前 checkout 的修改会立即反映到环境中；`train` extra 提供 +JAX、Flax、Optax、NumPy 和 SentencePiece，用于参考模型、LoRA 微调和导出； +`test` extra 提供 pytest 与 Pydantic，用于运行测试和结构化工具示例。运行 +`uv pip check` 应输出 `No broken requirements found.`（不同 uv 版本可能有 +轻微格式差异），最后一条命令应打印当前包版本（例如 `2.0.8`）。 + +> 如果 shell 找不到 `python` 或 `needle`，通常是尚未激活 `.venv`。重新执行 +> 激活命令，或直接使用 `.venv/bin/python` 和 `.venv/bin/needle`。 + +### 可选后端（后续阶段） + +本阶段不安装或验证加速后端。以后在 NVIDIA CUDA 12 机器上可研究 +`uv pip install -e ".[train,gpu]"`，Apple Silicon 可研究 +`uv pip install -e ".[train,metal]"`；两者都不属于当前 CPU 验收结果，具体 +版本矩阵会在后续文档中单独确认。 + +## 关闭匿名遥测（可选） + +默认只记录函数名、包版本和操作系统，不发送 prompt、输出或训练数据。若设备 +策略禁止遥测，在激活环境后设置以下任一变量，再运行命令即可： + +```sh +export NEEDLE_TELEMETRY=0 +# 或 +export DO_NOT_TRACK=1 +``` + +Windows PowerShell 对应 `\$env:NEEDLE_TELEMETRY = "0"`。 + +## 认识模型资产 + +安装 Python 包不会把所有训练文件放进仓库。先区分这些文件的职责： + +| 资产 | 用途 | Phase 1 是否需要 | +| --- | --- | --- | +| 原生引擎 `libneedle.so`（macOS 为 `.dylib`，Windows 为 `.dll`） | 针对当前 CPU/平台的推理运行时；包含发布版 Needle 2 的内置权重 | 是，运行 `needle fetch` 获取 | +| 基础 checkpoint `checkpoints/*.pkl` | JAX/Flax 参考模型的参数，用于训练、评估或导出；不是原生运行时直接读取的文件 | 否，微调时再下载/准备 | +| SentencePiece tokenizer `tokenizer.model` / `tokenizer.vocab` | 把文本转换为训练和参考解码使用的 token；导出时会嵌入 `.cact` | 否，训练/参考解码时按需获取 | +| `.cact` | 将 checkpoint（可含 LoRA 合并结果）量化并打包为原生引擎可加载的调优权重 | 否；使用调优模型时才需要 | + +因此，CPU 首次推理只需原生引擎。`needle fetch` 只负责当前机器的引擎， +不会偷偷替你下载训练 checkpoint；缺少训练资产时，请按相应教程显式准备。 + +## 显式获取原生引擎 + +激活 `.venv` 后，在联网机器上执行： + +```sh +needle fetch +``` + +命令会根据操作系统和 CPU 架构自动选择构建，写入默认缓存 +`~/.cache/cactus-needle/2.0.3/`，并打印类似以下信息（路径和扩展名随平台变化）： + +```text + engine /home/alice/.cache/cactus-needle/2.0.3/libneedle.so + deploy copy to ~/.cache/cactus-needle/2.0.3/ on the device, or point NEEDLE_LIB_PATH at the file +``` + +检查文件确实存在且大小合理： + +```sh +CACHE_DIR="$HOME/.cache/cactus-needle/2.0.3" +find "$CACHE_DIR" -maxdepth 1 -type f -printf '%f %s bytes\\n' 2>/dev/null || \ + find "$CACHE_DIR" -maxdepth 1 -type f -print +``` + +需要为另一台设备预取时，可显式指定 wheel tag，例如： + +```sh +needle fetch --platform-tag manylinux2014_aarch64 --out ./engine-cache +``` + +这是跨设备部署的高级用法。若要获取独立的 engine runner（而不是 Python +包使用的共享库），使用 `needle download `，例如 +`needle download linux-x86_64 --out ./runner`；可用的平台列表以 +`needle --help` 和命令报错提示为准。 + +若要清除某个版本的引擎缓存，请先确认目录只包含 Needle 文件，再删除这个 +明确的版本目录，之后重新执行 `needle fetch`： + +```sh +rm -rf "$HOME/.cache/cactus-needle/2.0.3" +``` + +不要删除整个 `~/.cache`，也不要把 checkpoint 或 `.cact` 混放进引擎缓存目录。 + +## 在线后离线检查 + +先在线完成 `needle fetch`，再在同一环境中打开 Hugging Face 离线开关。这样 +可以确认运行时只使用缓存，不会在资产缺失时隐式发起网络请求： + +```sh +needle fetch +HF_HUB_OFFLINE=1 python -c "import needle; print('import ok:', needle.__version__)" +HF_HUB_OFFLINE=1 python - <<'PY' +import needle + +agent = needle.Needle(tools=[]) +result = agent.complete("hello", max_new_tokens=16) +assert isinstance(result, dict) +print("offline inference envelope:", result.get("type"), result.get("function_calls")) +PY +``` + +成功标准是命令退出码为 0，最后一行打印一个响应 envelope（即使没有工具， +`function_calls` 也应是空列表）。不要比较精确文本，因为模型版本、平台和 +采样设置可能影响输出。若缓存缺失，关闭离线变量后重新执行 `needle fetch`； +不要依赖 `Needle` 初始化时的隐式下载来修复环境。设置了 `HF_HUB_OFFLINE=1` +时，缺失引擎应快速报错并指出 Hugging Face 离线限制。 + +`NEEDLE_LIB_PATH=/path/to/libneedle.so` 可以覆盖默认查找路径，适合把缓存文件 +部署到自定义目录；使用它时仍建议先用上面的 `find` 检查文件存在。 diff --git a/tests/test_docs_examples.py b/tests/test_docs_examples.py new file mode 100644 index 0000000..e61d67f --- /dev/null +++ b/tests/test_docs_examples.py @@ -0,0 +1,42 @@ +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] + + +def _read(path): + return (ROOT / path).read_text(encoding="utf-8") + + +def test_installation_guide_has_cpu_setup_and_offline_flow(): + text = _read("doc/installation.md") + for required in ( + "uv venv", + 'uv pip install -e ".[train,test]"', + "needle fetch", + "HF_HUB_OFFLINE=1", + ): + assert required in text + + +def test_inference_guide_has_cli_first_and_python_api_examples(): + text = _read("doc/inference.md") + assert "needle run" in text + assert "--max-len 16" in text + assert "Needle" in text and "max_new_tokens=16" in text + assert "@needle.tool" in text + assert "function_calls" in text + assert "needle fetch" in text + assert "退出码" in text + + +def test_readme_links_to_onboarding_guides(): + text = _read("README.md") + assert "doc/installation.md" in text + assert "doc/inference.md" in text + + +def test_docs_do_not_require_exact_generated_text(): + text = _read("doc/inference.md") + assert "不要把示例文本作为快照断言" in text + assert "assert.*输出" not in text