sokuto_runtime/engine/
batch.rs

1//! # 複数質問一括バッチ推論モジュール
2//!
3//! トークンレベルのプレフィックス共有と動的バッチ Collate により、
4//! 単一フォワードパスでの並列推論および決定プリミティブ解決を実現する。
5
6use indexmap::IndexMap;
7use sokuto_core::contract::CalibrationConfig;
8use sokuto_core::decision::evaluate_question_with_buf;
9use sokuto_core::gating::GatingConfig;
10use sokuto_core::schema::{Answer, Question, QuestionType};
11
12use crate::engine::coarse::{
13    CoarseScorer, CoarseToFineConfig, LexicalCoarseScorer, filter_top_candidates,
14    reconstruct_probabilities,
15};
16use crate::engine::gating::{apply_gating_to_answer_with_energy, calculate_choice_energy};
17use crate::engine::session::InferenceEngine;
18use crate::error::Result;
19use crate::hierarchical::{
20    CoarseToFineRouter, HierarchicalMapping, HierarchicalResult, HierarchicalRouterConfig,
21};
22use crate::tokenizer::{BatchTokenizedQuestions, JevTokenizer};
23
24/// サーバー層との連携において推奨されるデフォルトの最大バッチチャンクサイズ。
25///
26/// 質問数がこの値を超える場合、OOM (Out Of Memory) を防止するため
27/// マイクロバッチに分割して実行することが推奨される。
28pub const DEFAULT_MAX_BATCH_CHUNK_SIZE: usize = 16;
29
30/// バッチ推論および決定数理解決用のスクラッチパッドバッファ。
31///
32/// スレッドローカルまたはリクエストコンテキスト内で再利用することで、
33/// 推論ごとのメモリアロケーションを根絶する。
34#[derive(Debug, Default, Clone)]
35pub struct BatchScratchpad {
36    /// 平坦化生ロジット書き込み用バッファ (`[batch_size * num_options]`)。
37    pub flat_logits: Vec<f64>,
38    /// ソフトマックス確率分布算出用作業バッファ (`[max_options_per_q]`)。
39    pub probs_buf: Vec<f64>,
40}
41
42impl BatchScratchpad {
43    /// 指定されたキャパシティでスクラッチパッドを初期化する。
44    pub fn with_capacity(total_options: usize, max_options_per_q: usize) -> Self {
45        Self {
46            flat_logits: Vec::with_capacity(total_options),
47            probs_buf: Vec::with_capacity(max_options_per_q),
48        }
49    }
50
51    /// 要求されるサイズに合わせてバッファ領域を確保・リサイズする。
52    pub fn prepare(&mut self, total_options: usize, max_options_per_q: usize) {
53        if self.flat_logits.len() < total_options {
54            self.flat_logits.resize(total_options, 0.0);
55        }
56        if self.probs_buf.len() < max_options_per_q {
57            self.probs_buf.resize(max_options_per_q, 0.0);
58        }
59    }
60}
61
62impl InferenceEngine {
63    /// 事前パッキング済みバッチトークン列を受け取り、単一フォワードパスで各質問のロジットを算出する。
64    ///
65    /// 出力されたテンソルから各質問の真の候補数 $K_i$ 個のロジットのみをスライス抽出することで、
66    /// パディングスロット由来のダミーロジットを完全に除外して返却する。
67    ///
68    /// # 引数
69    /// - `batch`: パディングおよび平坦化済みのバッチトークナイズ結果。
70    ///
71    /// # 戻り値
72    /// 質問キーと有効ロジット列 (`Vec<f64>`) の順序付きマップ。
73    ///
74    /// # エラー
75    /// ONNX Runtime 推論処理またはテンソル検証に失敗した場合にエラーを返す。
76    pub fn forward_batch_tokenized(
77        &self,
78        batch: &BatchTokenizedQuestions,
79    ) -> Result<IndexMap<String, Vec<f64>>> {
80        let total_logits = batch.dims.batch_size * batch.dims.num_options;
81        let mut flat_logits = vec![0.0f64; total_logits];
82
83        self.forward_batch_tokenized_into(batch, &mut flat_logits)?;
84
85        let mut results = IndexMap::with_capacity(batch.dims.batch_size);
86        let num_options = batch.dims.num_options;
87
88        for (i, q_key) in batch.question_keys.iter().enumerate() {
89            let k_i = batch.candidate_counts[i];
90            let start_idx = i * num_options;
91            let valid_logits = flat_logits[start_idx..start_idx + k_i].to_vec();
92            results.insert(q_key.clone(), valid_logits);
93        }
94
95        Ok(results)
96    }
97
98    /// 事前パッキング済みバッチトークン列を受け取り、推論結果を指定されたスライスに書き込む。
99    ///
100    /// 借用ビュー (`TensorRef`) によりゼロコピーで ONNX Runtime に渡される。
101    ///
102    /// # 引数
103    /// - `batch`: パディングおよび平坦化済みのバッチトークナイズ結果。
104    /// - `out_logits`: 平坦化生ロジットの書き込み先バッファ (`[batch_size * num_options]`)。
105    pub fn forward_batch_tokenized_into(
106        &self,
107        batch: &BatchTokenizedQuestions,
108        out_logits: &mut [f64],
109    ) -> Result<()> {
110        self.forward_raw_into(
111            &batch.input_ids,
112            &batch.attention_mask,
113            &batch.op_indices,
114            batch.dims,
115            out_logits,
116        )
117    }
118
119    /// 共通 State と複数質問群を受け取り、トークンプレフィックス共有バッチ推論を実行する。
120    ///
121    /// State のトークナイズを 1 度だけ実行し、各質問を単一のバッチ次元に集約して
122    /// 1 回のフォワードパスで全質問の有効ロジットを返却する。
123    ///
124    /// # 引数
125    /// - `tokenizer`: Jev 高速トークナイザー。
126    /// - `state`: 共通文脈テキスト。
127    /// - `questions`: 質問識別子と質問定義のマップ。
128    ///
129    /// # 戻り値
130    /// 質問キーと有効ロジット列 (`Vec<f64>`) の順序付きマップ。
131    pub fn forward_batch_questions(
132        &self,
133        tokenizer: &JevTokenizer,
134        state: &str,
135        questions: &IndexMap<String, Question>,
136    ) -> Result<IndexMap<String, Vec<f64>>> {
137        let batch = tokenizer.encode_batch_questions(state, questions)?;
138        self.forward_batch_tokenized(&batch)
139    }
140
141    /// 共通 State、複数質問群、および較正設定を受け取り、単一フォワードパスで判定結果を一括導出する。
142    ///
143    /// # 処理フロー
144    /// 1. `tokenizer.encode_batch_questions` によるプレフィックス共有バッチ化。
145    /// 2. `forward_batch_tokenized_into` による単一フォワードパス実行。
146    /// 3. 各質問の真の候補数 $K_i$ によるスライス抽出とダミーロジット破棄。
147    /// 4. `sokuto_core::decision::evaluate_question_with_buf` による決定論的プリミティブ解決。
148    /// 5. 較正設定の閾値に基づき、透過的に確信度ゲーティングを適用して返却。
149    ///
150    /// # 引数
151    /// - `tokenizer`: Jev 高速トークナイザー。
152    /// - `state`: 共通文脈テキスト。
153    /// - `questions`: 質問識別子と質問定義のマップ。
154    /// - `calib_config`: 較正温度設定。
155    ///
156    /// # 戻り値
157    /// 質問キーと決定結果 (`Answer`) の順序付きマップ。
158    pub fn evaluate_batch_questions(
159        &self,
160        tokenizer: &JevTokenizer,
161        state: &str,
162        questions: &IndexMap<String, Question>,
163        calib_config: &CalibrationConfig,
164    ) -> Result<IndexMap<String, Answer>> {
165        self.evaluate_batch_questions_with_gating(tokenizer, state, questions, calib_config, None)
166    }
167
168    /// ゲーティング設定を明示指定して、複数質問一括バッチ推論を実行する。
169    ///
170    /// # 引数
171    /// - `tokenizer`: Jev 高速トークナイザー。
172    /// - `state`: 共通文脈テキスト。
173    /// - `questions`: 質問識別子と質問定義のマップ。
174    /// - `calib_config`: 較正温度設定。
175    /// - `gating_config`: 明示的なゲーティング設定 (任意、未指定時は較正設定から導出)。
176    pub fn evaluate_batch_questions_with_gating(
177        &self,
178        tokenizer: &JevTokenizer,
179        state: &str,
180        questions: &IndexMap<String, Question>,
181        calib_config: &CalibrationConfig,
182        gating_config: Option<&GatingConfig>,
183    ) -> Result<IndexMap<String, Answer>> {
184        let batch = tokenizer.encode_batch_questions(state, questions)?;
185        let total_logits = batch.dims.batch_size * batch.dims.num_options;
186        let mut flat_logits = vec![0.0f64; total_logits];
187
188        self.forward_batch_tokenized_into(&batch, &mut flat_logits)?;
189
190        let mut answers = IndexMap::with_capacity(batch.dims.batch_size);
191        let mut probs_buf = vec![0.0f64; batch.dims.num_options];
192        let num_options = batch.dims.num_options;
193
194        for (i, q_key) in batch.question_keys.iter().enumerate() {
195            let k_i = batch.candidate_counts[i];
196            let start_idx = i * num_options;
197            let valid_logits = &flat_logits[start_idx..start_idx + k_i];
198            let question = questions.get(q_key).expect("キーの一致が保証されている。");
199
200            let mut answer = evaluate_question_with_buf(
201                question,
202                valid_logits,
203                calib_config,
204                &mut probs_buf[..k_i],
205            )?;
206            let energy = calculate_choice_energy(question, valid_logits, calib_config);
207            apply_gating_to_answer_with_energy(
208                &mut answer,
209                energy,
210                Some(state),
211                Some(q_key),
212                Some(question),
213                gating_config,
214                calib_config,
215            );
216            answers.insert(q_key.clone(), answer);
217        }
218
219        Ok(answers)
220    }
221
222    /// スクラッチパッドバッファを再利用し、共通 State と複数質問群の判定結果を一括導出する。
223    ///
224    /// アロケーションを最小化したい高頻度 API サーバー環境向けの高速推論インターフェース。
225    ///
226    /// # 引数
227    /// - `tokenizer`: Jev 高速トークナイザー。
228    /// - `state`: 共通文脈テキスト。
229    /// - `questions`: 質問識別子と質問定義のマップ。
230    /// - `calib_config`: 較正温度設定。
231    /// - `scratchpad`: 再利用可能な作業用スクラッチパッドバッファ。
232    pub fn evaluate_batch_questions_with_scratchpad(
233        &self,
234        tokenizer: &JevTokenizer,
235        state: &str,
236        questions: &IndexMap<String, Question>,
237        calib_config: &CalibrationConfig,
238        scratchpad: &mut BatchScratchpad,
239    ) -> Result<IndexMap<String, Answer>> {
240        self.evaluate_batch_questions_with_scratchpad_and_gating(
241            tokenizer,
242            state,
243            questions,
244            calib_config,
245            None,
246            scratchpad,
247        )
248    }
249
250    /// スクラッチパッドバッファとゲーティング設定を指定してバッチ推論を実行する。
251    ///
252    /// # 引数
253    /// - `tokenizer`: Jev 高速トークナイザー。
254    /// - `state`: 共通文脈テキスト。
255    /// - `questions`: 質問識別子と質問定義のマップ。
256    /// - `calib_config`: 較正温度設定。
257    /// - `gating_config`: 明示的なゲーティング設定 (任意)。
258    /// - `scratchpad`: 再利用可能な作業用スクラッチパッドバッファ。
259    pub fn evaluate_batch_questions_with_scratchpad_and_gating(
260        &self,
261        tokenizer: &JevTokenizer,
262        state: &str,
263        questions: &IndexMap<String, Question>,
264        calib_config: &CalibrationConfig,
265        gating_config: Option<&GatingConfig>,
266        scratchpad: &mut BatchScratchpad,
267    ) -> Result<IndexMap<String, Answer>> {
268        let batch = tokenizer.encode_batch_questions(state, questions)?;
269        let total_logits = batch.dims.batch_size * batch.dims.num_options;
270        let num_options = batch.dims.num_options;
271
272        scratchpad.prepare(total_logits, num_options);
273
274        self.forward_batch_tokenized_into(&batch, &mut scratchpad.flat_logits[..total_logits])?;
275
276        let mut answers = IndexMap::with_capacity(batch.dims.batch_size);
277
278        for (i, q_key) in batch.question_keys.iter().enumerate() {
279            let k_i = batch.candidate_counts[i];
280            let start_idx = i * num_options;
281            let valid_logits = &scratchpad.flat_logits[start_idx..start_idx + k_i];
282            let question = questions.get(q_key).expect("キーの一致が保証されている。");
283
284            let mut answer = evaluate_question_with_buf(
285                question,
286                valid_logits,
287                calib_config,
288                &mut scratchpad.probs_buf[..k_i],
289            )?;
290            let energy = calculate_choice_energy(question, valid_logits, calib_config);
291            apply_gating_to_answer_with_energy(
292                &mut answer,
293                energy,
294                Some(state),
295                Some(q_key),
296                Some(question),
297                gating_config,
298                calib_config,
299            );
300            answers.insert(q_key.clone(), answer);
301        }
302
303        Ok(answers)
304    }
305
306    /// 共通 State と複数質問群を指定されたチャンクサイズで分割実行し、有効ロジットを返却する。
307    ///
308    /// 質問数が極端に多いリクエストにおいて、メモリ使用量 ($N \times L_{\max}$) の急増を抑制し
309    /// OOM (Out Of Memory) を防止するマイクロバッチ推論インターフェース。
310    ///
311    /// # 引数
312    /// - `tokenizer`: Jev 高速トークナイザー。
313    /// - `state`: 共通文脈テキスト。
314    /// - `questions`: 質問識別子と質問定義のマップ。
315    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
316    pub fn forward_batch_questions_chunked(
317        &self,
318        tokenizer: &JevTokenizer,
319        state: &str,
320        questions: &IndexMap<String, Question>,
321        chunk_size: usize,
322    ) -> Result<IndexMap<String, Vec<f64>>> {
323        let chunk_size = chunk_size.max(1);
324        if questions.len() <= chunk_size {
325            return self.forward_batch_questions(tokenizer, state, questions);
326        }
327
328        let state_ids = tokenizer.encode_state(state)?;
329        let mut results = IndexMap::with_capacity(questions.len());
330
331        for chunk_slice in questions.iter().collect::<Vec<_>>().chunks(chunk_size) {
332            let mut sub_questions = IndexMap::with_capacity(chunk_slice.len());
333            for (k, v) in chunk_slice {
334                sub_questions.insert((*k).clone(), (*v).clone());
335            }
336
337            let batch = tokenizer.encode_batch_with_pretokenized_state(
338                &state_ids,
339                &sub_questions,
340                tokenizer.max_sequence_length(),
341            )?;
342            let sub_results = self.forward_batch_tokenized(&batch)?;
343            results.extend(sub_results);
344        }
345
346        Ok(results)
347    }
348
349    /// 共通 State と複数質問群を指定されたチャンクサイズで分割実行し、判定結果を一括導出する。
350    ///
351    /// サーバー層 (sokuto-server) のガードレール等と連携し、大規模質問リクエストを
352    /// 安全な固定バッチサイズで順次推論・解決する。較正設定の閾値に基づき、各質問に透過的にゲーティングを適用する。
353    ///
354    /// # 引数
355    /// - `tokenizer`: Jev 高速トークナイザー。
356    /// - `state`: 共通文脈テキスト。
357    /// - `questions`: 質問識別子と質問定義のマップ。
358    /// - `calib_config`: 較正温度設定。
359    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
360    pub fn evaluate_batch_questions_chunked(
361        &self,
362        tokenizer: &JevTokenizer,
363        state: &str,
364        questions: &IndexMap<String, Question>,
365        calib_config: &CalibrationConfig,
366        chunk_size: usize,
367    ) -> Result<IndexMap<String, Answer>> {
368        self.evaluate_batch_questions_chunked_with_gating(
369            tokenizer,
370            state,
371            questions,
372            calib_config,
373            chunk_size,
374            None,
375        )
376    }
377
378    /// チャンク分割およびゲーティング設定を指定してバッチ推論を実行する。
379    ///
380    /// # 引数
381    /// - `tokenizer`: Jev 高速トークナイザー。
382    /// - `state`: 共通文脈テキスト。
383    /// - `questions`: 質問識別子と質問定義のマップ。
384    /// - `calib_config`: 較正温度設定。
385    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
386    /// - `gating_config`: 明示的なゲーティング設定 (任意)。
387    pub fn evaluate_batch_questions_chunked_with_gating(
388        &self,
389        tokenizer: &JevTokenizer,
390        state: &str,
391        questions: &IndexMap<String, Question>,
392        calib_config: &CalibrationConfig,
393        chunk_size: usize,
394        gating_config: Option<&GatingConfig>,
395    ) -> Result<IndexMap<String, Answer>> {
396        let chunk_size = chunk_size.max(1);
397        if questions.len() <= chunk_size {
398            return self.evaluate_batch_questions_with_gating(
399                tokenizer,
400                state,
401                questions,
402                calib_config,
403                gating_config,
404            );
405        }
406
407        let state_ids = tokenizer.encode_state(state)?;
408        let mut answers = IndexMap::with_capacity(questions.len());
409
410        for chunk_slice in questions.iter().collect::<Vec<_>>().chunks(chunk_size) {
411            let mut sub_questions = IndexMap::with_capacity(chunk_slice.len());
412            for (k, v) in chunk_slice {
413                sub_questions.insert((*k).clone(), (*v).clone());
414            }
415
416            let batch = tokenizer.encode_batch_with_pretokenized_state(
417                &state_ids,
418                &sub_questions,
419                tokenizer.max_sequence_length(),
420            )?;
421
422            let total_logits = batch.dims.batch_size * batch.dims.num_options;
423            let mut flat_logits = vec![0.0f64; total_logits];
424            self.forward_batch_tokenized_into(&batch, &mut flat_logits)?;
425
426            let mut probs_buf = vec![0.0f64; batch.dims.num_options];
427            let num_options = batch.dims.num_options;
428
429            for (i, q_key) in batch.question_keys.iter().enumerate() {
430                let k_i = batch.candidate_counts[i];
431                let start_idx = i * num_options;
432                let valid_logits = &flat_logits[start_idx..start_idx + k_i];
433                let question = sub_questions
434                    .get(q_key)
435                    .expect("キーの一致が保証されている。");
436
437                let mut answer = evaluate_question_with_buf(
438                    question,
439                    valid_logits,
440                    calib_config,
441                    &mut probs_buf[..k_i],
442                )?;
443                let energy = calculate_choice_energy(question, valid_logits, calib_config);
444                apply_gating_to_answer_with_energy(
445                    &mut answer,
446                    energy,
447                    Some(state),
448                    Some(q_key),
449                    Some(question),
450                    gating_config,
451                    calib_config,
452                );
453                answers.insert(q_key.clone(), answer);
454            }
455        }
456
457        Ok(answers)
458    }
459
460    /// 任意の類似度スコアラーを指定して、粗密 2 段階探索を適用した複数質問一括バッチ推論を実行する。
461    ///
462    /// # 処理フロー
463    /// 1. 各質問を走査し、`Choice` かつ候補数 $K > \text{threshold}$ の質問を個別に Top-M 候補にスクリーニング。
464    /// 2. 縮小質問群(および閾値以下の元質問群)をプレフィックス共有バッチ化。
465    /// 3. バッチ全体の最大候補数 $K_{\max}$ が圧縮された状態で単一フォワードパスを実行。
466    /// 4. 各質問の決定プリミティブを解決し、スクリーニングされた質問については除外候補を確率 0.0 として全候補確率マップを再構成。
467    ///
468    /// # 引数
469    /// - `tokenizer`: Jev 高速トークナイザー。
470    /// - `state`: 共通文脈テキスト。
471    /// - `questions`: 質問マップ。
472    /// - `calib_config`: 較正温度設定。
473    /// - `coarse_config`: 粗密探索設定。
474    /// - `scorer`: 類似度スコアラー実装。
475    pub fn evaluate_batch_questions_coarse_to_fine_with_scorer<S: CoarseScorer>(
476        &self,
477        tokenizer: &JevTokenizer,
478        state: &str,
479        questions: &IndexMap<String, Question>,
480        calib_config: &CalibrationConfig,
481        coarse_config: &CoarseToFineConfig,
482        scorer: &S,
483    ) -> Result<IndexMap<String, Answer>> {
484        coarse_config.validate()?;
485
486        if questions.is_empty() {
487            return self.evaluate_batch_questions(tokenizer, state, questions, calib_config);
488        }
489
490        // 各質問の前処理と縮小
491        let mut processed_questions = IndexMap::with_capacity(questions.len());
492        let mut original_keys_map: IndexMap<String, Vec<String>> = IndexMap::new();
493
494        for (q_key, question) in questions {
495            let should_trigger = question.question_type == QuestionType::Choice
496                && question
497                    .criteria
498                    .as_ref()
499                    .and_then(|c| c.as_map())
500                    .map(|map| coarse_config.should_trigger(map.len()))
501                    .unwrap_or(false);
502
503            if should_trigger {
504                let map = question
505                    .criteria
506                    .as_ref()
507                    .and_then(|c| c.as_map())
508                    .expect("should_trigger により Map が保証されている。");
509
510                let original_keys: Vec<String> = map.keys().cloned().collect();
511                let query = format!("State: {}\nInstructions: {}", state, question.instructions);
512
513                let filtered = filter_top_candidates(&query, map, coarse_config, scorer)?;
514                let sub_question =
515                    Question::new_choice(question.instructions.clone(), filtered.selected);
516
517                processed_questions.insert(q_key.clone(), sub_question);
518                original_keys_map.insert(q_key.clone(), original_keys);
519            } else {
520                processed_questions.insert(q_key.clone(), question.clone());
521            }
522        }
523
524        // バッチ推論を実行 (最大候補数 K_max が縮小された状態で単一フォワードパス)
525        let mut answers =
526            self.evaluate_batch_questions(tokenizer, state, &processed_questions, calib_config)?;
527
528        // スクリーニングされた質問の確率マップを全候補空間に復元し、ゲーティングを再適用
529        for (q_key, original_keys) in &original_keys_map {
530            if let Some(answer) = answers.get_mut(q_key)
531                && let Some(ref sub_probs) = answer.probabilities
532            {
533                let full_probs = reconstruct_probabilities(original_keys, sub_probs);
534                answer.probabilities = Some(full_probs);
535                let orig_q = questions.get(q_key);
536                let energy = answer.gating.as_ref().and_then(|g| g.energy);
537                apply_gating_to_answer_with_energy(
538                    answer,
539                    energy,
540                    Some(state),
541                    Some(q_key),
542                    orig_q,
543                    None,
544                    calib_config,
545                );
546            }
547        }
548
549        Ok(answers)
550    }
551
552    /// 既定の語彙スコアラー (`LexicalCoarseScorer`) を用いて、粗密 2 段階探索を適用した複数質問一括バッチ推論を実行する。
553    ///
554    /// # 引数
555    /// - `tokenizer`: Jev 高速トークナイザー。
556    /// - `state`: 共通文脈テキスト。
557    /// - `questions`: 質問マップ。
558    /// - `calib_config`: 較正温度設定。
559    /// - `coarse_config`: 粗密探索設定。
560    pub fn evaluate_batch_questions_coarse_to_fine(
561        &self,
562        tokenizer: &JevTokenizer,
563        state: &str,
564        questions: &IndexMap<String, Question>,
565        calib_config: &CalibrationConfig,
566        coarse_config: &CoarseToFineConfig,
567    ) -> Result<IndexMap<String, Answer>> {
568        let scorer = LexicalCoarseScorer::new(tokenizer);
569        self.evaluate_batch_questions_coarse_to_fine_with_scorer(
570            tokenizer,
571            state,
572            questions,
573            calib_config,
574            coarse_config,
575            &scorer,
576        )
577    }
578
579    /// 任意の類似度スコアラーを指定して、粗密 2 段階探索とマイクロバッチチャンキングを統合した安全推論を実行する。
580    ///
581    /// 候補数が膨大な質問を事前に Top-$M$ 件へスクリーニングした上で、
582    /// 共通 State を 1 度だけエンコードし、質問群を指定されたチャンクサイズごとに順次フォワードパスへ供給する。
583    ///
584    /// # 引数
585    /// - `tokenizer`: Jev 高速トークナイザー。
586    /// - `state`: 共通文脈テキスト。
587    /// - `questions`: 質問マップ。
588    /// - `calib_config`: 較正温度設定。
589    /// - `coarse_config`: 粗密探索設定。
590    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
591    /// - `scorer`: 類似度スコアラー実装。
592    #[allow(clippy::too_many_arguments)]
593    pub fn evaluate_batch_questions_coarse_to_fine_chunked_with_scorer<S: CoarseScorer>(
594        &self,
595        tokenizer: &JevTokenizer,
596        state: &str,
597        questions: &IndexMap<String, Question>,
598        calib_config: &CalibrationConfig,
599        coarse_config: &CoarseToFineConfig,
600        chunk_size: usize,
601        scorer: &S,
602    ) -> Result<IndexMap<String, Answer>> {
603        self.evaluate_batch_questions_coarse_to_fine_chunked_with_gating_and_scorer(
604            tokenizer,
605            state,
606            questions,
607            calib_config,
608            coarse_config,
609            chunk_size,
610            None,
611            scorer,
612        )
613    }
614
615    /// 任意の類似度スコアラーおよび明示的なゲーティング設定を指定して、粗密 2 段階探索とマイクロバッチチャンキングを統合した安全推論を実行する。
616    ///
617    /// # 引数
618    /// - `tokenizer`: Jev 高速トークナイザー。
619    /// - `state`: 共通文脈テキスト。
620    /// - `questions`: 質問マップ。
621    /// - `calib_config`: 較正温度設定。
622    /// - `coarse_config`: 粗密探索設定。
623    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
624    /// - `gating_config`: 明示的なゲーティング設定 (任意)。
625    /// - `scorer`: 類似度スコアラー実装。
626    #[allow(clippy::too_many_arguments)]
627    pub fn evaluate_batch_questions_coarse_to_fine_chunked_with_gating_and_scorer<
628        S: CoarseScorer,
629    >(
630        &self,
631        tokenizer: &JevTokenizer,
632        state: &str,
633        questions: &IndexMap<String, Question>,
634        calib_config: &CalibrationConfig,
635        coarse_config: &CoarseToFineConfig,
636        chunk_size: usize,
637        gating_config: Option<&GatingConfig>,
638        scorer: &S,
639    ) -> Result<IndexMap<String, Answer>> {
640        coarse_config.validate()?;
641
642        if questions.is_empty() {
643            return Ok(IndexMap::new());
644        }
645
646        let effective_chunk_size = chunk_size.max(1);
647
648        // 1. 各質問の前処理と粗スクリーニング (Top-M 抽出)
649        let mut processed_questions = IndexMap::with_capacity(questions.len());
650        let mut original_keys_map: IndexMap<String, Vec<String>> = IndexMap::new();
651
652        for (q_key, question) in questions {
653            let should_trigger = question.question_type == QuestionType::Choice
654                && question
655                    .criteria
656                    .as_ref()
657                    .and_then(|c| c.as_map())
658                    .map(|map| coarse_config.should_trigger(map.len()))
659                    .unwrap_or(false);
660
661            if should_trigger {
662                let map = question
663                    .criteria
664                    .as_ref()
665                    .and_then(|c| c.as_map())
666                    .expect("should_trigger により Map が保証されている。");
667
668                let original_keys: Vec<String> = map.keys().cloned().collect();
669                let query = format!("State: {}\nInstructions: {}", state, question.instructions);
670
671                let filtered = filter_top_candidates(&query, map, coarse_config, scorer)?;
672                let sub_question =
673                    Question::new_choice(question.instructions.clone(), filtered.selected);
674
675                processed_questions.insert(q_key.clone(), sub_question);
676                original_keys_map.insert(q_key.clone(), original_keys);
677            } else {
678                processed_questions.insert(q_key.clone(), question.clone());
679            }
680        }
681
682        // 2. 共通 State を 1 度だけエンコードし全チャンクで再利用
683        let state_ids = tokenizer.encode_state(state)?;
684        let mut answers = IndexMap::with_capacity(questions.len());
685        let mut coarse_energies: IndexMap<String, Option<f64>> =
686            IndexMap::with_capacity(original_keys_map.len());
687
688        // 3. マイクロバッチ分割逐次推論
689        for chunk_slice in processed_questions
690            .iter()
691            .collect::<Vec<_>>()
692            .chunks(effective_chunk_size)
693        {
694            let mut sub_questions = IndexMap::with_capacity(chunk_slice.len());
695            for (k, v) in chunk_slice {
696                sub_questions.insert((*k).clone(), (*v).clone());
697            }
698
699            let batch = tokenizer.encode_batch_with_pretokenized_state(
700                &state_ids,
701                &sub_questions,
702                tokenizer.max_sequence_length(),
703            )?;
704
705            let total_logits = batch.dims.batch_size * batch.dims.num_options;
706            let mut flat_logits = vec![0.0f64; total_logits];
707            self.forward_batch_tokenized_into(&batch, &mut flat_logits)?;
708
709            let mut probs_buf = vec![0.0f64; batch.dims.num_options];
710            let num_options = batch.dims.num_options;
711
712            for (i, q_key) in batch.question_keys.iter().enumerate() {
713                let k_i = batch.candidate_counts[i];
714                let start_idx = i * num_options;
715                let valid_logits = &flat_logits[start_idx..start_idx + k_i];
716                let question = sub_questions
717                    .get(q_key)
718                    .expect("キーの一致が保証されている。");
719
720                let mut answer = evaluate_question_with_buf(
721                    question,
722                    valid_logits,
723                    calib_config,
724                    &mut probs_buf[..k_i],
725                )?;
726
727                let energy = calculate_choice_energy(question, valid_logits, calib_config);
728
729                // 縮小されていない質問にはここでゲーティングを適用
730                if !original_keys_map.contains_key(q_key) {
731                    let orig_q = questions.get(q_key);
732                    apply_gating_to_answer_with_energy(
733                        &mut answer,
734                        energy,
735                        Some(state),
736                        Some(q_key),
737                        orig_q,
738                        gating_config,
739                        calib_config,
740                    );
741                } else {
742                    coarse_energies.insert(q_key.clone(), energy);
743                }
744                answers.insert(q_key.clone(), answer);
745            }
746        }
747
748        // 4. スクリーニングされた質問の確率マップを全候補空間に復元し、ゲーティングを適用
749        for (q_key, original_keys) in &original_keys_map {
750            if let Some(answer) = answers.get_mut(q_key)
751                && let Some(ref sub_probs) = answer.probabilities
752            {
753                let full_probs = reconstruct_probabilities(original_keys, sub_probs);
754                answer.probabilities = Some(full_probs);
755                let orig_q = questions.get(q_key);
756                let energy = coarse_energies.get(q_key).copied().flatten();
757                apply_gating_to_answer_with_energy(
758                    answer,
759                    energy,
760                    Some(state),
761                    Some(q_key),
762                    orig_q,
763                    gating_config,
764                    calib_config,
765                );
766            }
767        }
768
769        Ok(answers)
770    }
771
772    /// 既定の語彙スコアラー (`LexicalCoarseScorer`) を用いて、粗密 2 段階探索とマイクロバッチチャンキングを統合した推論を実行する。
773    ///
774    /// # 引数
775    /// - `tokenizer`: Jev 高速トークナイザー。
776    /// - `state`: 共通文脈テキスト。
777    /// - `questions`: 質問マップ。
778    /// - `calib_config`: 較正温度設定。
779    /// - `coarse_config`: 粗密探索設定。
780    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
781    pub fn evaluate_batch_questions_coarse_to_fine_chunked(
782        &self,
783        tokenizer: &JevTokenizer,
784        state: &str,
785        questions: &IndexMap<String, Question>,
786        calib_config: &CalibrationConfig,
787        coarse_config: &CoarseToFineConfig,
788        chunk_size: usize,
789    ) -> Result<IndexMap<String, Answer>> {
790        let scorer = LexicalCoarseScorer::new(tokenizer);
791        self.evaluate_batch_questions_coarse_to_fine_chunked_with_scorer(
792            tokenizer,
793            state,
794            questions,
795            calib_config,
796            coarse_config,
797            chunk_size,
798            &scorer,
799        )
800    }
801
802    /// 既定の語彙スコアラーとゲーティング設定を指定して、粗密 2 段階探索とマイクロバッチチャンキングを統合した推論を実行する。
803    ///
804    /// # 引数
805    /// - `tokenizer`: Jev 高速トークナイザー。
806    /// - `state`: 共通文脈テキスト。
807    /// - `questions`: 質問マップ。
808    /// - `calib_config`: 較正温度設定。
809    /// - `coarse_config`: 粗密探索設定。
810    /// - `chunk_size`: 1 回のフォワードパスで処理する最大質問数。
811    /// - `gating_config`: 明示的なゲーティング設定 (任意)。
812    #[allow(clippy::too_many_arguments)]
813    pub fn evaluate_batch_questions_coarse_to_fine_chunked_with_gating(
814        &self,
815        tokenizer: &JevTokenizer,
816        state: &str,
817        questions: &IndexMap<String, Question>,
818        calib_config: &CalibrationConfig,
819        coarse_config: &CoarseToFineConfig,
820        chunk_size: usize,
821        gating_config: Option<&GatingConfig>,
822    ) -> Result<IndexMap<String, Answer>> {
823        let scorer = LexicalCoarseScorer::new(tokenizer);
824        self.evaluate_batch_questions_coarse_to_fine_chunked_with_gating_and_scorer(
825            tokenizer,
826            state,
827            questions,
828            calib_config,
829            coarse_config,
830            chunk_size,
831            gating_config,
832            &scorer,
833        )
834    }
835
836    /// 粗密二段階階層ルーティング (Coarse-to-Fine Hierarchical Routing) を用いて複数質問を一括評価する。
837    ///
838    /// 共通の階層オントロジーマッピングとルーター設定を各質問に適用し、
839    /// 大分類オントロジー、ODIR Dirichlet 較正、Soft-Beam、エントロピー連動温度較正による低遅延・高精度推論を実行する。
840    ///
841    /// # 引数
842    /// - `tokenizer`: Jev 高速トークナイザー。
843    /// - `state`: 共通文脈テキスト。
844    /// - `questions`: 評価対象の質問群マップ。
845    /// - `mapping`: 大分類・細分類の階層オントロジー定義。
846    /// - `calib_config`: 較正温度設定。
847    /// - `config`: 階層ルーター設定。
848    /// - `gating_config`: ゲーティング閾値設定 (任意)。
849    ///
850    /// # 戻り値
851    /// 各質問キーと階層推論結果 (`HierarchicalResult`) の順序付きマップ。
852    #[allow(clippy::too_many_arguments)]
853    pub fn evaluate_batch_questions_hierarchical(
854        &self,
855        tokenizer: &JevTokenizer,
856        state: &str,
857        questions: &IndexMap<String, Question>,
858        mapping: &HierarchicalMapping,
859        calib_config: &CalibrationConfig,
860        config: HierarchicalRouterConfig,
861        gating_config: Option<&GatingConfig>,
862    ) -> Result<IndexMap<String, HierarchicalResult>> {
863        let router = CoarseToFineRouter::new(self, tokenizer, mapping, calib_config, config)?;
864        let mut results = IndexMap::with_capacity(questions.len());
865
866        for (q_key, question) in questions {
867            let res = router.route(state, &question.instructions, gating_config)?;
868            results.insert(q_key.clone(), res);
869        }
870
871        Ok(results)
872    }
873}
sokuto_runtime/engine/batch.rs

sokuto_runtime/engine/
batch.rs
