ai00-core 0.5.1

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
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

use bit_set::BitSet;
use bnf_sampler::sampler::{AcceptTokenResult, PossibleTokensResult, Sampler};

use super::Transformer;

#[derive(Debug)]
pub struct BnfSampler {
    sampler: Sampler,
    current_token_ids: BitSet,
}

impl BnfSampler {
    pub fn new(mut sampler: Sampler) -> Self {
        let current_token_ids = match sampler.all_possible_next_tokens(None) {
            Ok(PossibleTokensResult::Continue(tokens)) => tokens.clone(),
            _ => BitSet::new(),
        };
        Self {
            sampler,
            current_token_ids,
        }
    }

    #[inline]
    pub fn current_token_ids(&self) -> &BitSet {
        &self.current_token_ids
    }
}

impl std::ops::Deref for BnfSampler {
    type Target = Sampler;

    fn deref(&self) -> &Self::Target {
        &self.sampler
    }
}

impl std::ops::DerefMut for BnfSampler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.sampler
    }
}

impl Transformer for BnfSampler {
    fn transform(&self, output: &mut [f32]) {
        output
            .iter_mut()
            .enumerate()
            .filter(|&(token, _)| !self.current_token_ids().contains(token))
            .for_each(|(_, logits)| *logits = f32::MIN)
    }

    fn update(&mut self, token: u16) -> bool {
        let token = Some(token as u32);
        let accept = self.accept_a_token(token).expect("invalid input token");
        self.current_token_ids = match self.sampler.all_possible_next_tokens(None) {
            Ok(PossibleTokensResult::Continue(tokens)) => tokens.clone(),
            _ => BitSet::new(),
        };
        match accept {
            AcceptTokenResult::Continue => false,
            AcceptTokenResult::End | AcceptTokenResult::Failed => true,
        }
    }
}