skippy_protocol/
tokenizer.rs1use 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}