ai00-core 0.5.12

An implementation of the RWKV language model in pure WebGPU.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

pub mod bnf;
pub mod mirostat;
pub mod nucleus;
pub mod typical;

mod radix;

pub trait Sampler {
    /// Initialize the sampler state.
    fn init(&mut self, model_tokens: &[u16]);
    /// Update the raw model output.
    fn transform(&self, output: &mut [f32]);
    /// Select one token from the distribution, and also update the state.
    fn sample(&mut self, probs: &[f32]) -> u16;
}

pub trait Formatter {
    /// Update the raw model output.
    fn transform(&self, output: &mut [f32]);
    /// Update the internal state after a token is chosen. Return if the state machine is halt.
    fn update(&mut self, token: u16) -> bool;
}