1use indexmap::IndexMap;
7use serde::{Deserialize, Serialize};
8use sokuto_core::contract::calibration::CalibrationConfig;
9use sokuto_core::gating::{DecisionRoute, GatingConfig, GatingMetadata};
10use sokuto_core::math::{normalized_entropy, top_margin};
11use sokuto_core::schema::{Answer, Criteria, Question, QuestionType};
12
13use crate::engine::InferenceEngine;
14use crate::engine::cache::{CachedStateFeatures, DualSessionEngine};
15use crate::engine::coarse::{HierarchicalCoarseScorer, ModelDrivenCoarseScorer};
16use crate::error::{Result, RuntimeError};
17use crate::hierarchical::calibration::{
18 DEFAULT_COARSE_POWER_ALPHA, DEFAULT_ENTROPY_TEMP_GAMMA, DirichletCalibrator,
19 combine_and_reconstruct_probabilities, compute_fine_temperature, compute_normalized_entropy,
20 temperature_scaled_softmax,
21};
22use crate::hierarchical::mapping::HierarchicalMapping;
23use crate::tokenizer::JevTokenizer;
24
25pub const DEFAULT_BEAM_MARGIN_THRESHOLD: f64 = 0.35;
27
28pub const DEFAULT_ESCALATE_CONFIDENCE_THRESHOLD: f64 = 0.25;
30
31pub const DEFAULT_CATCHALL_THRESHOLD: f64 = 0.10;
33
34#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
36pub struct HierarchicalRouterConfig {
37 pub beam_margin_threshold: f64,
39 pub escalate_confidence_threshold: f64,
41 pub catchall_threshold: f64,
43 pub entropy_temp_gamma: f64,
45 pub coarse_power_alpha: f64,
47 pub enable_soft_beam: bool,
49 pub enable_early_exit: bool,
51 pub enable_catchall_bias: bool,
53 pub coarse_calibrator: Option<DirichletCalibrator>,
55}
56
57impl Default for HierarchicalRouterConfig {
58 fn default() -> Self {
59 Self {
60 beam_margin_threshold: DEFAULT_BEAM_MARGIN_THRESHOLD,
61 escalate_confidence_threshold: DEFAULT_ESCALATE_CONFIDENCE_THRESHOLD,
62 catchall_threshold: DEFAULT_CATCHALL_THRESHOLD,
63 entropy_temp_gamma: DEFAULT_ENTROPY_TEMP_GAMMA,
64 coarse_power_alpha: DEFAULT_COARSE_POWER_ALPHA,
65 enable_soft_beam: true,
66 enable_early_exit: true,
67 enable_catchall_bias: true,
68 coarse_calibrator: None,
69 }
70 }
71}
72
73impl HierarchicalRouterConfig {
74 pub fn with_coarse_calibrator(mut self, calibrator: DirichletCalibrator) -> Self {
76 self.coarse_calibrator = Some(calibrator);
77 self
78 }
79 pub fn validate(&self) -> Result<()> {
81 if !(0.0..=1.0).contains(&self.beam_margin_threshold) {
82 return Err(RuntimeError::InvalidQuestion(
83 "beam_margin_threshold は 0.0〜1.0 の範囲である必要があります。".to_string(),
84 ));
85 }
86 if !(0.0..=1.0).contains(&self.escalate_confidence_threshold) {
87 return Err(RuntimeError::InvalidQuestion(
88 "escalate_confidence_threshold は 0.0〜1.0 の範囲である必要があります。"
89 .to_string(),
90 ));
91 }
92 if !(0.0..=1.0).contains(&self.catchall_threshold) {
93 return Err(RuntimeError::InvalidQuestion(
94 "catchall_threshold は 0.0〜1.0 の範囲である必要があります。".to_string(),
95 ));
96 }
97 if self.entropy_temp_gamma < 0.0 {
98 return Err(RuntimeError::InvalidQuestion(
99 "entropy_temp_gamma は 0.0 以上である必要があります。".to_string(),
100 ));
101 }
102 if !(0.0..=2.0).contains(&self.coarse_power_alpha) {
103 return Err(RuntimeError::InvalidQuestion(
104 "coarse_power_alpha は 0.0〜2.0 の範囲である必要があります。".to_string(),
105 ));
106 }
107 Ok(())
108 }
109}
110
111#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
113pub struct HierarchicalExecutionTrace {
114 pub top1_coarse: String,
116 pub top1_coarse_prob: f64,
118 pub top2_coarse: Option<String>,
120 pub top2_coarse_prob: Option<f64>,
122 pub coarse_margin: f64,
124 pub coarse_normalized_entropy: f64,
126 pub coarse_composite_confidence: f64,
128 pub early_exit_triggered: bool,
130 pub early_exit_reason: Option<String>,
132 pub soft_beam_triggered: bool,
134 pub selected_coarse_clusters: Vec<String>,
136 pub evaluated_fine_count: usize,
138 pub fine_temperature: f64,
140 pub catchall_bias_triggered: bool,
142}
143
144#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
146pub struct HierarchicalResult {
147 pub answer: Answer,
149 pub trace: HierarchicalExecutionTrace,
151}
152
153impl HierarchicalResult {
154 pub fn into_answer(self) -> Answer {
156 self.answer
157 }
158
159 pub fn answer(&self) -> &Answer {
161 &self.answer
162 }
163
164 pub fn trace(&self) -> &HierarchicalExecutionTrace {
166 &self.trace
167 }
168}
169
170pub enum HierarchicalEngine<'a> {
172 Single(&'a InferenceEngine),
174 Dual {
176 engine: &'a DualSessionEngine,
178 cached_state: &'a CachedStateFeatures,
180 },
181}
182
183pub struct CoarseToFineRouter<'a> {
185 backend: HierarchicalEngine<'a>,
186 tokenizer: &'a JevTokenizer,
187 mapping: &'a HierarchicalMapping,
188 calib_config: &'a CalibrationConfig,
189 config: HierarchicalRouterConfig,
190}
191
192impl<'a> CoarseToFineRouter<'a> {
193 pub fn new(
195 engine: &'a InferenceEngine,
196 tokenizer: &'a JevTokenizer,
197 mapping: &'a HierarchicalMapping,
198 calib_config: &'a CalibrationConfig,
199 config: HierarchicalRouterConfig,
200 ) -> Result<Self> {
201 Self::with_backend(
202 HierarchicalEngine::Single(engine),
203 tokenizer,
204 mapping,
205 calib_config,
206 config,
207 )
208 }
209
210 pub fn new_dual(
212 engine: &'a DualSessionEngine,
213 cached_state: &'a CachedStateFeatures,
214 tokenizer: &'a JevTokenizer,
215 mapping: &'a HierarchicalMapping,
216 calib_config: &'a CalibrationConfig,
217 config: HierarchicalRouterConfig,
218 ) -> Result<Self> {
219 Self::with_backend(
220 HierarchicalEngine::Dual {
221 engine,
222 cached_state,
223 },
224 tokenizer,
225 mapping,
226 calib_config,
227 config,
228 )
229 }
230
231 pub fn with_backend(
233 backend: HierarchicalEngine<'a>,
234 tokenizer: &'a JevTokenizer,
235 mapping: &'a HierarchicalMapping,
236 calib_config: &'a CalibrationConfig,
237 config: HierarchicalRouterConfig,
238 ) -> Result<Self> {
239 config.validate()?;
240 mapping.validate()?;
241
242 if let Some(ref calibrator) = config.coarse_calibrator
243 && calibrator.num_classes != mapping.total_coarse_count()
244 {
245 return Err(RuntimeError::InvalidQuestion(format!(
246 "DirichletCalibrator のクラス数 ({}) がマッピングの大分類数 ({}) と一致しません。",
247 calibrator.num_classes,
248 mapping.total_coarse_count()
249 )));
250 }
251
252 Ok(Self {
253 backend,
254 tokenizer,
255 mapping,
256 calib_config,
257 config,
258 })
259 }
260
261 pub fn route(
272 &self,
273 state: &str,
274 instructions: &str,
275 gating_config: Option<&GatingConfig>,
276 ) -> Result<HierarchicalResult> {
277 let coarse_probabilities = match self.backend {
279 HierarchicalEngine::Single(single_engine) => {
280 let coarse_scorer = ModelDrivenCoarseScorer::new(single_engine, self.tokenizer)
281 .with_calibration(self.calib_config);
282 coarse_scorer.score_coarse(state, instructions, &self.mapping.coarse_categories)?
283 }
284 HierarchicalEngine::Dual {
285 engine,
286 cached_state,
287 } => {
288 let coarse_question = Question {
289 question_type: QuestionType::Choice,
290 instructions: instructions.to_string(),
291 criteria: Some(sokuto_core::schema::Criteria::Map(
292 self.mapping.coarse_categories.clone(),
293 )),
294 };
295 let raw_logits =
296 engine.forward_question_raw(self.tokenizer, cached_state, &coarse_question)?;
297 let base_temp = self
298 .calib_config
299 .get_temperature(QuestionType::Choice, self.mapping.coarse_categories.len());
300 let probs = temperature_scaled_softmax(&raw_logits, base_temp);
301 let mut map = IndexMap::with_capacity(self.mapping.coarse_categories.len());
302 for (i, key) in self.mapping.coarse_categories.keys().enumerate() {
303 map.insert(key.clone(), probs[i]);
304 }
305 map
306 }
307 };
308
309 let coarse_probabilities = if let Some(ref calibrator) = self.config.coarse_calibrator {
311 let all_coarse = self.mapping.all_coarse_keys();
312 let uncalibrated: Vec<f64> = all_coarse
313 .iter()
314 .map(|k| coarse_probabilities.get(k).copied().unwrap_or(0.0))
315 .collect();
316 let calibrated_vec = calibrator.calibrate(&uncalibrated)?;
317 let mut calibrated_map = IndexMap::with_capacity(all_coarse.len());
318 for (k, p) in all_coarse.iter().zip(calibrated_vec) {
319 calibrated_map.insert(k.clone(), p);
320 }
321 calibrated_map
322 } else {
323 coarse_probabilities
324 };
325
326 let all_coarse_keys = self.mapping.all_coarse_keys();
328 let mut sorted_coarse: Vec<(&str, f64)> = all_coarse_keys
329 .iter()
330 .map(|k| {
331 (
332 k.as_str(),
333 coarse_probabilities.get(k).copied().unwrap_or(0.0),
334 )
335 })
336 .collect();
337
338 sorted_coarse.sort_by(|a, b| {
339 b.1.partial_cmp(&a.1)
340 .unwrap_or(std::cmp::Ordering::Equal)
341 .then_with(|| {
342 let idx_a = all_coarse_keys
343 .iter()
344 .position(|k| k == a.0)
345 .unwrap_or(usize::MAX);
346 let idx_b = all_coarse_keys
347 .iter()
348 .position(|k| k == b.0)
349 .unwrap_or(usize::MAX);
350 idx_a.cmp(&idx_b)
351 })
352 });
353
354 if sorted_coarse.is_empty() {
355 return Err(RuntimeError::InvalidQuestion(
356 "第 1 パスの大分類確率分布が空です。".to_string(),
357 ));
358 }
359
360 let top1_coarse = sorted_coarse[0].0.to_string();
361 let top1_prob = sorted_coarse[0].1;
362 let (top2_coarse, top2_prob): (Option<String>, Option<f64>) = if sorted_coarse.len() > 1 {
363 (
364 Some(sorted_coarse[1].0.to_string()),
365 Some(sorted_coarse[1].1),
366 )
367 } else {
368 (None, None)
369 };
370
371 let coarse_probs_slice: Vec<f64> = sorted_coarse.iter().map(|(_, p)| *p).collect();
372 let (_, norm_entropy) = compute_normalized_entropy(&coarse_probs_slice);
373
374 let coarse_margin: f64 = if let Some(p2) = top2_prob {
375 (top1_prob - p2).max(0.0)
376 } else {
377 1.0
378 };
379
380 let composite_confidence: f64 = top1_prob * (1.0 - norm_entropy);
382
383 if self.config.enable_early_exit
385 && composite_confidence < self.config.escalate_confidence_threshold
386 {
387 let trace = HierarchicalExecutionTrace {
388 top1_coarse: top1_coarse.clone(),
389 top1_coarse_prob: top1_prob,
390 top2_coarse: top2_coarse.clone(),
391 top2_coarse_prob: top2_prob,
392 coarse_margin,
393 coarse_normalized_entropy: norm_entropy,
394 coarse_composite_confidence: composite_confidence,
395 early_exit_triggered: true,
396 early_exit_reason: Some(format!(
397 "複合確信度 ({:.4}) が閾値 ({:.4}) 未満のため Early-Exit しました。",
398 composite_confidence, self.config.escalate_confidence_threshold
399 )),
400 soft_beam_triggered: false,
401 selected_coarse_clusters: vec![top1_coarse.clone()],
402 evaluated_fine_count: 0,
403 fine_temperature: 1.0,
404 catchall_bias_triggered: false,
405 };
406
407 let all_fine_keys = self.mapping.all_fine_keys();
409 let uniform_prob = 1.0 / (all_fine_keys.len() as f64);
410 let mut full_probs = IndexMap::with_capacity(all_fine_keys.len());
411 for k in all_fine_keys {
412 full_probs.insert(k.clone(), uniform_prob);
413 }
414
415 let fallback_gating = GatingMetadata {
416 route: DecisionRoute::ConfirmOrEscalate,
417 confidence: composite_confidence,
418 entropy: Some(1.0),
419 margin: Some(0.0),
420 energy: None,
421 is_ood: false,
422 reason: "Hierarchical Early-Exit: composite confidence below threshold".to_string(),
423 escalation: None,
424 };
425
426 let fallback_answer =
427 Answer::choice(all_fine_keys[0].clone(), full_probs, composite_confidence)
428 .with_gating(fallback_gating);
429
430 return Ok(HierarchicalResult {
431 answer: fallback_answer,
432 trace,
433 });
434 }
435
436 let mut selected_clusters = vec![top1_coarse.clone()];
438 let mut soft_beam_triggered = false;
439
440 if self.config.enable_soft_beam
441 && coarse_margin < self.config.beam_margin_threshold
442 && let Some(ref c2) = top2_coarse
443 {
444 selected_clusters.push(c2.clone());
445 soft_beam_triggered = true;
446 }
447
448 let mut fine_candidates = IndexMap::new();
449 for cluster in &selected_clusters {
450 if let Some(f_keys) = self.mapping.fine_keys_for_coarse(cluster) {
451 for fk in f_keys {
452 if let Some(desc) = self.mapping.fine_criteria.get(fk) {
453 fine_candidates.insert(fk.clone(), desc.clone());
454 }
455 }
456 }
457 }
458
459 if fine_candidates.is_empty() {
460 return Err(RuntimeError::InvalidQuestion(
461 "選択された大分類クラスタ配下の細分類候補が空です。".to_string(),
462 ));
463 }
464
465 let fine_question = Question {
467 question_type: QuestionType::Choice,
468 instructions: instructions.to_string(),
469 criteria: Some(Criteria::Map(fine_candidates.clone())),
470 };
471
472 let raw_logits = match self.backend {
474 HierarchicalEngine::Single(single_engine) => {
475 single_engine.forward_question_raw(self.tokenizer, state, &fine_question)?
476 }
477 HierarchicalEngine::Dual {
478 engine,
479 cached_state,
480 } => engine.forward_question_raw(self.tokenizer, cached_state, &fine_question)?,
481 };
482
483 let fine_temperature =
485 compute_fine_temperature(norm_entropy, self.config.entropy_temp_gamma);
486
487 let base_temp = self
489 .calib_config
490 .get_temperature(QuestionType::Choice, fine_candidates.len());
491 let effective_temp = (base_temp * fine_temperature).max(0.01);
492
493 let temp_scaled_probs = temperature_scaled_softmax(&raw_logits, effective_temp);
495
496 let mut evaluated_fine_probs = IndexMap::with_capacity(fine_candidates.len());
497 for (i, (k, _)) in fine_candidates.iter().enumerate() {
498 evaluated_fine_probs.insert(k.clone(), temp_scaled_probs[i]);
499 }
500
501 let full_probabilities = combine_and_reconstruct_probabilities(
502 self.mapping.all_fine_keys(),
503 &self.mapping.fine_to_coarse,
504 &coarse_probabilities,
505 &evaluated_fine_probs,
506 self.config.coarse_power_alpha,
507 )?;
508
509 let all_fine_keys = self.mapping.all_fine_keys();
511 let mut sorted_full: Vec<(&str, f64)> = all_fine_keys
512 .iter()
513 .map(|k| {
514 (
515 k.as_str(),
516 full_probabilities.get(k).copied().unwrap_or(0.0),
517 )
518 })
519 .collect();
520
521 sorted_full.sort_by(|a, b| {
522 b.1.partial_cmp(&a.1)
523 .unwrap_or(std::cmp::Ordering::Equal)
524 .then_with(|| {
525 let idx_a = all_fine_keys
526 .iter()
527 .position(|k| k == a.0)
528 .unwrap_or(usize::MAX);
529 let idx_b = all_fine_keys
530 .iter()
531 .position(|k| k == b.0)
532 .unwrap_or(usize::MAX);
533 idx_a.cmp(&idx_b)
534 })
535 });
536
537 let default_winner = sorted_full[0].0.to_string();
538 let default_winner_prob = sorted_full[0].1;
539
540 let mut final_winner = default_winner.clone();
541 let mut catchall_bias_triggered = false;
542
543 if self.config.enable_catchall_bias {
544 for (key, &prob) in &full_probabilities {
546 let desc = self
547 .mapping
548 .fine_criteria
549 .get(key)
550 .map(|s| s.as_str())
551 .unwrap_or("");
552 if self.mapping.is_negative_candidate(key, desc)
553 && prob >= self.config.catchall_threshold
554 {
555 final_winner = key.clone();
556 catchall_bias_triggered = true;
557 break;
558 }
559 }
560 }
561
562 let all_probs_slice: Vec<f64> = all_fine_keys
564 .iter()
565 .map(|k| full_probabilities.get(k).copied().unwrap_or(0.0))
566 .collect();
567
568 let final_norm_entropy = normalized_entropy(&all_probs_slice).unwrap_or(0.0);
569 let final_margin = top_margin(&all_probs_slice).unwrap_or(0.0);
570
571 let final_confidence = if catchall_bias_triggered {
572 full_probabilities
573 .get(&final_winner)
574 .copied()
575 .unwrap_or(default_winner_prob)
576 } else {
577 default_winner_prob
578 };
579
580 let fine_energy = crate::engine::gating::calculate_choice_energy(
582 &fine_question,
583 &raw_logits,
584 self.calib_config,
585 );
586
587 let default_gating_config = GatingConfig::default();
589 let gc = gating_config.unwrap_or(&default_gating_config);
590 let final_route =
591 if final_confidence >= gc.high_threshold && final_margin >= gc.top_margin_threshold {
592 DecisionRoute::AutoExecute
593 } else if final_confidence < gc.low_threshold {
594 DecisionRoute::Fallback
595 } else {
596 DecisionRoute::ConfirmOrEscalate
597 };
598
599 let gating = GatingMetadata {
600 route: final_route,
601 confidence: final_confidence,
602 entropy: Some(final_norm_entropy),
603 margin: Some(final_margin),
604 energy: fine_energy,
605 is_ood: false,
606 reason: format!("Hierarchical decision routing (route={})", final_route),
607 escalation: None,
608 };
609
610 let answer =
611 Answer::choice(final_winner, full_probabilities, final_confidence).with_gating(gating);
612
613 let trace = HierarchicalExecutionTrace {
614 top1_coarse,
615 top1_coarse_prob: top1_prob,
616 top2_coarse,
617 top2_coarse_prob: top2_prob,
618 coarse_margin,
619 coarse_normalized_entropy: norm_entropy,
620 coarse_composite_confidence: composite_confidence,
621 early_exit_triggered: false,
622 early_exit_reason: None,
623 soft_beam_triggered,
624 selected_coarse_clusters: selected_clusters,
625 evaluated_fine_count: fine_candidates.len(),
626 fine_temperature,
627 catchall_bias_triggered,
628 };
629
630 Ok(HierarchicalResult { answer, trace })
631 }
632}