Metadata-Version: 2.4
Name: neuratensor
Version: 4.3.0
Summary: Neuromorphic sequence models in pure PyTorch — spiking neurons + selective state-space model, exact streaming inference, CPU/CUDA/MPS
Author-email: "Neuramorphic, Inc." <info@neuramorphic.ai>
License: Apache-2.0
Project-URL: Homepage, https://neuramorphic.ai
Project-URL: Repository, https://github.com/Neuramorphic/neuratensor
Project-URL: Bug Reports, https://github.com/Neuramorphic/neuratensor/issues
Project-URL: Changelog, https://github.com/Neuramorphic/neuratensor/blob/main/CHANGELOG.md
Keywords: neuromorphic,snn,spiking-neural-networks,ssm,mamba,state-space-model,pytorch,deep-learning,streaming-inference
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: License :: OSI Approved :: Apache Software 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: Programming Language :: Python :: 3.13
Classifier: Operating System :: OS Independent
Classifier: Typing :: Typed
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1.0
Requires-Dist: safetensors>=0.4.0
Provides-Extra: server
Requires-Dist: fastapi>=0.100.0; extra == "server"
Requires-Dist: uvicorn[standard]>=0.23.0; extra == "server"
Requires-Dist: pydantic>=2.0.0; extra == "server"
Provides-Extra: hub
Requires-Dist: huggingface_hub>=0.20.0; extra == "hub"
Provides-Extra: dev
Requires-Dist: pytest>=7.4.0; extra == "dev"
Requires-Dist: httpx>=0.24.0; extra == "dev"
Requires-Dist: fastapi>=0.100.0; extra == "dev"
Requires-Dist: ruff>=0.4.0; extra == "dev"
Requires-Dist: build>=1.0.0; extra == "dev"
Requires-Dist: twine>=4.0.0; extra == "dev"
Dynamic: license-file

# NeuraTensor

Neuromorphic sequence models in pure PyTorch: **leaky integrate-and-fire spiking neurons** feeding a **Mamba-style selective state-space model**, with optional sliding-window attention. Runs wherever PyTorch runs (CPU, CUDA, Apple Silicon), no compiled extensions.

```bash
pip install neuratensor
```

Python 3.9+ and PyTorch 2.1+. Tested on PyTorch 2.8.

> **Status: beta.** This package provides the architecture, training loop, LoRA, serving and serialization. It does **not** ship pretrained weights: `from_preset` gives a randomly initialised model to train or benchmark. See [What this is not](#what-this-is-not).

---

## Quickstart

```python
import torch
from neuratensor import NeuraTensor

model = NeuraTensor.from_preset("nano", vocab_size=256)   # random init
logits, states, metrics = model(torch.randint(0, 256, (1, 64)))
print(f"spike rate {metrics['spike_rate']:.3f}")
```

Train on a text file (UTF-8 bytes by default, or pass your own tokenizer to `TextDataset`):

```python
from neuratensor.training import Trainer, TrainConfig

trainer = Trainer(model, TrainConfig(max_steps=2000, batch_size=16, max_seq_len=256), train_data="data.txt")
trainer.train()
model.save_pretrained("my-model")          # config.json + model.safetensors
```

Serve or stream it:

```python
from neuratensor import InferenceEngine

engine = InferenceEngine.from_pretrained("my-model")
out = engine.generate(prompt_ids, max_new_tokens=128, temperature=0.8, top_k=50, seed=0)
for token in engine.stream(prompt_ids, max_new_tokens=128):   # [B, 1] per step
    ...
```

---

## Architecture

```
tokens → Embedding → N × FusedSpikeBlock → RMSNorm → LM head (tied)
                         │
   x → RMSNorm → Linear → [LIF spikes] → [Selective SSM] → × SiLU(gate) → Linear → + residual
                                  optional sliding-window attention after every attn_every-th block
```

* **LIF neurons**: membrane `V ← β·V + I`, binary spike at `V ≥ θ`, soft reset. β learned per neuron; fast-sigmoid surrogate gradient. The scan runs in float32 whatever the model dtype.
* **Selective SSM**: input-dependent `B, C, Δ`; HiPPO-style diagonal `A`; zero-order-hold discretisation. Solved with an exact chunked parallel-prefix scan (stable for any Δ) and activation checkpointing during training.
* **Attention**: causal sliding window with RoPE and a KV cache bounded by the window.
* **Exact streaming**: feeding tokens one at a time (or in any split) reproduces the full forward pass to ~1e-6. Every component, including attention, carries state. This is covered by tests.

### Presets

| Preset   | Params | d_model | Layers | d_state | Attention layers | Max seq |
|----------|--------|---------|--------|---------|------------------|---------|
| `pico`   | 3.5M   | 96      | 4      | 8       | 0                | 512     |
| `nano`   | 14.7M  | 288     | 6      | 16      | 1                | 1024    |
| `small`  | 74.4M  | 640     | 12     | 32      | 2                | 2048    |
| `medium` | 333M   | 1152    | 20     | 64      | 4                | 4096    |
| `large`  | 1.35B  | 2048    | 28     | 128     | 4                | 8192    |

Counts are exact for the 32k vocabulary and tied head (`ModelConfig.n_params_estimate()` is verified against the real model in the tests). Override any field: `NeuraTensor.from_preset("nano", vocab_size=256)`.

---

## Spike sparsity

`metrics["spike_rate"]` is the mean fraction of neurons firing; `estimated_sparsity_pct` is `100 × (1 − rate)`. A freshly initialised model fires on roughly 7-13% of neurons. Training moves this.

Add `TrainConfig(spike_regularization=λ)` to penalise the spike rate (differentiable through the surrogate gradient). In our tests, λ=0.5 cut the rate on training data by about half; very large λ drives the network almost silent, so tune it against your loss.

**Sparsity does not currently make inference faster.** The projections are dense matmuls; v4.0.4's claimed "sparse output projection" never triggered and returned wrong results when forced, so it was removed in 4.3.0. Treat sparsity as a model property (and a prerequisite for neuromorphic hardware or custom sparse kernels), not as a measured speedup.

---

## Performance

Measured on one Apple-silicon laptop, CPU, float32, batch 1, median of several runs. **Your numbers will differ**: run `neuratensor benchmark small` on your hardware.

| Model (forward) | v4.0.4 | v4.3.0 |
|---|---|---|
| `pico`, 64 tokens | 7 ms | 6 ms |
| `nano`, 512 tokens | 99 ms | 111 ms |
| `small`, 64 tokens | 118 ms | 66 ms |
| `small`, 1024 tokens | 2659 ms | 1194 ms |
| `pico` train step (fwd+bwd, 4×256) | 1139 ms | 451 ms |

Larger models and longer sequences benefit most (`small` is about 2× faster); `nano` at T=512 is roughly on par (within run-to-run noise on this machine). The LIF recurrence is inherently sequential, so throughput is bounded by a Python loop over time. A fused CUDA/Triton kernel is the obvious next step and is not implemented.

---

## Custom models, LoRA, streaming state

```python
from neuratensor import NeuraTensorModel, ModelConfig
from neuratensor.training import LoRAConfig, apply_lora, merge_lora

model = NeuraTensorModel(ModelConfig(vocab_size=50_000, d_model=768, n_layers=16, d_state=32,
                                     attn_every=4, attn_heads=8, attn_window=512))

apply_lora(model, LoRAConfig(r=16, alpha=32))     # freezes base weights
# ... fine-tune ...
merge_lora(model)                                  # required before save_pretrained

# streaming: carry state across calls
engine = InferenceEngine(model)
for chunk in token_chunks:
    logits = engine.run(chunk).logits              # continues from previous chunk
engine.reset_state()
```

---

## Loading and saving

| | |
|---|---|
| `model.save_pretrained(dir)` | `config.json` + `model.safetensors` |
| `NeuraTensorModel.from_pretrained(path_or_repo)` | local directory, legacy `.pt`, or Hugging Face repo id (`pip install 'neuratensor[hub]'`) |
| `InferenceEngine.save(path)` / `.from_checkpoint(path)` | single `.pt` file |

Checkpoints are loaded with `torch.load(weights_only=True)`; a malicious pickle is rejected rather than executed (tested). Older PyTorch (< 2.4) may not be able to load legacy `.pt` files that embed the config dataclass; re-save them with `save_pretrained`.

---

## Inference server

```bash
pip install 'neuratensor[server]'
export NEURATENSOR_API_KEY=change-me           # optional bearer-token auth
neuratensor serve --checkpoint my-model --port 8080
```

`GET /health`, `GET /info`, `POST /infer`, `POST /generate`, `GET /metrics` (Prometheus text). Binds to `127.0.0.1` by default; pass `--host 0.0.0.0` only together with an API key. Requests are validated (batch, length, token range) and model calls run off the event loop.

```bash
docker build -t neuratensor . && docker run -p 8080:8080 -e NEURATENSOR_API_KEY=change-me neuratensor
```

---

## CLI

```bash
neuratensor info                 # device + backend
neuratensor presets              # presets with parameter counts
neuratensor benchmark small      # latency / throughput
neuratensor serve --preset small
neuratensor telemetry [status|enable|disable]
```

---

## Telemetry

Anonymous usage statistics (SDK/Python version, OS, CPU arch, preset, device type, batch/sequence length, latency, spike rate, and a random install ID; never inputs, outputs or personal data) are sent when you use `InferenceEngine`. **Nothing is sent on `import neuratensor`.** A notice is printed once to stderr before the first event.

Off automatically in CI and when `DO_NOT_TRACK=1`. Disable explicitly with any of:

```bash
neuratensor telemetry disable
export NEURATENSOR_TELEMETRY=0
```

Source: [`neuratensor/telemetry.py`](neuratensor/telemetry.py).

---

## What this is not

* Not a pretrained language model: no weights are distributed, and `from_pretrained` only loads what you point it at.
* No claim of accuracy parity with Transformers or Mamba. A small, reproducible check in the test suite shows the SSM can learn a long-range copy task (the v4.0.4 build could not), nothing more.
* No energy-efficiency or hardware-speedup claims.

## Development

```bash
pip install -e '.[dev,server]'
pytest
```

CI runs the suite on Python 3.9-3.12. Releases are published from GitHub Actions via PyPI Trusted Publishing.

## License

Apache 2.0, see [LICENSE](LICENSE). Contact: support@neuramorphic.ai
