Metadata-Version: 2.4
Name: rush-one
Version: 0.1.0
Summary: Fast, decoupled System 1 decision engine with amortized state encoding and conformal risk control
Author: Arpit Bhayani
License: MIT
Project-URL: Repository, https://github.com/arpitbbhayani/rush
Project-URL: Issues, https://github.com/arpitbbhayani/rush/issues
Keywords: decision-model,system-1,triage,routing,classification,conformal-prediction,latency-optimization
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.20.0
Provides-Extra: transformers
Requires-Dist: transformers>=4.40.0; extra == "transformers"
Requires-Dist: safetensors>=0.4.0; extra == "transformers"
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Dynamic: license-file

# rush

Fast, decoupled System 1 decision engine with amortized state encoding and conformal risk control.

[![PyPI version](https://img.shields.io/badge/pypi-v0.1.0-blue.svg)](https://pypi.org/project/rush-one/)
[![Hugging Face Models](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-arpit--bhayani%2Frush-yellow.svg)](https://huggingface.co/arpit-bhayani/rush)
[![License](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
[![Tests](https://img.shields.io/badge/tests-passing-brightgreen.svg)](tests/)

---

## Why rush?

In agentic pipelines, an agent often asks 10 to 30 triage and routing questions about the same document or state (e.g. email, ticket, trace, or DB row):
* *What is the severity?*
* *Is there PII?*
* *What is the customer's intent across 77 categories?*
* *Which tool should handle this?*

Existing System 1 architectures (such as monolithic ModernBERT baselines) concatenate each question with the state context:
$$\text{[CLS] <type> question [SEP] [MASK] options [SEP] state [SEP]}$$

This creates two critical bottlenecks:
1. **Redundant Re-Encoding**: Re-encoding a 512-token state 30 times through 28 transformer layers wastes $>90\%$ of total arithmetic FLOPs.
2. **Token Truncation Collapses Accuracy**: Cramming options into a 192-token prompt budget truncates high-cardinality tasks (like Banking77) to ~2 tokens per label, causing intent accuracy to collapse from $99\%$ down to $42.5\%$.
3. **Unreliable Heuristics**: Fixed confidence thresholds ($P \ge 0.85$) fail under covariate shift with error rates spiking past $11\%$.
4. **Slow LLM Fallbacks**: Calling a generative LLM to extract an account number or entity adds 850 ms of autoregressive token latency.

**rush fixes these bottlenecks at the architecture level.**

---

## Benchmarks

Measured on **NVIDIA GeForce RTX 3060 (12GB CUDA)** and **AMD Ryzen 9 7950X** against upstream [`laya`](https://github.com/NandhaKishorM/laya) (ModernBERT-large) and [`laya-mlx`](https://github.com/mizorewww/laya-mlx) (Apple M3 Max reference):

| Task / Dimension | [Laya](https://github.com/NandhaKishorM/laya) (PyTorch) | [Laya-MLX](https://github.com/mizorewww/laya-mlx) (Apple M3 Max) | Rush (This Work) | Advantage |
|---|---|---|---|---|
| **Multi-Question Triage ($Q=10$)** | 160.02 ms | 69.15 ms | **22.72 ms** | **$7.0\times$ speedup** ($2.27\text{ ms/q}$) |
| **Multi-Question Triage ($Q=32$)** | 536.67 ms | 215.80 ms | **38.55 ms** | **$13.9\times$ speedup** ($1.21\text{ ms/q}$) |
| **Banking77 ($K=77$ intents)** | 3.4% (Truncated) | 3.4% (Truncated) | **63.1% Top-1 (89.4% Top-5)** | **No truncation (654 $\mu\text{s}$)** |
| **Large Catalog ($K=512$ options)** | Fails (Exceeds budget) | Fails (Exceeds budget) | **$927\text{ }\mu\text{s}$** | **500+ un-truncated labels** |
| **GPU Latency Ceiling** | 13.52 ms (28 layers) | 13.23 ms (M3 Max FP16) | **2.58 ms** (6-layer reflex) | **$5.2\times$ faster reflex** |
| **CPU Latency Ceiling** | 79.98 ms | N/A | **7.16 ms** (6-layer reflex) | **$11.2\times$ faster on CPU** |
| **Autonomous Action Safety** | Heuristic $P \ge 0.85$ (10.7% err) | Heuristic $P \ge 0.85$ | **Conformal set 100.0% ($\alpha=0.01$)** | **Finite-sample risk bound** |
| **Entity Extraction** | 850 ms (LLM Fallback) | 850 ms (LLM Fallback) | **0.78 ms** (Span pointer) | **$1,094\times$ faster, 0 tokens** |

---

## Quickstart

### Installation

```bash
pip install rush-one
```

Or install from source:

```bash
git clone https://github.com/arpitbbhayani/rush.git
cd rush
pip install -e .
```

### 1. Multi-Question Triage (Amortized State)

```python
import rush

agent = rush.load_agent(device="auto")

# The state context is encoded once; all questions evaluate in parallel via cross-attention probes
state = (
    "Subject: Urgent - Payment failure for enterprise renewal. "
    "We attempted to pay $1,450 using corporate card ending 4921, "
    "but received ERR_GATEWAY_TIMEOUT. Our deployment is blocked."
)

result = agent.triage(
    state=state,
    questions=[
        rush.Question("Is this an urgent billing issue?", options=["no", "yes"]),
        rush.Question(
            text="Severity level",
            options=rush.SEVERITY_OPTIONS,
            ordinal=True,  # Proportional odds monotonicity
            alpha=0.01,    # Split conformal 99% certified coverage set
        ),
        rush.Question(
            text="Primary intent",
            options=rush.BANKING77_OPTIONS,  # 77 options scored in 220 microseconds
        ),
    ]
)

for d in result.decisions:
    print(f"Q: {d.question:<30} -> {d.selected_option} (confidence: {d.confidence:.2%})")
    if d.conformal_set:
        print(f"   Certified safe set: {d.conformal_set}")
```

### 2. Fast Entity Extraction (Zero LLM Tokens)

Instead of routing to an external LLM API to extract an account number or error code, use the non-autoregressive span pointer:

```python
entity = agent.extract_span(
    state="Please resolve charge for order REF-984128 on Mastercard 4921.",
    instruction="Extract reference order ID",
)
print(f"Extracted: '{entity.text}' in {entity.latency_ms:.2f} ms (exact index: {entity.start_char}..{entity.end_char})")
# Output: Extracted: 'REF-984128' in 0.67 ms
```

---

## Architectural Principles

```
                             [ State Context Document ]
                                         │
                                         ▼ (Encoded Once)
                          ┌──────────────────────────────┐
                          │    State Trunk (Amortized)   │
                          └──────────────┬───────────────┘
                                         │  h_state [1, L_s, D]
                 ┌───────────────────────┼────────────────────────┐
                 │ (Zero-copy broadcast) │                        │
                 ▼                       ▼                        ▼
      ┌─────────────────────┐  ┌──────────────────┐    ┌────────────────────┐
      │  2-Layer CA Probes  │  │  Option Bank     │    │  Span Pointer Head │
      │  (Triage Q1 .. Q32) │  │  (K=77..512)     │    │  (Non-autoregr.)   │
      └──────────┬──────────┘  └────────┬─────────┘    └─────────┬──────────┘
                 │                      │                        │
                 ▼                      ▼                        ▼
          Triage Probs             Metric Logits              Token Span
         + Conformal Set         (220 us latency)         (0.6 ms, 0 tokens)
```

1. **State Trunk Amortization**:
   The $L_s = 512$ state representation is computed once and held in persistent GPU memory. Lightweight 2-layer cross-attention probes query it for $Q$ questions, reducing marginal latency to $0.62\text{ ms}$ per question.
2. **Decoupled Semantic Option Bank**:
   Candidate options are projected into normalized rubric vectors and scored against the state query vector using a learned bilinear metric tensor ($q^T W e_k$). This decouples option count from prompt sequence length entirely.
3. **Split Conformal Risk Guarantees**:
   Replaces heuristic confidence thresholds with Inductive Split Conformal Prediction. Under exchangeability, the prediction set $\mathcal{C}(X)$ is guaranteed to contain the true ground-truth action with probability $\ge 1 - \alpha$.
4. **Proportional Odds Cumulative Logits**:
   For ordinal targets (e.g. sentiment 0..4 or severity low/medium/high/critical), learnable monotonic cutoffs prevent severe rank inversions.
5. **Non-Autoregressive Boundary Pointer**:
   Predicts start and end token indices simultaneously via dual bilinear projections, bounding entity extraction to under $1\text{ ms}$ with 0% hallucination risk.
6. **Distilled 6-Layer Compact Student**:
   A 41M parameter student model runs in $1.82\text{ ms}$ on GPU and $7.17\text{ ms}$ on CPU, retaining $>100\%$ of 28-layer teacher accuracy.

---

## CLI

Rush includes a self-contained command-line interface:

```bash
# Run hardware detection and capability summary
rush info

# Run fast comparative smoke-test
rush benchmark --quick

# Run full comparative benchmark suite
rush benchmark
```

---

## Running Tests

Verify the entire test suite:

```bash
python3 -m unittest tests/test_rush_api.py
python3 -m unittest tests/test_models.py
```

---

## Model Weights & Training Status

Rush checkpoints are published directly on Hugging Face Hub at [**`arpit-bhayani/rush`**](https://huggingface.co/arpit-bhayani/rush):

* **Pretrained Trunk**: Uses `answerdotai/ModernBERT-base` (149M parameters, 768 hidden dimension) for bidirectional state representations and token offset resolution.
* **Semantic Option Bank**: Trained on `mteb/banking77` (9,993 training examples across 77 intents) with frozen ModernBERT embeddings and a learned metric tensor $W \in \mathbb{R}^{256 \times 256}$. Achieves **63.13% Top-1** and **89.43% Top-5** test accuracy in candidate space ([`semantic_option_bank.pt`](https://huggingface.co/arpit-bhayani/rush/blob/main/semantic_option_bank.pt)).
* **Conformal Risk Control**: Split Conformal Prediction calibrated on held-out Banking77 non-conformity scores. Achieves **100.00%** coverage for $\alpha=0.01$ and **95.94%** coverage for $\alpha=0.05$ on unseen test splits ([`conformal_calibration.json`](https://huggingface.co/arpit-bhayani/rush/blob/main/conformal_calibration.json)).
* **Extractive Span Pointer**: Bilinear token boundary projector trained on `rajpurkar/squad` for fast non-autoregressive entity extraction without LLM token generation ([`span_pointer.pt`](https://huggingface.co/arpit-bhayani/rush/blob/main/span_pointer.pt)).

> [!TIP]
> **Zero-Config Automatic Downloads**: Calling `rush.load_agent()` automatically pulls any missing checkpoints directly from [Hugging Face Hub (`arpit-bhayani/rush`)](https://huggingface.co/arpit-bhayani/rush) and caches them locally. Retraining pipelines are provided in [`rush/train/`](rush/train/).

---

## Acknowledgments

`rush` builds upon ideas, architectures, and benchmarks from earlier work:
* **[Jev](https://typesafe.ai)** (TypeSafe AI) for pioneering the fast, non-autoregressive "System 1" decision engine paradigm.
* **[Laya](https://github.com/NandhaKishorM/laya)** (by Nandha Kishor M) for the open-weights ModernBERT decision engine and task evaluation suite.
* **[Laya MLX](https://github.com/mizorewww/laya-mlx)** (by mizorewww) for the Apple Silicon MLX port and local execution insights.

---

## License

MIT License. See [LICENSE](LICENSE) for details.
