Struct Encoder

Source

pub struct Encoder { /* private fields */ }

Expand description

载入好的编码器。

Implementations§

Source §

impl Encoder

Source

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

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

Examples found in repository ?

examples/formula_enc_check.rs (line 82)

79fn main() {
80    let a: Vec<String> = std::env::args().skip(1).collect();
81    let t = Instant::now();
82    let enc = Encoder::load(std::path::Path::new(&a[0])).unwrap();
83    println!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
84    let img = a.get(1).map_or_else(synthetic, |p| read_ppm(p));
85    let t = Instant::now();
86    let x = preprocess(&img);
87    let pre_ms = t.elapsed().as_secs_f64() * 1e3;
88    stats("preprocess", &x, pre_ms);
89    let t0 = Instant::now();
90    let mut last = t0;
91    let m = enc
92        .encode_traced(&x, |name, t| {
93            stats(name, t, last.elapsed().as_secs_f64() * 1e3);
94            last = Instant::now();
95        })
96        .unwrap();
97    println!(
98        "编码合计 {:.0} ms，输出 {:?}",
99        t0.elapsed().as_secs_f64() * 1e3,
100        m.shape
101    );
102    // 不带逐段统计再跑一次，得到干净的耗时
103    let t = Instant::now();
104    let m2 = enc.encode(&x).unwrap();
105    println!("第二次编码 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
106    assert_eq!(m.data, m2.data, "两次编码结果应逐位相同");
107}

Source

pub fn from_weights(w: &Weights, cfg: Config) -> Result<Encoder, String>

从已读入的权重建（解码器可共用同一份 Weights）。用到的参数都复制出来，之后可丢掉 w。

Examples found in repository ?

examples/formula_check.rs (line 52)

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 config(&self) -> &Config

Source

pub fn encode(&self, x: &Tensor) -> Result<Tensor, String>

x 是 preprocess 的输出 [3, 768, 768]，返回解码器交叉注意力的 memory [144, 512]。

Examples found in repository ?

examples/formula_enc_check.rs (line 104)

79fn main() {
80    let a: Vec<String> = std::env::args().skip(1).collect();
81    let t = Instant::now();
82    let enc = Encoder::load(std::path::Path::new(&a[0])).unwrap();
83    println!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
84    let img = a.get(1).map_or_else(synthetic, |p| read_ppm(p));
85    let t = Instant::now();
86    let x = preprocess(&img);
87    let pre_ms = t.elapsed().as_secs_f64() * 1e3;
88    stats("preprocess", &x, pre_ms);
89    let t0 = Instant::now();
90    let mut last = t0;
91    let m = enc
92        .encode_traced(&x, |name, t| {
93            stats(name, t, last.elapsed().as_secs_f64() * 1e3);
94            last = Instant::now();
95        })
96        .unwrap();
97    println!(
98        "编码合计 {:.0} ms，输出 {:?}",
99        t0.elapsed().as_secs_f64() * 1e3,
100        m.shape
101    );
102    // 不带逐段统计再跑一次，得到干净的耗时
103    let t = Instant::now();
104    let m2 = enc.encode(&x).unwrap();
105    println!("第二次编码 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
106    assert_eq!(m.data, m2.data, "两次编码结果应逐位相同");
107}

More examples

Hide additional examples

examples/formula_check.rs (line 87)

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 encode_traced( &self, x: &Tensor, trace: impl FnMut(&str, &Tensor), ) -> Result<Tensor, String>

同 encode，每一段的输出都交给 trace(名字, 张量) 看一眼（检查用）。

Examples found in repository ?

examples/formula_enc_check.rs (lines 92-95)

79fn main() {
80    let a: Vec<String> = std::env::args().skip(1).collect();
81    let t = Instant::now();
82    let enc = Encoder::load(std::path::Path::new(&a[0])).unwrap();
83    println!("载入 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
84    let img = a.get(1).map_or_else(synthetic, |p| read_ppm(p));
85    let t = Instant::now();
86    let x = preprocess(&img);
87    let pre_ms = t.elapsed().as_secs_f64() * 1e3;
88    stats("preprocess", &x, pre_ms);
89    let t0 = Instant::now();
90    let mut last = t0;
91    let m = enc
92        .encode_traced(&x, |name, t| {
93            stats(name, t, last.elapsed().as_secs_f64() * 1e3);
94            last = Instant::now();
95        })
96        .unwrap();
97    println!(
98        "编码合计 {:.0} ms，输出 {:?}",
99        t0.elapsed().as_secs_f64() * 1e3,
100        m.shape
101    );
102    // 不带逐段统计再跑一次，得到干净的耗时
103    let t = Instant::now();
104    let m2 = enc.encode(&x).unwrap();
105    println!("第二次编码 {:.0} ms", t.elapsed().as_secs_f64() * 1e3);
106    assert_eq!(m.data, m2.data, "两次编码结果应逐位相同");
107}