Skip to main content

Decoder

Struct Decoder 

Source
pub struct Decoder { /* private fields */ }
Expand description

载入好的解码器。

Implementations§

Source§

impl Decoder

Source

pub fn load(dir: &Path) -> Result<Decoder, String>

dir 是模型目录:model.safetensors、config.json、generation_config.json。

Source

pub fn from_weights( w: &Weights, text: &Value, generation: &Value, ) -> Result<Decoder, String>

从已读入的权重构造(和图像编码器共用一次读文件时用)。 text 是 config.json 的 text_config,generation 是 generation_config.json。

Examples found in repository?
examples/formula_check.rs (lines 46-50)
40fn main() {
41    let a: Vec<String> = std::env::args().skip(1).collect();
42    let dir = Path::new(&a[0]);
43    let t = Instant::now();
44    let cfg = json(dir, "config.json");
45    let w = Weights::from_bytes(&std::fs::read(dir.join("model.safetensors")).unwrap()).unwrap();
46    let dec = Decoder::from_weights(
47        &w,
48        &cfg["text_config"],
49        &json(dir, "generation_config.json"),
50    )
51    .unwrap();
52    let enc = Encoder::from_weights(&w, Config::from_json(&cfg).unwrap()).unwrap();
53    drop(w);
54    let tok = Tokenizer::load(dir).unwrap();
55    println!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
56
57    if a.len() == 1 {
58        // 固定种子的随机 memory
59        let mut s = 0x9E37_79B9_7F4A_7C15u64;
60        let data = (0..144 * 512)
61            .map(|_| {
62                s ^= s << 13;
63                s ^= s >> 7;
64                s ^= s << 17;
65                (s >> 40) as f32 / (1u64 << 23) as f32 - 1.0
66            })
67            .collect();
68        let mem = Tensor::new(vec![144, 512], data);
69        let t = Instant::now();
70        let ids = dec.generate(&mem, 200).unwrap();
71        let ms = t.elapsed().as_secs_f64() * 1e3;
72        println!(
73            "随机 memory:{} 个 token,{ms:.0} ms,{:.1} ms/token",
74            ids.len(),
75            ms / (ids.len() + 1) as f64
76        );
77        println!("{}", to_latex(&tok.decode(&ids)));
78        return;
79    }
80
81    for p in &a[1..] {
82        let img = read_ppm(p);
83        let t = Instant::now();
84        let x = preprocess(&img);
85        let t_pre = t.elapsed().as_secs_f64() * 1e3;
86        let t = Instant::now();
87        let mem = enc.encode(&x).unwrap();
88        let t_enc = t.elapsed().as_secs_f64() * 1e3;
89        let t = Instant::now();
90        let ids = dec.generate(&mem, dec.max_length()).unwrap();
91        let t_dec = t.elapsed().as_secs_f64() * 1e3;
92        let name = Path::new(p).file_name().unwrap().to_string_lossy();
93        println!(
94            "{name}\t预处理 {t_pre:.0} ms\t编码 {t_enc:.0} ms\t解码 {t_dec:.0} ms({} 个 token,{:.1} ms/token)",
95            ids.len(),
96            t_dec / (ids.len() + 1) as f64
97        );
98        println!("  原始:{}", tok.decode(&ids));
99        println!("  LaTeX:{}", to_latex(&tok.decode(&ids)));
100    }
101    // 内存峰值(Linux 才有这个文件)
102    if let Ok(s) = std::fs::read_to_string("/proc/self/status")
103        && let Some(l) = s.lines().find(|l| l.starts_with("VmHWM"))
104    {
105        println!("{l}");
106    }
107}
Source

pub fn max_length(&self) -> usize

generation_config.json 的 max_length(含起始 token 的序列总长),传给 generate。

