tekken/config.rs
1use crate::audio::AudioConfig;
2use crate::special_tokens::SpecialTokenInfo;
3use serde::{Deserialize, Serialize};
4
5/// Information about a vocabulary token.
6///
7/// This struct contains metadata about a single token in the vocabulary,
8/// including its rank (position), byte representation, and optional string form.
9///
10/// # Fields
11///
12/// * `rank` - Position of the token in the vocabulary (used as token ID)
13/// * `token_bytes` - Base64-encoded byte representation of the token
14/// * `token_str` - Optional human-readable string representation
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct TokenInfo {
17 /// The position of this token in the vocabulary (used as token ID).
18 pub rank: usize,
19 /// Base64-encoded byte representation of the token.
20 pub token_bytes: String,
21 /// Optional human-readable string representation of the token.
22 pub token_str: Option<String>,
23}
24
25/// Configuration parameters for a Tekken tokenizer.
26///
27/// This struct contains the core configuration needed to initialize a tokenizer,
28/// including the regex pattern for tokenization, vocabulary sizes, and version information.
29///
30/// # Fields
31///
32/// * `pattern` - Regex pattern used for tokenization
33/// * `num_vocab_tokens` - Number of regular vocabulary tokens
34/// * `default_vocab_size` - Default total vocabulary size including special tokens
35/// * `default_num_special_tokens` - Default number of special tokens
36/// * `version` - Tokenizer version string (e.g., "v7")
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct TekkenConfig {
39 /// Regex pattern used for tokenization.
40 pub pattern: String,
41 /// Number of regular vocabulary tokens (excluding special tokens).
42 pub num_vocab_tokens: usize,
43 /// Default total vocabulary size including special tokens.
44 pub default_vocab_size: usize,
45 /// Default number of special tokens.
46 pub default_num_special_tokens: usize,
47 /// Tokenizer version string (e.g., "v7", "v11", "v13").
48 pub version: String,
49}
50
51/// Configuration for image processing (placeholder).
52///
53/// This struct is reserved for future image processing capabilities.
54/// Currently minimal as audio processing is the primary multimodal focus.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct ImageConfig {
57 // Image config fields would go here
58 // For now, we'll keep it minimal as audio is the focus
59}
60
61/// Complete model data loaded from a tokenizer configuration file.
62///
63/// This struct represents the entire configuration and data needed to initialize
64/// a Tekken tokenizer, typically loaded from a JSON file like `tekken.json`.
65///
66/// # Fields
67///
68/// * `vocab` - All vocabulary tokens with their metadata
69/// * `special_tokens` - Optional special token definitions
70/// * `config` - Core tokenizer configuration
71/// * `audio` - Optional audio processing configuration
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct ModelData {
74 /// All vocabulary tokens with their metadata.
75 pub vocab: Vec<TokenInfo>,
76 /// Optional special token definitions (uses defaults if None).
77 pub special_tokens: Option<Vec<SpecialTokenInfo>>,
78 /// Core tokenizer configuration parameters.
79 pub config: TekkenConfig,
80 /// Optional audio processing configuration for multimodal support.
81 pub audio: Option<AudioConfig>,
82}
83
84/// Enumeration of supported tokenizer versions.
85///
86/// Different versions may have different vocabulary sizes, special tokens,
87/// and processing capabilities. This enum provides a type-safe way to
88/// handle version-specific behavior.
89///
90/// # Supported Versions
91///
92/// * `V3` - Early version with basic functionality
93/// * `V7` - Version with enhanced special tokens and audio support
94/// * `V11` - Updated version with additional features
95/// * `V13` - Latest version with full multimodal capabilities
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub enum TokenizerVersion {
98 V3,
99 V7,
100 V11,
101 V13,
102}
103
104impl TokenizerVersion {
105 /// Parses a version string into a `TokenizerVersion`.
106 ///
107 /// # Arguments
108 ///
109 /// * `s` - Version string (e.g., "v7", "v11")
110 ///
111 /// # Returns
112 ///
113 /// The corresponding `TokenizerVersion` if recognized, None otherwise.
114 ///
115 /// # Examples
116 ///
117 /// ```rust
118 /// use tekken::config::TokenizerVersion;
119 ///
120 /// assert_eq!(TokenizerVersion::from_string("v7"), Some(TokenizerVersion::V7));
121 /// assert_eq!(TokenizerVersion::from_string("invalid"), None);
122 /// ```
123 #[must_use]
124 pub fn from_string(s: &str) -> Option<Self> {
125 match s {
126 "v3" => Some(Self::V3),
127 "v7" => Some(Self::V7),
128 "v11" => Some(Self::V11),
129 "v13" => Some(Self::V13),
130 _ => None,
131 }
132 }
133
134 /// Returns the string representation of the version.
135 ///
136 /// # Returns
137 ///
138 /// The version string (e.g., "v7", "v11").
139 ///
140 /// # Examples
141 ///
142 /// ```rust
143 /// use tekken::config::TokenizerVersion;
144 ///
145 /// assert_eq!(TokenizerVersion::V7.as_str(), "v7");
146 /// assert_eq!(TokenizerVersion::V13.as_str(), "v13");
147 /// ```
148 #[must_use]
149 pub const fn as_str(&self) -> &'static str {
150 match self {
151 Self::V3 => "v3",
152 Self::V7 => "v7",
153 Self::V11 => "v11",
154 Self::V13 => "v13",
155 }
156 }
157}