multiscreen-rs 0.3.0

A Rust implementation of the Multiscreen neural language model — training and inference powered by Burn.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

# multiscreen-rs

A Rust implementation of the Multiscreen neural language model — training and inference — powered by [Burn](https://github.com/tracel-ai/burn). Based on [Screening Is Enough](https://arxiv.org/pdf/2604.01178).

## Quickstart

```bash
# Train (10M params, ~25 min on CUDA)
cargo run --release --features cuda --example train_with_tokenizer -- \
    --budget 10m --steps 10000 --batch-size 16 --seq-len 256

# Chat
cargo run --release --features cuda --example chat_with_tokenizer -- --run-dir runs/my-model
```

Data and tokenizer bundled in `examples/data/`. No setup needed.

---

## CLI

### Train

```bash
# Quick test (CPU, ~1 min)
cargo run --release --example train_with_tokenizer -- --budget 1m --steps 500

# 10M on CUDA (~25 min)
cargo run --release --features cuda --example train_with_tokenizer -- \
    --budget 10m --steps 10000 --batch-size 16 --seq-len 256

# 50M on CUDA (~2-3 hrs)
cargo run --release --features cuda --example train_with_tokenizer -- \
    --budget 50m --steps 10000 --batch-size 4 --seq-len 256

# Your own data (put tokenizer.model + .jsonl in a directory)
cargo run --release --features cuda --example train_with_tokenizer -- \
    --train-dir /path/to/data --run-dir runs/custom --budget 10m --steps 50000
```

### Chat

```bash
# Interactive (streaming)
cargo run --release --features cuda --example chat_with_tokenizer -- --run-dir runs/my-model

# One-shot
cargo run --release --features cuda --example chat_with_tokenizer -- \
    --run-dir runs/my-model --prompt "Once upon a time"

# Custom sampling
cargo run --release --features cuda --example chat_with_tokenizer -- \
    --run-dir runs/my-model --temperature 0.8 --top-k 40 --max-new-tokens 256
```

> Run with `--help` to see all options (training & chat).

---

## Data Formats

Put training data alongside `tokenizer.model` in any directory:

| Format | Description |
|---|---|
| `.txt` | Blank lines separate samples; split 60/40 into prompt/response |
| `.jsonl` text | `{"text": "your text here"}` — one sample per line |
| `.jsonl` chat | `{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}` — OpenAI format |

> **Tip:** Aim for ≥10× more tokens than parameters. 10M params → ~83M+ tokens.

---

## Library Usage

```toml
[dependencies]
multiscreen-rs = "0.2"
# multiscreen-rs = { version = "0.2", features = ["cuda"] }
```

```rust
use multiscreen_rs::prelude::*;

// Training
let mut trainer = Trainer::builder()
    .vocab_size(1000)
    .budget(ParameterBudget::Params10M)
    .device(auto_device()?)
    .batch_size(16).seq_len(128).steps(50_000)
    .build()?;

let report = trainer.train_on_token_sequences(&sequences)?;
// or: trainer.train_on_chat_sequences(&chat_pairs)?;

// Inference
let model = ChatModel::load("checkpoints/latest.mpk")?;
let tokens = model.generate(&[1, 2, 3], GenerationConfig::default())?;

model.generate_stream(&[1, 2, 3], GenerationConfig::default(), |id, _| {
    print!("{id} "); true
})?;
```

---

## Architecture

```
Transformer:    score = QK^T/√d → softmax → attn · V
Multiscreen:    alpha = clamp(1 - r*(1 - QK^T), 0, 1)²   (trim-and-square)
                relevance = alpha * causal_cosine_window   (NO softmax)
                out = tanh_norm(relevance · V) * silu_gate(g)
```

No softmax → no attention dilution at long contexts. Advantages show at seq_len 4K+.

Parameter budgets: `1m` `5m` `10m` `50m` `100m`

## License

[MIT](LICENSE)