Examples found in repository?
examples/formula_check.rs (line 90)
40fn main() {
41    let a: Vec<String> = std::env::args().skip(1).collect();
42    let dir = Path::new(&a[0]);
43    let t = Instant::now();
44    let cfg = json(dir, "config.json");
45    let w = Weights::from_bytes(&std::fs::read(dir.join("model.safetensors")).unwrap()).unwrap();
46    let dec = Decoder::from_weights(
47        &w,
48        &cfg["text_config"],
49        &json(dir, "generation_config.json"),
50    )
51    .unwrap();
52    let enc = Encoder::from_weights(&w, Config::from_json(&cfg).unwrap()).unwrap();
53    drop(w);
54    let tok = Tokenizer::load(dir).unwrap();
55    println!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
56
57    if a.len() == 1 {
58        // 固定种子的随机 memory
59        let mut s = 0x9E37_79B9_7F4A_7C15u64;
60        let data = (0..144 * 512)
61            .map(|_| {
62                s ^= s << 13;
63                s ^= s >> 7;
64                s ^= s << 17;
65                (s >> 40) as f32 / (1u64 << 23) as f32 - 1.0
66            })
67            .collect();
68        let mem = Tensor::new(vec![144, 512], data);
69        let t = Instant::now();
70        let ids = dec.generate(&mem, 200).unwrap();
71        let ms = t.elapsed().as_secs_f64() * 1e3;
72        println!(
73            "随机 memory:{} 个 token,{ms:.0} ms,{:.1} ms/token",
74            ids.len(),
75            ms / (ids.len() + 1) as f64
76        );
77        println!("{}", to_latex(&tok.decode(&ids)));
78        return;
79    }
80
81    for p in &a[1..] {
82        let img = read_ppm(p);
83        let t = Instant::now();
84        let x = preprocess(&img);
85        let t_pre = t.elapsed().as_secs_f64() * 1e3;
86        let t = Instant::now();
87        let mem = enc.encode(&x).unwrap();
88        let t_enc = t.elapsed().as_secs_f64() * 1e3;
89        let t = Instant::now();
90        let ids = dec.generate(&mem, dec.max_length()).unwrap();
91        let t_dec = t.elapsed().as_secs_f64() * 1e3;
92        let name = Path::new(p).file_name().unwrap().to_string_lossy();
93        println!(
94            "{name}\t预处理 {t_pre:.0} ms\t编码 {t_enc:.0} ms\t解码 {t_dec:.0} ms({} 个 token,{:.1} ms/token)",
95            ids.len(),
96            t_dec / (ids.len() + 1) as f64
97        );
98        println!("  原始:{}", tok.decode(&ids));
99        println!("  LaTeX:{}", to_latex(&tok.decode(&ids)));
100    }
101    // 内存峰值(Linux 才有这个文件)
102    if let Ok(s) = std::fs::read_to_string("/proc/self/status")
103        && let Some(l) = s.lines().find(|l| l.starts_with("VmHWM"))
104    {
105        println!("{l}");
106    }
107}
Source

pub fn generate( &self, memory: &Tensor, max_len: usize, ) -> Result<Vec<u32>, String>

贪心生成。memory 是编码器输出 [L, d_model](参考实现的 encoder_hidden_states)。

max_len 的含义同 generate(max_length=…):序列总长的上限,含起始 token;序列长到 max_len − 1 时下一个 token 强制为 forced_eos_token_id (ForcedEOSTokenLogitsProcessor)。超过位置向量表的长度时取表长。

返回生成的 token,不含起始 token 和结尾的 eos。

Examples found in repository?
examples/formula_check.rs (line 70)
40fn main() {
41    let a: Vec<String> = std::env::args().skip(1).collect();
42    let dir = Path::new(&a[0]);
43    let t = Instant::now();
44    let cfg = json(dir, "config.json");
45    let w = Weights::from_bytes(&std::fs::read(dir.join("model.safetensors")).unwrap()).unwrap();
46    let dec = Decoder::from_weights(
47        &w,
48        &cfg["text_config"],
49        &json(dir, "generation_config.json"),
50    )
51    .unwrap();
52    let enc = Encoder::from_weights(&w, Config::from_json(&cfg).unwrap()).unwrap();
53    drop(w);
54    let tok = Tokenizer::load(dir).unwrap();
55    println!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
56
57    if a.len() == 1 {
58        // 固定种子的随机 memory
59        let mut s = 0x9E37_79B9_7F4A_7C15u64;
60        let data = (0..144 * 512)
61            .map(|_| {
62                s ^= s << 13;
63                s ^= s >> 7;
64                s ^= s << 17;
65                (s >> 40) as f32 / (1u64 << 23) as f32 - 1.0
66            })
67            .collect();
68        let mem = Tensor::new(vec![144, 512], data);
69        let t = Instant::now();
70        let ids = dec.generate(&mem, 200).unwrap();
71        let ms = t.elapsed().as_secs_f64() * 1e3;
72        println!(
73            "随机 memory:{} 个 token,{ms:.0} ms,{:.1} ms/token",
74            ids.len(),
75            ms / (ids.len() + 1) as f64
76        );
77        println!("{}", to_latex(&tok.decode(&ids)));
78        return;
79    }
80
81    for p in &a[1..] {
82        let img = read_ppm(p);
83        let t = Instant::now();
84        let x = preprocess(&img);
85        let t_pre = t.elapsed().as_secs_f64() * 1e3;
86        let t = Instant::now();
87        let mem = enc.encode(&x).unwrap();
88        let t_enc = t.elapsed().as_secs_f64() * 1e3;
89        let t = Instant::now();
90        let ids = dec.generate(&mem, dec.max_length()).unwrap();
91        let t_dec = t.elapsed().as_secs_f64() * 1e3;
92        let name = Path::new(p).file_name().unwrap().to_string_lossy();
93        println!(
94            "{name}\t预处理 {t_pre:.0} ms\t编码 {t_enc:.0} ms\t解码 {t_dec:.0} ms({} 个 token,{:.1} ms/token)",
95            ids.len(),
96            t_dec / (ids.len() + 1) as f64
97        );
98        println!("  原始:{}", tok.decode(&ids));
99        println!("  LaTeX:{}", to_latex(&tok.decode(&ids)));
100    }
101    // 内存峰值(Linux 才有这个文件)
102    if let Ok(s) = std::fs::read_to_string("/proc/self/status")
103        && let Some(l) = s.lines().find(|l| l.starts_with("VmHWM"))
104    {
105        println!("{l}");
106    }
107}

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.