Skip to content

About

A local microGPT implementation that splits tasks into subtasks

Resources

Stars

6 stars

Watchers

0 watching

Forks

Latest commit

 

History

17 Commits

Folders and files

Repository files navigation

solverforge-gpt

A minimal, zero-dependency, CPU-trainable GPT-2-style transformer written entirely in Rust. It does one thing: takes a task description and returns a list of subtasks.

Three domain-specific models are trained independently -- work, software engineering, and creative -- with a TF-IDF router that automatically selects the right one at inference time.

Inspired by Karpathy's microgpt, reimplemented from scratch in Rust with a custom autograd engine.

Why zero dependencies?

The entire stack -- autograd, transformer, BPE tokenizer, Adam optimizer, RNG, binary serialization -- is built on top of Rust std and nothing else. The [dependencies] section in Cargo.toml is empty. This keeps the binary small, the build fast, and the code auditable.

Architecture

Each domain model is a decoder-only transformer (~1.5M parameters, ~4.3 MB on disk):

  • 4 layers, 4 attention heads, 128-dimensional embeddings
  • RMSNorm (pre-norm), ReLU MLP, residual connections
  • Block size of 512 tokens
  • KV cache for efficient autoregressive inference
  • Gradient clipping (max norm 1.0) and warmup + cosine LR decay
  • Custom BPE tokenizer per domain (~1000 merges)

The domain router uses TF-IDF cosine similarity to classify incoming tasks.

Quick start

This repository does not currently include pre-trained weights/; you need to train or provide them first. To run inference after training:

# build
cargo build --release

# single task
cargo run --bin infer --release -- --weights weights/ --task "Build a REST API for user authentication"

# interactive REPL
cargo run --bin infer --release -- --weights weights/ --interactive

# pipe from stdin
echo "Plan a team offsite" | cargo run --bin infer --release -- --weights weights/

Or use Make:

make all      # build release binaries
make infer    # launch interactive REPL

Training

Train all three domains in parallel, then build the combined router:

make train

Train a single domain:

make train-swe
make train-work
make train-creative

Rebuild the router after training:

make router

Training parameters

All configurable via environment variables or Make arguments:

Variable Default Description
STEPS 30000 Training steps per domain
VOCAB_SIZE 1000 BPE vocabulary size
LR 5e-4 Learning rate
EVAL_EVERY 500 Evaluate every N steps
WEIGHTS weights Output directory
DATA data Input directory
make train STEPS=50000 LR=3e-4

For more stable inference, you can lower sampling randomness or force greedy decoding:

cargo run --bin infer --release -- --weights weights/ --task "Build a REST API for user authentication" --temperature 0.2
cargo run --bin infer --release -- --task "Build a REST API for user authentication" --temperature 0

Reality check: the corpora in data/ combine 535 hand-written examples with synthetic examples accepted by a Jev relevance gate, taking each domain to 1,537–2,693 examples (12× overall, see Dataset generation). That is enough for the ~1.5M-parameter decoder to learn coherent task→subtask structure, but the synthetic examples are template-derived, so the models stay domain-shaped rather than general.

Training data format

Training data lives in data/ as plain text files (swe.txt, work.txt, creative.txt). Each example follows this format:

DOMAIN: swe
TASK: Build a REST API for user authentication
SUB: Define API schema and authentication endpoints
SUB: Implement JWT token generation and validation
SUB: Write integration tests for auth flow

DOMAIN: swe
TASK: Migrate database from PostgreSQL to MySQL
SUB: Audit existing schema and queries for compatibility
SUB: Rewrite PostgreSQL-specific SQL to MySQL dialect

Examples are separated by blank lines.

Dataset generation

data/{swe,work,creative}.txt are the canonical training corpora and the only data path make train reads. They combine 535 hand-written examples with synthetic examples accepted by a TypeSafe Jev relevance gate (Noul ≥ 0.84 against the paired task, task realism ≥ 0.6). The pipeline lives in scripts/:

  • gen_candidates.py — mines task/subtask fragments from data/ and synthesizes candidate variants
  • jev_gate.py — judges candidates in batched requests; controls mode samples the corpus and other domains to check judge discrimination
  • expand_dataset.py — runs the gated expansion; every judgment is journaled to data/expand_work/
  • assemble_expanded.py — rebuilds data/{swe,work,creative}.txt from the journals at any --threshold, with no API calls

data/expand_work/ and data/gen/ are the provenance record: each judged candidate is stored with its score and the source task it was paired with, so the corpus can be re-assembled or audited after the fact without re-spending API tokens.

Library usage

The crate exposes Decomposer (single domain) and MultiDecomposer (auto-routing) for use as a library:

use microgpt::{MultiDecomposer, Domain};

let md = MultiDecomposer::load("weights/").unwrap();

// auto-detect domain
let subtasks = md.split("Build a login page with OAuth").unwrap();
for s in &subtasks {
    println!("- {}", s);
}

// explicit domain
let subtasks = md.split_in_domain("Write a sonnet about rain", Domain::Creative).unwrap();

Project structure

src/
  lib.rs          Public API (Decomposer, MultiDecomposer)
  model.rs        Transformer model (training + inference with KV cache)
  tensor.rs       Autograd engine (reverse-mode autodiff, Adam optimizer, gradient clipping)
  tokenizer.rs    BPE tokenizer (training + encode/decode)
  router.rs       TF-IDF domain classifier
  rng.rs          xorshift64 PRNG with Box-Muller sampling
  bin/
    train.rs      Training CLI
    infer.rs      Inference CLI (single / interactive / pipe)
    build_router.rs  Router construction CLI
data/             Canonical training corpora (swe, work, creative) and generation artifacts
weights/          Pre-trained model weights and tokenizers
scripts/          Training shell scripts and the Jev-gated dataset pipeline

Make targets

make all              Build release binaries
make train            Train all domains + build router
make train-swe        Train software engineering domain
make train-work       Train work/business domain
make train-creative   Train creative domain
make router           Build domain router from training data
make infer            Interactive inference REPL
make clean            Delete all weights and tokenizer files
make help             Show all targets and variables

About

A local microGPT implementation that splits tasks into subtasks

Resources

Stars

6 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages