Metadata-Version: 2.4
Name: mcr-attention
Version: 2.0.0
Summary: MCR-Attention V2.0: Multi-Scale Compressed Recurrent Attention Model for efficient language modeling with adaptive attention routing.
Author-email: MCR-Attention Authors <mcr-attention@example.com>
License: MIT
Project-URL: Homepage, https://github.com/bueormnew/MCRA
Project-URL: Repository, https://github.com/bueormnew/MCRA
Project-URL: Documentation, https://github.com/bueormnew/MCRA#readme
Project-URL: Issues, https://github.com/bueormnew/MCRA/issues
Keywords: attention,recurrent,language-model,transformer,pytorch,deep-learning
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: hypothesis>=6.0; extra == "dev"
Provides-Extra: benchmark
Requires-Dist: matplotlib>=3.5; extra == "benchmark"
Provides-Extra: all
Requires-Dist: mcr-attention[benchmark,dev]; extra == "all"
Dynamic: license-file

# MCR-Attention V2.0

**Multi-Scale Compressed Recurrent Attention Model**

A PyTorch library for building, training, and deploying language models using multi-scale recurrent memories with adaptive attention routing. MCR-Attention replaces traditional self-attention with O(N) linear recurrences at multiple time scales, enabling efficient processing of long sequences without the quadratic cost of transformers.

## Features

- **Multi-scale recurrent memory** — K independent decay scales capture patterns from short-range to long-range context
- **Parallel associative scan** — O(N) time, O(log N) parallel depth for training
- **Streaming inference** — Single-token updates with O(1) per-step cost
- **Attention routing** — Learned soft-attention over compressed memory vectors
- **Rotary position encoding (RoPE)** — Position-aware representations
- **Simple high-level API** — Create, train, save, load, and generate in a few lines

## Installation

```bash
pip install mcr-attention
```

For development (includes pytest, hypothesis):
```bash
pip install mcr-attention[dev]
```

For benchmarking (adds matplotlib):
```bash
pip install mcr-attention[all]
```

## Quick Start

```python
from mcr_attention import MCRConfig, MCRAttention, create_model, save_model, load_model

# Create a model
config = MCRConfig(vocab_size=32000, d_model=512)
model = create_model(config)

# Or use kwargs directly
model = create_model(vocab_size=32000, d_model=512, num_scales=8)
```

## API Overview

### Core Classes

| Class | Description |
|-------|-------------|
| `MCRConfig` | Model hyperparameters (vocab, dimensions, scales, etc.) |
| `MCRAttention` | The full model (embedding → memory → router → head) |
| `Trainer` | High-level training loop with warmup and checkpointing |
| `GenerationConfig` | Settings for token generation (temp, top-k, top-p) |

### Core Functions

| Function | Description |
|----------|-------------|
| `create_model(config, **kwargs)` | Create a new model instance |
| `save_model(model, path)` | Save model weights + config to directory |
| `load_model(path, device)` | Load model from directory |
| `generate(model, prompt, config)` | Generate tokens from a prompt |

## Configuration Reference

```python
from mcr_attention import MCRConfig

config = MCRConfig(
    vocab_size=32000,    # Token vocabulary size
    d_model=512,         # Embedding dimension
    d_state=128,         # Recurrent hidden state dimension per scale
    num_scales=8,        # Number of memory scales (K)
    scale_min=256,       # Shortest window length
    scale_factor=2,      # Geometric factor between scales
    rope_base=10000.0,   # RoPE base frequency
    dropout=0.1,         # Dropout probability
    use_input_gate=True, # Learned input gate per scale
    tie_embeddings=False, # Tie embedding weights with language head
    mlp_hidden=2048,     # Language head MLP hidden dim
    max_seq_len=32768,   # Maximum sequence length
)
```

## Training

