web-rwkv 0.10.20

An implementation of the RWKV language model in pure WebGPU.
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

use half::f16;
use serde::{Deserialize, Serialize};

use super::{JobInfo, JobInput};
use crate::tensor::{TensorCpu, TensorInit};

pub mod rnn;
pub mod vision;

pub use rnn::{
    Rnn, RnnChunk, RnnChunkBatch, RnnInfo, RnnInfoBatch, RnnInput, RnnInputBatch, RnnIter,
    RnnOption, RnnOutput, RnnOutputBatch, RnnRedirect,
};

pub trait Infer: Send + Sync + 'static {
    type Info: JobInfo;
    type Input: JobInput;
    type Output: Send + Sync + 'static;
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Token {
    Token(u32),
    Embed(TensorCpu<f16>),
}

impl Default for Token {
    fn default() -> Self {
        Self::Token(0)
    }
}

impl From<u32> for Token {
    fn from(value: u32) -> Self {
        Self::Token(value)
    }
}

impl From<u16> for Token {
    fn from(value: u16) -> Self {
        Self::Token(value as u32)
    }
}

impl From<Vec<f16>> for Token {
    fn from(value: Vec<f16>) -> Self {
        Self::Embed(TensorCpu::from_data_1d(value))
    }
}

impl From<Vec<f32>> for Token {
    fn from(value: Vec<f32>) -> Self {
        let value: Vec<_> = value.into_iter().map(f16::from_f32).collect();
        Self::Embed(TensorCpu::from_data_1d(value))
    }
}