Skip to main content

skippy_protocol/
tokenizer.rs

1use serde::{Deserialize, Serialize};
2
3pub const MAX_TOKENIZE_INPUT_BYTES: usize = 1_048_576;
4
5#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
6pub struct TokenizerIdentity {
7    pub model_id: String,
8    pub source_model_sha256: String,
9    pub tokenizer_id: String,
10}
11
12#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
13pub struct TokenizeRequest {
14    pub expected_identity: TokenizerIdentity,
15    pub text: String,
16    #[serde(default)]
17    pub add_special: bool,
18    #[serde(default)]
19    pub include_token_pieces: bool,
20}
21
22#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
23pub struct TokenizeResponse {
24    pub identity: TokenizerIdentity,
25    pub token_ids: Vec<i32>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub token_pieces: Option<Vec<Vec<u8>>>,
28}
29
30#[cfg(test)]
31mod tests {
32    use super::*;
33
34    #[test]
35    fn tokenizer_wire_contract_is_exact() {
36        let identity = TokenizerIdentity {
37            model_id: "model".to_owned(),
38            source_model_sha256: "a".repeat(64),
39            tokenizer_id: format!("gguf-source-sha256:{}", "a".repeat(64)),
40        };
41        let request: TokenizeRequest = serde_json::from_value(serde_json::json!({
42            "expected_identity": identity,
43            "text": "hello",
44        }))
45        .unwrap();
46        assert!(!request.add_special);
47        assert!(!request.include_token_pieces);
48
49        let response = TokenizeResponse {
50            identity: request.expected_identity,
51            token_ids: vec![1, 2],
52            token_pieces: None,
53        };
54        assert_eq!(
55            serde_json::to_value(response).unwrap(),
56            serde_json::json!({
57                "identity": {
58                    "model_id": "model",
59                    "source_model_sha256": "a".repeat(64),
60                    "tokenizer_id": format!("gguf-source-sha256:{}", "a".repeat(64)),
61                },
62                "token_ids": [1, 2],
63            })
64        );
65        assert_eq!(MAX_TOKENIZE_INPUT_BYTES, 1_048_576);
66    }
67}