```python
from mcr_attention import MCRConfig, create_model, Trainer
from mcr_attention.api import TrainingConfig
import torch

# Create model
model = create_model(vocab_size=32000, d_model=512)

# Define data source
def data_fn():
    """Return (input_ids [B, N], targets [B]) each call."""
    tokens = torch.randint(0, 32000, (16, 256))
    return tokens[:, :-1], tokens[:, -1]

# Train
config = TrainingConfig(
    learning_rate=5e-4,
    num_steps=1000,
    warmup_steps=100,
    save_dir="checkpoints/",
    save_interval=500,
    device="auto",
)
trainer = Trainer(model, config)
stats = trainer.fit(data_fn)

print(f"Final loss: {stats['final_loss']:.4f}")
```

### Training with Callbacks

```python
losses = []

def track_loss(step, loss, model):
    losses.append(loss)

stats = trainer.fit(data_fn, callbacks=[track_loss])
```

### Training with Evaluation

```python
def eval_fn(model):
    model.eval()
    # Run your evaluation logic
    return {"accuracy": 0.85, "perplexity": 12.3}

stats = trainer.fit(data_fn, eval_fn=eval_fn)
```

## Fine-tuning

```python
from mcr_attention import load_model
from mcr_attention.finetune import finetune, FinetuneConfig, freeze_all_except

# Load pretrained model
model = load_model("pretrained_model/")

# Fine-tune with frozen embeddings and memory
config = FinetuneConfig(
    learning_rate=1e-4,
    num_steps=500,
    freeze_embeddings=True,
    freeze_memory=True,
)
stats = finetune(model, data_fn, config, save_path="finetuned_model/")

# Or freeze everything except the head
freeze_all_except(model, ["head"])
```

## Saving and Loading

```python
from mcr_attention import save_model, load_model, create_model

model = create_model(vocab_size=32000, d_model=512)

# Save (creates directory with config.json, model.pt, metadata.json)
save_model(model, "my_model/", metadata={"trained_on": "my_dataset"})

# Load
model = load_model("my_model/", device="auto")  # auto-detects cuda/mps/cpu
```

## Generation

```python
from mcr_attention import load_model, GenerationConfig
from mcr_attention.api import generate

model = load_model("my_model/")

# Basic generation
tokens = generate(model, prompt=[1, 2, 3, 4, 5], max_tokens=100)

# With sampling controls
config = GenerationConfig(
    max_tokens=200,
    temperature=0.8,
    top_k=50,
    top_p=0.9,
    repetition_penalty=1.2,
    seed=42,
)
tokens = generate(model, prompt=[1, 2, 3], config=config)

# Greedy decoding
tokens = generate(model, prompt=[1, 2, 3], temperature=0.0)
```

## CLI Usage

```bash
# Show model info
mcr-attention info path/to/model/

# Run speed benchmark
mcr-attention benchmark --lengths 32 128 512 2048 --d-model 256

# Generate tokens
mcr-attention generate path/to/model/ --prompt 1 2 3 4 5 --max-tokens 50 --temperature 0.8
```

## Architecture

MCR-Attention processes sequences through four stages:

```
Token IDs → [Embedding + RoPE] → [Multi-Scale Memory] → [Attention Router] → [Language Head] → Logits
```

1. **Embedding Layer**: Token embeddings with Rotary Position Encoding for position-awareness.

2. **Multi-Scale Memory**: K independent recurrent scales, each with a geometric window length (L_k = scale_min × scale_factor^(k-1)). Each scale runs a linear recurrence with learned decay:
   ```
   h_t = a ⊙ h_(t-1) + B @ e_t
   ```
   Parallelized via associative scan during training.

3. **Attention Router**: 1×K soft attention where the query is formed from the last token embedding and the longest-scale memory. Attends over all K compressed memory vectors.

4. **Language Head**: Two-layer MLP (GELU activation) projecting to vocabulary logits. Supports optional embedding weight tying.

### Complexity

| Mode | Time | Space |
|------|------|-------|
| Training (parallel scan) | O(N) | O(N × K × d_state) |
| Inference (streaming) | O(1) per token | O(K × d_state) |

## Development

```bash
# Clone and install in dev mode
git clone https://github.com/bueormnew/MCRA.git
cd MCRA
pip install -e ".[dev]"

# Run tests
pytest tests/ -v

# Run specific test
pytest tests/test_model.py -v
```

## License

MIT — see [LICENSE](LICENSE) for details.
