multiscreen-rs 0.2.1

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

use burn::{
    module::Module,
    tensor::{backend::Backend, Int, Tensor},
};

/// Minimal language-model trait used by the Burn-backed Multiscreen model.
#[allow(dead_code)]
pub trait LanguageModel<B: Backend> {
    fn forward(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3>;
}

/// Trainable Burn language model.
#[allow(dead_code)]
pub trait TrainableLanguageModel<B: Backend>: LanguageModel<B> + Module<B> {
    fn parameter_count(&self) -> usize {
        self.num_params()
    }
}