Skip to main content

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}