ALBERT (A Lite BERT) masked language model implementation in Rust using Burn.

ALBERT uses factorized embedding parameterization and cross-layer parameter sharing to reduce model size while maintaining performance.

Supports all v2 variants from HuggingFace:

use burn::backend::ndarray::NdArray;
use albert_burn::{AlbertMaskedLM, AlbertVariant, tokenize_batch};

type B = NdArray<f32>;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = Default::default();

    // Load pretrained model and tokenizer (downloads from HuggingFace)
    let (model, tokenizer) = AlbertMaskedLM::<B>::pretrained(&device, AlbertVariant::BaseV2, None)?;

    // Tokenize input with [MASK] token
    let sentence = "The capital of France is [MASK].";
    let (input_ids, attention_mask) = tokenize_batch::<B>(&tokenizer, &[sentence], &device);

    // Forward pass returns logits over vocabulary
    let logits = model.forward(input_ids, attention_mask, None);

    Ok(())
}

• pretrained - Enables model download utilities (default)
• ndarray - NdArray backend

Backend features:

• wgpu - WebGPU backend
• cuda - CUDA backend
• tch-cpu - LibTorch CPU backend
• tch-gpu - LibTorch GPU backend

Run the fill-mask inference example:

cargo run --example inference --features "pretrained,ndarray" --release

Specify a variant:

cargo run --example inference --features "pretrained,ndarray" --release -- xxlarge

Prompt: "The capital of France is [MASK]."

BaseV2 (12M params):

LargeV2 (18M params):

XLargeV2 (60M params):

XXLargeV2 (235M params):

Integration tests (requires model download):

cargo test --features "pretrained,ndarray" -- --ignored

Tests verify logit values, top-5 predictions, statistics, and per-position L2 norms across 3 sentences against Python HuggingFace reference at 5e-4 relative tolerance.

MIT OR Apache-2.0