sokuto_core/contract/
model_spec.rs

1//! # ONNX モデル入出力テンソル仕様および契約定義モジュール
2//!
3//! Python 学習・エクスポート側と Rust ランタイム推論エンジン側で
4//! 厳密に共有される入出力テンソル名、形状制約、および特殊トークン仕様を定義する。
5
6use crate::error::{CoreError, Result};
7
8/// 入力トークン列テンソル名 (`int64[batch_size, sequence_length]`)。
9pub const TENSOR_INPUT_IDS: &str = "input_ids";
10
11/// アテンションマスクテンソル名 (`int64[batch_size, sequence_length]`)。
12pub const TENSOR_ATTENTION_MASK: &str = "attention_mask";
13
14/// オプションマーカー位置インデックステンソル名 (`int64[batch_size, num_options]`)。
15pub const TENSOR_OP_INDICES: &str = "op_indices";
16
17/// 決定ヘッド出力ロジットテンソル名 (`float32[batch_size, num_options]`)。
18pub const TENSOR_LOGITS: &str = "logits";
19
20/// `State` エンコーダが出力し、`Query` 評価器が参照する最終層隠れ状態テンソル名 (`float32[batch_size, state_seq_len, hidden_size]`)。
21pub const TENSOR_STATE_HIDDEN_STATES: &str = "state_hidden_states";
22
23/// `State` 系列の有効トークンマスクテンソル名 (`int64[batch_size, state_seq_len]`)。
24pub const TENSOR_STATE_MASK: &str = "state_mask";
25
26/// 質問系列の入力トークンID列テンソル名 (`int64[batch_size, query_seq_len]`)。
27pub const TENSOR_QUERY_IDS: &str = "query_ids";
28
29/// 質問系列のアテンションマスクテンソル名 (`int64[batch_size, query_seq_len]`)。
30pub const TENSOR_QUERY_MASK: &str = "query_mask";
31
32/// 各候補の先頭に付与されるオプションマーカートークン文字列。
33pub const TOKEN_OPTION_MARKER: &str = "[OP]";
34
35/// システムで許容される最大入力トークン系列長。
36pub const MAX_SEQUENCE_LENGTH: usize = 8192;
37
38/// システムで許容される単一質問あたりの最大候補数(Choice 型上限)。
39pub const MAX_NUM_OPTIONS: usize = 255;
40
41/// 最小候補数。
42pub const MIN_NUM_OPTIONS: usize = 1;
43
44/// 検証済みのモデル入力次元情報。
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct ModelInputDimensions {
47    /// バッチサイズ(質問並列数)。
48    pub batch_size: usize,
49    /// 入力系列長(パディング含む)。
50    pub sequence_length: usize,
51    /// 評価候補数。
52    pub num_options: usize,
53}
54
55impl ModelInputDimensions {
56    /// 新規次元情報を構築し、仕様制約を満たしているか検証する。
57    pub fn new(batch_size: usize, sequence_length: usize, num_options: usize) -> Result<Self> {
58        if batch_size == 0 {
59            return Err(CoreError::MathError {
60                message: "バッチサイズは 1 以上である必要があります。".to_string(),
61            });
62        }
63        if sequence_length == 0 || sequence_length > MAX_SEQUENCE_LENGTH {
64            return Err(CoreError::MathError {
65                message: format!(
66                    "入力系列長は 1〜{MAX_SEQUENCE_LENGTH} の範囲内である必要がありますが、{sequence_length} が指定されました。"
67                ),
68            });
69        }
70        if !(MIN_NUM_OPTIONS..=MAX_NUM_OPTIONS).contains(&num_options) {
71            return Err(CoreError::InvalidChoiceCount { count: num_options });
72        }
73
74        Ok(Self {
75            batch_size,
76            sequence_length,
77            num_options,
78        })
79    }
80
81    /// 各入力テンソルの形状スライスを検証し、一致する次元情報を返却する。
82    pub fn validate_tensor_shapes(
83        input_ids_shape: &[usize],
84        attention_mask_shape: &[usize],
85        op_indices_shape: &[usize],
86    ) -> Result<Self> {
87        if input_ids_shape.len() != 2 {
88            return Err(CoreError::MathError {
89                message: format!(
90                    "{TENSOR_INPUT_IDS} は 2 次元テンソル [batch, seq_len] である必要があります: {input_ids_shape:?}。"
91                ),
92            });
93        }
94        if attention_mask_shape != input_ids_shape {
95            return Err(CoreError::MathError {
96                message: format!(
97                    "{TENSOR_ATTENTION_MASK} の形状 {attention_mask_shape:?} は {TENSOR_INPUT_IDS} の形状 {input_ids_shape:?} と一致する必要があります。"
98                ),
99            });
100        }
101        if op_indices_shape.len() != 2 {
102            return Err(CoreError::MathError {
103                message: format!(
104                    "{TENSOR_OP_INDICES} は 2 次元テンソル [batch, num_options] である必要があります: {op_indices_shape:?}。"
105                ),
106            });
107        }
108        if op_indices_shape[0] != input_ids_shape[0] {
109            return Err(CoreError::MathError {
110                message: format!(
111                    "{TENSOR_OP_INDICES} のバッチサイズ {} は {TENSOR_INPUT_IDS} のバッチサイズ {} と一致する必要があります。",
112                    op_indices_shape[0], input_ids_shape[0]
113                ),
114            });
115        }
116
117        Self::new(input_ids_shape[0], input_ids_shape[1], op_indices_shape[1])
118    }
119
120    /// 出力 logits テンソルの形状が入力次元と整合しているか検証する。
121    pub fn validate_output_shape(&self, logits_shape: &[usize]) -> Result<()> {
122        if logits_shape.len() != 2 {
123            return Err(CoreError::MathError {
124                message: format!(
125                    "{TENSOR_LOGITS} は 2 次元テンソル [batch, num_options] である必要があります: {logits_shape:?}。"
126                ),
127            });
128        }
129        if logits_shape[0] != self.batch_size || logits_shape[1] != self.num_options {
130            return Err(CoreError::MathError {
131                message: format!(
132                    "{TENSOR_LOGITS} の形状 {logits_shape:?} は期待される形状 [{}, {}] と一致しません。",
133                    self.batch_size, self.num_options
134                ),
135            });
136        }
137        Ok(())
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn test_tensor_shape_validation_success() {
147        let input_ids = [2, 128];
148        let attention_mask = [2, 128];
149        let op_indices = [2, 4];
150
151        let dims =
152            ModelInputDimensions::validate_tensor_shapes(&input_ids, &attention_mask, &op_indices)
153                .unwrap();
154
155        assert_eq!(dims.batch_size, 2);
156        assert_eq!(dims.sequence_length, 128);
157        assert_eq!(dims.num_options, 4);
158
159        // 出力 logits [2, 4] の検証成功
160        assert!(dims.validate_output_shape(&[2, 4]).is_ok());
161
162        // 出力 logits の形状不一致エラー
163        assert!(dims.validate_output_shape(&[2, 5]).is_err());
164    }
165
166    #[test]
167    fn test_tensor_shape_validation_mismatch() {
168        // attention_mask の系列長不一致
169        let res = ModelInputDimensions::validate_tensor_shapes(&[1, 64], &[1, 32], &[1, 3]);
170        assert!(res.is_err());
171
172        // バッチサイズ不一致
173        let res2 = ModelInputDimensions::validate_tensor_shapes(&[2, 64], &[2, 64], &[1, 3]);
174        assert!(res2.is_err());
175    }
176}
sokuto_core/contract/model_spec.rs

sokuto_core/contract/
model_spec.rs
