Crate multiscreen_rs 

multiscreen-rs

A Rust implementation of the Multiscreen neural language model — training and inference — powered by Burn. Based on Screening Is Enough.

Quick Start

Training

use multiscreen_rs::prelude::*;

fn main() -> multiscreen_rs::Result<()> {
    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 sequences = vec![vec![1, 2, 3, 4], vec![1, 2, 5, 4]];
    let report = trainer.train_on_token_sequences(&sequences)?;
    // or: trainer.train_on_chat_sequences(&chat_pairs)?;
    Ok(())
}

Inference

use multiscreen_rs::prelude::*;

fn main() -> multiscreen_rs::Result<()> {
    let model = ChatModel::load("checkpoints/latest.mpk")?;

    // One-shot
    let tokens = model.generate(&[1, 2, 3], GenerationConfig::default())?;

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

Feature Flags

Enable the cuda feature for NVIDIA GPU acceleration. Default uses Burn Flex for CPU training with auto SIMD detection.

Re-exports

pub use device::cpu;

pub use device::auto_device;

pub use device::cuda;

pub use inference::ChatModel;

pub use inference::GenerationConfig;

pub use training::ParameterBudget;

pub use training::Trainer;

pub use training::TrainingReport;

Modules

device

Device abstraction for CPU and CUDA.

inference

High-level inference API.

prelude

training

High-level training API with a builder pattern.

Structs

EvaluationResult

Result of evaluating a model on held-out sequences.

InferenceConfig

Controls deterministic inference behavior for the lightweight transition engine (MultiscreenEngine).

InferenceOutput

Output from token inference.

Int

A type-level representation of the kind of a int tensor.

ModelInferenceConfig

Greedy generation options.

ModelTrainingConfig

Training options for token-sequence training.

ModelTrainingReport

MultiscreenConfig

Full engine configuration.

MultiscreenEngine

Main user-facing engine.

MultiscreenModel

Burn-backed neural Multiscreen language model ported from multiscreen-testing.

MultiscreenModelConfig

Paper-faithful neural Multiscreen model configuration.

MultiscreenModelOutput

Screen

A screen over a contiguous token span.

ScreenConfig

Controls how a token sequence is split into screens.

ScreenLayout

Fully materialized screen and tile layout for a token sequence.

ScreeningGridConfig

Logical N_L x N_H screening grid from the Multiscreen paper.

Tensor

A tensor with a given backend, shape and data type.

TensorData

Data structure for tensors.

Tile

A tile over a contiguous token span, mapped to paper-style layer/head slots.

TileConfig

Controls how screens are split into sliding tiles.

TokenSpan

Half-open token span [start, end).

TrainReport

Summary returned after training.

TrimConfig

Multi-screen scoring parameters.

Enums

Error

Top-level crate error.

MultiscreenParameterBudget

Supported neural Multiscreen parameter budgets.

TrainInput

User-provided training input.

Traits

AutodiffBackend

Trait that allows a backend to support autodiff.

Backend

This trait defines all types and functions needed for a backend to be used with burn.

Functions

causal_softmask

Causal softmask for query/key token positions.

causal_trim_relevance

Combined causal trim relevance for a query/key pair.

cross_entropy_loss_with_mask

default_device

Chooses the default device based on the active backend.

device_label

Human-readable label for the active backend device.

trim_and_square

Trim-and-square relevance gate from the multiscreen experiment.

Type Aliases

DefaultAutodiffBackend

Default training backend. Burn’s Autodiff decorator adds backpropagation.

DefaultBackend

Default CPU backend (Flex) — used when the cuda feature is not enabled.

DefaultMultiscreenModel

Convenience alias for the default Burn Flex autodiff model.

Device

Device type for the active backend.

GridConfigDeprecated

MultiScreenConfigDeprecated

MultiScreenEngineDeprecated

Result

Result alias used by the crate.