sokuto_runtime/tokenizer/
mod.rs

1//! # 高速トークナイザー統合モジュール
2//!
3//! Hugging Face `tokenizers` クレートをラップし、`tokenizer.json` のロード、
4//! 特殊トークンおよび `[OP]` マーカーの検証、State 優先トランケーション (State-Priority Truncation)、
5//! ならびに ONNX Runtime 入力用テンソルバッファの高速生成を行う。
6
7pub mod batch;
8pub mod prompt;
9
10pub use batch::BatchTokenizedQuestions;
11
12use std::path::Path;
13
14use sokuto_core::contract::model_spec::{MAX_SEQUENCE_LENGTH, TOKEN_OPTION_MARKER};
15use sokuto_core::schema::Question;
16
17use crate::error::{Result, RuntimeError};
18use prompt::format_prompt;
19
20/// 単一の質問に対するトークナイズ結果構造体。
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct TokenizedQuestion {
23    /// 入力トークン列テンソル用データ (`int64[sequence_length]`)。
24    pub input_ids: Vec<i64>,
25    /// アテンションマスクテンソル用データ (`int64[sequence_length]`)。
26    pub attention_mask: Vec<i64>,
27    /// 候補出現位置インデックステンソル用データ (`int64[num_options]`)。
28    pub op_indices: Vec<i64>,
29    /// プロンプト出現順に整列した候補識別子キー列。
30    pub option_keys: Vec<String>,
31}
32
33/// Jev システム用高速トークナイザー。
34#[derive(Debug, Clone)]
35pub struct JevTokenizer {
36    /// 内部の Hugging Face トークナイザーインスタンス。
37    inner: tokenizers::Tokenizer,
38    /// `[OP]` オプションマーカートークンの ID。
39    op_token_id: u32,
40    /// パディングトークンの ID。
41    pad_token_id: Option<u32>,
42    /// 先頭特殊トークン ID 列 (`[CLS]` または `<s>` 等)。
43    leading_special_tokens: Vec<u32>,
44    /// 末尾特殊トークン ID 列 (`[SEP]` または `</s>` 等)。
45    trailing_special_tokens: Vec<u32>,
46    /// 固定プレフィックス `"State: "` のキャッシュ済みトークン ID 列。
47    state_prefix_tokens: Vec<u32>,
48    /// 許容最大入力系列長。
49    max_sequence_length: usize,
50}
51
52impl JevTokenizer {
53    /// トークナイザー定義ファイル (`tokenizer.json`) からインスタンスを初期化する。
54    ///
55    /// 最大系列長にはデフォルトの `MAX_SEQUENCE_LENGTH` (8192) が適用される。
56    ///
57    /// # エラー
58    /// ファイルのロード失敗、`[OP]` トークン不在、または `[OP]` トークン分割テストに失敗した場合にエラーを返す。
59    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
60        Self::from_file_with_max_length(path, MAX_SEQUENCE_LENGTH)
61    }
62
63    /// 最大系列長を指定してトークナイザー定義ファイルからインスタンスを初期化する。
64    ///
65    /// # エラー
66    /// ファイルのロード失敗、`[OP]` トークン不在、または `[OP]` トークン分割テストに失敗した場合にエラーを返す。
67    pub fn from_file_with_max_length(
68        path: impl AsRef<Path>,
69        max_sequence_length: usize,
70    ) -> Result<Self> {
71        let path_ref = path.as_ref();
72        let inner = tokenizers::Tokenizer::from_file(path_ref).map_err(|e| {
73            RuntimeError::TokenizerLoadError(format!(
74                "ファイル '{}' のロードに失敗しました: {e}",
75                path_ref.display()
76            ))
77        })?;
78
79        Self::from_tokenizer_with_max_length(inner, max_sequence_length)
80    }
81
82    /// 既存の `tokenizers::Tokenizer` インスタンスから初期化する。
83    pub fn from_tokenizer_with_max_length(
84        inner: tokenizers::Tokenizer,
85        max_sequence_length: usize,
86    ) -> Result<Self> {
87        // [OP] マーカーの ID 取得
88        let op_token_id =
89            inner
90                .token_to_id(TOKEN_OPTION_MARKER)
91                .ok_or(RuntimeError::MissingSpecialToken {
92                    token: TOKEN_OPTION_MARKER,
93                })?;
94
95        // [OP] トークンが単一トークンとしてエンコードされるかの分割テスト検証
96        let op_encoding = inner
97            .encode(TOKEN_OPTION_MARKER, false)
98            .map_err(|e| RuntimeError::TokenizerEncodeError(e.to_string()))?;
99        let op_ids = op_encoding.get_ids();
100        if op_ids != [op_token_id] {
101            return Err(RuntimeError::InvalidOptionMarkerEncoding {
102                expected: op_token_id,
103                actual: op_ids.to_vec(),
104            });
105        }
106
107        // 先頭特殊トークン (CLS / BOS) の同定。
108        // ModernBERT-ja は `<cls>` を使用し、BERT系は `[CLS]`、RoBERTa/Gemma等は `<s>` / `[BOS]` を使用する。
109        let leading_special_tokens: Vec<u32> = inner
110            .token_to_id("<cls>")
111            .or_else(|| inner.token_to_id("[CLS]"))
112            .or_else(|| inner.token_to_id("<s>"))
113            .or_else(|| inner.token_to_id("[BOS]"))
114            .map(|id| vec![id])
115            .unwrap_or_default();
116
117        // 末尾特殊トークン (SEP / EOS) の同定。
118        // ModernBERT-ja は `<sep>` を使用し、BERT系は `[SEP]`、RoBERTa/Gemma等は `</s>` / `[EOS]` を使用する。
119        let trailing_special_tokens: Vec<u32> = inner
120            .token_to_id("<sep>")
121            .or_else(|| inner.token_to_id("[SEP]"))
122            .or_else(|| inner.token_to_id("</s>"))
123            .or_else(|| inner.token_to_id("[EOS]"))
124            .map(|id| vec![id])
125            .unwrap_or_default();
126
127        // パディングトークン (PAD) の同定。
128        let pad_token_id: Option<u32> = inner
129            .token_to_id("<pad>")
130            .or_else(|| inner.token_to_id("[PAD]"));
131
132        // 固定 Prefix ("State: ") のトークン列を事前エンコード・キャッシュ
133        let state_prefix_encoding = inner
134            .encode(prompt::PREFIX_STATE, false)
135            .map_err(|e| RuntimeError::TokenizerEncodeError(e.to_string()))?;
136        let state_prefix_tokens = state_prefix_encoding.get_ids().to_vec();
137
138        Ok(Self {
139            inner,
140            op_token_id,
141            pad_token_id,
142            leading_special_tokens,
143            trailing_special_tokens,
144            state_prefix_tokens,
145            max_sequence_length,
146        })
147    }
148
149    /// `[OP]` マーカーのトークン ID を取得する。
150    pub fn op_token_id(&self) -> u32 {
151        self.op_token_id
152    }
153
154    /// パディングトークン ID を取得する。
155    pub fn pad_token_id(&self) -> Option<u32> {
156        self.pad_token_id
157    }
158
159    /// 先頭特殊トークン ID 列を取得する。
160    pub fn leading_special_tokens(&self) -> &[u32] {
161        &self.leading_special_tokens
162    }
163
164    /// 末尾特殊トークン ID 列を取得する。
165    pub fn trailing_special_tokens(&self) -> &[u32] {
166        &self.trailing_special_tokens
167    }
168
169    /// 固定プレフィックス `"State: "` のキャッシュ済みトークン ID 列を取得する。
170    pub fn state_prefix_tokens(&self) -> &[u32] {
171        &self.state_prefix_tokens
172    }
173
174    /// 設定されている最大系列長を取得する。
175    pub fn max_sequence_length(&self) -> usize {
176        self.max_sequence_length
177    }
178
179    /// 内部の `tokenizers::Tokenizer` への参照を取得する。
180    pub fn inner(&self) -> &tokenizers::Tokenizer {
181        &self.inner
182    }
183
184    /// 文字列を特殊トークン付加なしでエンコードする。
185    fn encode_text(&self, text: &str) -> Result<Vec<u32>> {
186        let encoding = self
187            .inner
188            .encode(text, false)
189            .map_err(|e| RuntimeError::TokenizerEncodeError(e.to_string()))?;
190        Ok(encoding.get_ids().to_vec())
191    }
192
193    /// State 文字列を事前にトークナイズする (プレフィックス共有用)。
194    pub fn encode_state(&self, state: &str) -> Result<Vec<u32>> {
195        self.encode_text(state)
196    }
197
198    /// 単一の State と Question から ONNX 入力用バッファを生成する。
199    pub fn encode_question(&self, state: &str, question: &Question) -> Result<TokenizedQuestion> {
200        self.encode_question_with_max_len(state, question, self.max_sequence_length)
201    }
202
203    /// 系列長制約を明示指定して、単一の State と Question をエンコードする。
204    pub fn encode_question_with_max_len(
205        &self,
206        state: &str,
207        question: &Question,
208        max_length: usize,
209    ) -> Result<TokenizedQuestion> {
210        let state_ids = self.encode_state(state)?;
211        self.encode_with_pretokenized_state(&state_ids, question, max_length)
212    }
213
214    /// 事前エンコード済み State トークン列を用いてエンコードする (State 優先トランケーション適用)。
215    ///
216    /// # State 優先トランケーションの不変条件
217    /// 文末の Criteria (すべての `[OP]` マーカー群) および Instructions を絶対に保護し、
218    /// 系列長が上限を超える場合は State のみ末尾からスライスして切り詰める。
219    ///
220    /// # エラー
221    /// - 固定部 (特殊トークン + Prefix + Instructions + Criteria) が `max_length` 以上の場合は `PromptExceedsMaxLength`。
222    /// - 検出された `[OP]` マーカー数と候補数が不一致の場合は `OptionMarkerCountMismatch`。
223    pub fn encode_with_pretokenized_state(
224        &self,
225        state_ids: &[u32],
226        question: &Question,
227        max_length: usize,
228    ) -> Result<TokenizedQuestion> {
229        let formatted = format_prompt(question)?;
230
231        // キャッシュされた固定 Prefix ("State: ") のトークン列を再利用
232        let prefix_ids = &self.state_prefix_tokens;
233        let suffix_ids = self.encode_text(&formatted.suffix)?;
234
235        let fixed_len = self.leading_special_tokens.len()
236            + prefix_ids.len()
237            + suffix_ids.len()
238            + self.trailing_special_tokens.len();
239
240        if fixed_len >= max_length {
241            return Err(RuntimeError::PromptExceedsMaxLength {
242                fixed_len,
243                max_len: max_length,
244            });
245        }
246
247        // State を収容可能な残余スロット数を算出
248        let allowed_state_len = max_length - fixed_len;
249        let truncated_state_ids = if state_ids.len() > allowed_state_len {
250            &state_ids[..allowed_state_len]
251        } else {
252            state_ids
253        };
254
255        let total_len = fixed_len + truncated_state_ids.len();
256        let mut input_ids = Vec::with_capacity(total_len);
257
258        // トークン列の結合: [CLS] + Prefix + Truncated_State + Suffix + [SEP]
259        for &id in &self.leading_special_tokens {
260            input_ids.push(id as i64);
261        }
262        for &id in prefix_ids {
263            input_ids.push(id as i64);
264        }
265        for &id in truncated_state_ids {
266            input_ids.push(id as i64);
267        }
268        for &id in &suffix_ids {
269            input_ids.push(id as i64);
270        }
271        for &id in &self.trailing_special_tokens {
272            input_ids.push(id as i64);
273        }
274
275        let attention_mask = vec![1i64; total_len];
276
277        // [OP] マーカー出現インデックスの高速抽出
278        let op_token_id_i64 = self.op_token_id as i64;
279        let mut op_indices = Vec::with_capacity(formatted.option_keys.len());
280
281        for (idx, &token_id) in input_ids.iter().enumerate() {
282            if token_id == op_token_id_i64 {
283                op_indices.push(idx as i64);
284            }
285        }
286
287        // 候補数との完全一致検証
288        if op_indices.len() != formatted.option_keys.len() {
289            return Err(RuntimeError::OptionMarkerCountMismatch {
290                detected: op_indices.len(),
291                expected: formatted.option_keys.len(),
292            });
293        }
294
295        Ok(TokenizedQuestion {
296            input_ids,
297            attention_mask,
298            op_indices,
299            option_keys: formatted.option_keys,
300        })
301    }
302
303    /// 単一の質問定義から `QueryEvaluator` 入力用トークン列および `[OP]` マーカー位置を生成する。
304    ///
305    /// # 戻り値
306    /// - `TokenizedQuestion`: Query 部分の `input_ids`、`attention_mask`、`op_indices`、`option_keys`。
307    pub fn encode_query(&self, question: &Question) -> Result<TokenizedQuestion> {
308        self.encode_query_with_max_len(question, self.max_sequence_length)
309    }
310
311    /// 系列長制約を指定して、単一の質問定義から `QueryEvaluator` 入力用トークン列を生成する。
312    pub fn encode_query_with_max_len(
313        &self,
314        question: &Question,
315        max_length: usize,
316    ) -> Result<TokenizedQuestion> {
317        let formatted = format_prompt(question)?;
318        let suffix_ids = self.encode_text(&formatted.suffix)?;
319
320        let mut query_tokens = Vec::with_capacity(
321            self.leading_special_tokens.len()
322                + suffix_ids.len()
323                + self.trailing_special_tokens.len(),
324        );
325        query_tokens.extend_from_slice(&self.leading_special_tokens);
326        query_tokens.extend_from_slice(&suffix_ids);
327        query_tokens.extend_from_slice(&self.trailing_special_tokens);
328
329        if query_tokens.len() > max_length {
330            return Err(RuntimeError::PromptExceedsMaxLength {
331                fixed_len: query_tokens.len(),
332                max_len: max_length,
333            });
334        }
335
336        let seq_len = query_tokens.len();
337        let input_ids: Vec<i64> = query_tokens.iter().map(|&id| id as i64).collect();
338        let attention_mask = vec![1i64; seq_len];
339
340        let mut op_indices = Vec::with_capacity(formatted.option_keys.len());
341        let op_token_id_i64 = self.op_token_id as i64;
342        for (idx, &token_id) in input_ids.iter().enumerate() {
343            if token_id == op_token_id_i64 {
344                op_indices.push(idx as i64);
345            }
346        }
347
348        if op_indices.len() != formatted.option_keys.len() {
349            return Err(RuntimeError::OptionMarkerCountMismatch {
350                detected: op_indices.len(),
351                expected: formatted.option_keys.len(),
352            });
353        }
354
355        Ok(TokenizedQuestion {
356            input_ids,
357            attention_mask,
358            op_indices,
359            option_keys: formatted.option_keys,
360        })
361    }
362
363    /// `State` 文字列を `StateEncoder` 入力用トークン列 (`input_ids`, `attention_mask`) としてエンコードする。
364    pub fn encode_state_tensor(&self, state: &str) -> Result<(Vec<i64>, Vec<i64>)> {
365        let text_tokens = self.encode_state(state)?;
366        let mut tokens = Vec::with_capacity(
367            self.leading_special_tokens.len()
368                + self.state_prefix_tokens.len()
369                + text_tokens.len()
370                + self.trailing_special_tokens.len(),
371        );
372        tokens.extend_from_slice(&self.leading_special_tokens);
373        tokens.extend_from_slice(&self.state_prefix_tokens);
374        tokens.extend_from_slice(&text_tokens);
375        tokens.extend_from_slice(&self.trailing_special_tokens);
376
377        if tokens.len() > self.max_sequence_length {
378            tokens.truncate(self.max_sequence_length);
379        }
380
381        let seq_len = tokens.len();
382        let input_ids: Vec<i64> = tokens.iter().map(|&t| t as i64).collect();
383        let attention_mask = vec![1i64; seq_len];
384        Ok((input_ids, attention_mask))
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    use indexmap::IndexMap;
392    use std::path::PathBuf;
393
394    /// ワークスペース内の default tokenizer.json のパスを取得するヘルパー。
395    fn default_tokenizer_path() -> Option<PathBuf> {
396        let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
397        let path = manifest_dir
398            .parent()
399            .unwrap()
400            .parent()
401            .unwrap()
402            .join("models")
403            .join("default")
404            .join("tokenizer.json");
405
406        if !path.exists() {
407            if std::env::var("CI").is_ok() {
408                panic!(
409                    "CI 環境で必須トークナイザーファイルが見つかりません: {:?}",
410                    path
411                );
412            }
413            eprintln!("スキップ: tokenizer.json が見つかりません: {:?}", path);
414            return None;
415        }
416
417        Some(path)
418    }
419
420    #[test]
421    fn test_load_default_tokenizer() {
422        let Some(path) = default_tokenizer_path() else {
423            return;
424        };
425
426        let tokenizer = JevTokenizer::from_file(&path).unwrap();
427        assert_eq!(tokenizer.max_sequence_length(), 8192);
428        assert!(!tokenizer.leading_special_tokens().is_empty());
429        assert!(!tokenizer.trailing_special_tokens().is_empty());
430        assert!(!tokenizer.state_prefix_tokens().is_empty());
431        assert!(tokenizer.pad_token_id().is_some());
432
433        // [OP] トークン ID が正しく解決されていること
434        let op_id = tokenizer.op_token_id();
435        assert!(op_id > 0);
436    }
437
438    #[test]
439    fn test_tokenize_choice_question() {
440        let Some(path) = default_tokenizer_path() else {
441            return;
442        };
443
444        let tokenizer = JevTokenizer::from_file(&path).unwrap();
445
446        let mut criteria = IndexMap::new();
447        criteria.insert(
448            "card_arrival".to_string(),
449            "Card delivery status".to_string(),
450        );
451        criteria.insert("lost_card".to_string(), "Reporting lost card".to_string());
452        criteria.insert("pin_reset".to_string(), "Resetting PIN code".to_string());
453
454        let question = Question::new_choice("意図を分類せよ。", criteria);
455        let state = "I lost my card yesterday.";
456
457        let tokenized = tokenizer.encode_question(state, &question).unwrap();
458
459        assert_eq!(tokenized.input_ids.len(), tokenized.attention_mask.len());
460        assert_eq!(
461            tokenized.option_keys,
462            vec!["card_arrival", "lost_card", "pin_reset"]
463        );
464        assert_eq!(tokenized.op_indices.len(), 3);
465
466        // 各 [OP] 位置のトークン ID が op_token_id と一致すること
467        let op_id_i64 = tokenizer.op_token_id() as i64;
468        for &idx in &tokenized.op_indices {
469            assert_eq!(tokenized.input_ids[idx as usize], op_id_i64);
470        }
471    }
472
473    #[test]
474    fn test_tokenize_noul_question() {
475        let Some(path) = default_tokenizer_path() else {
476            return;
477        };
478
479        let tokenizer = JevTokenizer::from_file(&path).unwrap();
480        let question = Question::new_noul("言明「空は青い」の真偽を判定せよ。");
481        let state = "The sky is blue.";
482
483        let tokenized = tokenizer.encode_question(state, &question).unwrap();
484
485        assert_eq!(tokenized.option_keys, vec!["true", "false"]);
486        assert_eq!(tokenized.op_indices.len(), 2);
487
488        let op_id_i64 = tokenizer.op_token_id() as i64;
489        for &idx in &tokenized.op_indices {
490            assert_eq!(tokenized.input_ids[idx as usize], op_id_i64);
491        }
492    }
493
494    #[test]
495    fn test_state_priority_truncation() {
496        let Some(path) = default_tokenizer_path() else {
497            return;
498        };
499
500        let tokenizer = JevTokenizer::from_file(&path).unwrap();
501
502        let mut criteria = IndexMap::new();
503        criteria.insert("opt_a".to_string(), "First option description".to_string());
504        criteria.insert("opt_b".to_string(), "Second option description".to_string());
505        criteria.insert("opt_c".to_string(), "Third option description".to_string());
506        criteria.insert("opt_d".to_string(), "Fourth option description".to_string());
507
508        let question = Question::new_choice("分類せよ。", criteria);
509
510        // 非常に長い State テキスト
511        let long_state =
512            "This is an extremely long transaction log with hundreds of details. ".repeat(50);
513
514        // 系列長上限を狭い 128 に設定
515        let max_len = 128;
516        let tokenized = tokenizer
517            .encode_question_with_max_len(&long_state, &question, max_len)
518            .unwrap();
519
520        // 系列長が上限以下に収まっていること
521        assert!(tokenized.input_ids.len() <= max_len);
522
523        // 【重要】State が切り詰められても、4候補すべての [OP] マーカーが完全に保持されていること
524        assert_eq!(tokenized.op_indices.len(), 4);
525        assert_eq!(tokenized.option_keys.len(), 4);
526
527        let op_id_i64 = tokenizer.op_token_id() as i64;
528        for &idx in &tokenized.op_indices {
529            assert_eq!(tokenized.input_ids[idx as usize], op_id_i64);
530        }
531    }
532
533    #[test]
534    fn test_prompt_exceeds_max_length_error() {
535        let Some(path) = default_tokenizer_path() else {
536            return;
537        };
538
539        let tokenizer = JevTokenizer::from_file(&path).unwrap();
540
541        let mut criteria = IndexMap::new();
542        for i in 0..10 {
543            criteria.insert(format!("opt_{i}"), format!("詳細な候補説明その{i}"));
544        }
545        let question = Question::new_choice("分類せよ。", criteria);
546
547        // 固定部より小さい極端な max_len (例えば 10) を指定
548        let res = tokenizer.encode_question_with_max_len("短い状態", &question, 10);
549        assert!(matches!(
550            res,
551            Err(RuntimeError::PromptExceedsMaxLength { .. })
552        ));
553    }
554}
sokuto_runtime/tokenizer/mod.rs

sokuto_runtime/tokenizer/
mod.rs
