realizar 0.8.4

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154


#[cfg(test)]
mod tests {
    use super::*;
    use crate::apr::AprV2Model;

    #[test]
    fn test_factory_builds_valid_apr() {
        let data = build_executable_pygmy_apr();
        let model = AprV2Model::from_bytes(data);
        assert!(model.is_ok());
    }

    #[test]
    fn test_factory_apr_forward_works() {
        let data = build_executable_pygmy_apr();
        let model = AprV2Model::from_bytes(data).unwrap();
        let logits = model.forward(&[1]);
        assert!(logits.is_ok());
    }

    // =========================================================================
    // GH-194: Weight Tying Tests (token_embd.weight as lm_head)
    // =========================================================================

    /// GH-194: GGUF naming convention model must load successfully
    #[test]
    fn test_gh194_gguf_names_model_loads() {
        let data = build_executable_pygmy_apr_gguf_names();
        let model = AprV2Model::from_bytes(data);
        assert!(model.is_ok(), "GH-194: GGUF-named model must load");
    }

    /// GH-194: GGUF model with token_embd.weight (no lm_head) must find lm_head tensor
    #[test]
    fn test_gh194_gguf_names_finds_lm_head_via_token_embd() {
        let data = build_executable_pygmy_apr_gguf_names();
        let model = AprV2Model::from_bytes(data).unwrap();

        // The model should find token_embd.weight when looking for lm_head
        let lm_head_candidates = [
            "lm_head.weight",
            "output.weight",
            "model.embed_tokens.weight",
            "embed_tokens.weight",
            "token_embd.weight",
        ];

        let found = lm_head_candidates
            .iter()
            .any(|&name| model.get_tensor(name).is_some());

        assert!(
            found,
            "GH-194: Model must find at least one lm_head candidate. \
             token_embd.weight should be found for weight tying."
        );
    }

    /// GH-194: GGUF model forward pass must work with weight tying
    #[test]
    fn test_gh194_gguf_names_forward_works() {
        let data = build_executable_pygmy_apr_gguf_names();
        let model = AprV2Model::from_bytes(data).unwrap();

        // Forward should work even without explicit lm_head.weight
        let result = model.forward(&[1]);
        assert!(
            result.is_ok(),
            "GH-194: Forward must work with token_embd.weight as lm_head. Error: {:?}",
            result.err()
        );

        // Logits should have vocab_size elements
        let logits = result.unwrap();
        assert_eq!(
            logits.len(),
            10,
            "GH-194: Logits must have vocab_size=10 elements"
        );
    }

    /// GH-194: embed_tokens.weight tied model must work
    #[test]
    fn test_gh194_embed_tied_forward_works() {
        let data = build_executable_pygmy_apr_embed_tied();
        let model = AprV2Model::from_bytes(data).unwrap();

        let result = model.forward(&[1]);
        assert!(
            result.is_ok(),
            "GH-194: Forward must work with embed_tokens.weight as lm_head. Error: {:?}",
            result.err()
        );
    }

    /// GH-194: Tensor count must match expected count (no silent drops)
    #[test]
    fn test_gh194_tensor_count_preserved() {
        // HuggingFace naming: 12 tensors
        let hf_data = build_executable_pygmy_apr();
        let hf_model = AprV2Model::from_bytes(hf_data).unwrap();
        assert_eq!(
            hf_model.tensor_count(),
            12,
            "HuggingFace model must have 12 tensors"
        );

        // GGUF naming: 11 tensors (no lm_head due to weight tying)
        let gguf_data = build_executable_pygmy_apr_gguf_names();
        let gguf_model = AprV2Model::from_bytes(gguf_data).unwrap();
        assert_eq!(
            gguf_model.tensor_count(),
            11,
            "GGUF model must have 11 tensors (weight tying)"
        );

        // Embed-tied naming: 11 tensors (no lm_head)
        let tied_data = build_executable_pygmy_apr_embed_tied();
        let tied_model = AprV2Model::from_bytes(tied_data).unwrap();
        assert_eq!(
            tied_model.tensor_count(),
            11,
            "Embed-tied model must have 11 tensors"
        );
    }

    /// GH-194: All tensor naming conventions must produce valid logits
    #[test]
    fn test_gh194_all_naming_conventions_produce_valid_logits() {
        let models = vec![
            ("HuggingFace", build_executable_pygmy_apr()),
            ("GGUF", build_executable_pygmy_apr_gguf_names()),
            ("EmbedTied", build_executable_pygmy_apr_embed_tied()),
        ];

        for (name, data) in models {
            let model =
                AprV2Model::from_bytes(data).unwrap_or_else(|_| panic!("{name} model must load"));
            let logits = model
                .forward(&[1])
                .unwrap_or_else(|_| panic!("{name} forward must work"));

            // Logits must be valid floats (no NaN/Inf)
            assert!(
                logits.iter().all(|&x| x.is_finite()),
                "{name}: Logits must be finite (no NaN/Inf)"
            );

            // Logits must have correct size
            assert_eq!(logits.len(), 10, "{name}: Logits must have vocab_size=10");
        }
    }
}