1use crate::error::{AprFormatError, Result};
15use crate::model_card::ModelCard;
16use serde::{Deserialize, Serialize};
17use std::collections::HashMap;
18
19pub const MAGIC: [u8; 4] = [0x41, 0x50, 0x52, 0x4E];
21
22pub const FORMAT_VERSION: (u8, u8) = (1, 0);
24
25pub const HEADER_SIZE: usize = 32;
27
28pub const MAX_UNCOMPRESSED_SIZE: u32 = 1024 * 1024 * 1024;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33#[repr(u16)]
34pub enum ModelType {
35 LinearRegression = 0x0001,
37 LogisticRegression = 0x0002,
39 DecisionTree = 0x0003,
41 RandomForest = 0x0004,
43 GradientBoosting = 0x0005,
45 KMeans = 0x0006,
47 Pca = 0x0007,
49 NaiveBayes = 0x0008,
51 Knn = 0x0009,
53 Svm = 0x000A,
55 NgramLm = 0x0010,
57 Tfidf = 0x0011,
59 CountVectorizer = 0x0012,
61 NeuralSequential = 0x0020,
63 NeuralCustom = 0x0021,
65 ContentRecommender = 0x0030,
67 MixtureOfExperts = 0x0040,
69 Custom = 0x00FF,
71}
72
73impl ModelType {
74 #[must_use]
76 pub fn from_u16(value: u16) -> Option<Self> {
77 match value {
78 0x0001 => Some(Self::LinearRegression),
79 0x0002 => Some(Self::LogisticRegression),
80 0x0003 => Some(Self::DecisionTree),
81 0x0004 => Some(Self::RandomForest),
82 0x0005 => Some(Self::GradientBoosting),
83 0x0006 => Some(Self::KMeans),
84 0x0007 => Some(Self::Pca),
85 0x0008 => Some(Self::NaiveBayes),
86 0x0009 => Some(Self::Knn),
87 0x000A => Some(Self::Svm),
88 0x0010 => Some(Self::NgramLm),
89 0x0011 => Some(Self::Tfidf),
90 0x0012 => Some(Self::CountVectorizer),
91 0x0020 => Some(Self::NeuralSequential),
92 0x0021 => Some(Self::NeuralCustom),
93 0x0030 => Some(Self::ContentRecommender),
94 0x0040 => Some(Self::MixtureOfExperts),
95 0x00FF => Some(Self::Custom),
96 _ => None,
97 }
98 }
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
103#[repr(u8)]
104pub enum Compression {
105 None = 0x00,
107 #[default]
109 ZstdDefault = 0x01,
110 ZstdMax = 0x02,
112 Lz4 = 0x03,
114}
115
116impl Compression {
117 #[must_use]
119 pub fn from_u8(value: u8) -> Option<Self> {
120 match value {
121 0x00 => Some(Self::None),
122 0x01 => Some(Self::ZstdDefault),
123 0x02 => Some(Self::ZstdMax),
124 0x03 => Some(Self::Lz4),
125 _ => None,
126 }
127 }
128}
129
130#[derive(Debug, Clone, Copy, Default)]
132pub struct Flags(u8);
133
134impl Flags {
135 pub const ENCRYPTED: u8 = 0b0000_0001;
137 pub const SIGNED: u8 = 0b0000_0010;
139 pub const STREAMING: u8 = 0b0000_0100;
141 pub const LICENSED: u8 = 0b0000_1000;
143 pub const TRUENO_NATIVE: u8 = 0b0001_0000;
145 pub const QUANTIZED: u8 = 0b0010_0000;
147 pub const HAS_MODEL_CARD: u8 = 0b0100_0000;
149
150 #[must_use]
152 pub fn new() -> Self {
153 Self(0)
154 }
155
156 #[must_use]
158 pub fn with_encrypted(mut self) -> Self {
159 self.0 |= Self::ENCRYPTED;
160 self
161 }
162
163 #[must_use]
165 pub fn with_signed(mut self) -> Self {
166 self.0 |= Self::SIGNED;
167 self
168 }
169
170 #[must_use]
172 pub fn with_streaming(mut self) -> Self {
173 self.0 |= Self::STREAMING;
174 self
175 }
176
177 #[must_use]
179 pub fn with_licensed(mut self) -> Self {
180 self.0 |= Self::LICENSED;
181 self
182 }
183
184 #[must_use]
186 pub fn with_trueno_native(mut self) -> Self {
187 self.0 |= Self::TRUENO_NATIVE;
188 self
189 }
190
191 #[must_use]
193 pub fn with_quantized(mut self) -> Self {
194 self.0 |= Self::QUANTIZED;
195 self
196 }
197
198 #[must_use]
200 pub fn with_model_card(mut self) -> Self {
201 self.0 |= Self::HAS_MODEL_CARD;
202 self
203 }
204
205 #[must_use]
207 pub fn is_encrypted(self) -> bool {
208 self.0 & Self::ENCRYPTED != 0
209 }
210
211 #[must_use]
213 pub fn is_signed(self) -> bool {
214 self.0 & Self::SIGNED != 0
215 }
216
217 #[must_use]
219 pub fn is_streaming(self) -> bool {
220 self.0 & Self::STREAMING != 0
221 }
222
223 #[must_use]
225 pub fn is_licensed(self) -> bool {
226 self.0 & Self::LICENSED != 0
227 }
228
229 #[must_use]
231 pub fn is_trueno_native(self) -> bool {
232 self.0 & Self::TRUENO_NATIVE != 0
233 }
234
235 #[must_use]
237 pub fn is_quantized(self) -> bool {
238 self.0 & Self::QUANTIZED != 0
239 }
240
241 #[must_use]
243 pub fn has_model_card(self) -> bool {
244 self.0 & Self::HAS_MODEL_CARD != 0
245 }
246
247 #[must_use]
249 pub fn bits(self) -> u8 {
250 self.0
251 }
252
253 #[must_use]
255 pub fn from_bits(bits: u8) -> Self {
256 Self(bits & 0b0111_1111)
257 }
258}
259
260#[derive(Debug, Clone)]
262pub struct Header {
263 pub magic: [u8; 4],
265 pub version: (u8, u8),
267 pub model_type: ModelType,
269 pub metadata_size: u32,
271 pub payload_size: u32,
273 pub uncompressed_size: u32,
275 pub compression: Compression,
277 pub flags: Flags,
279 pub quality_score: u8,
282}
283
284impl Header {
285 #[must_use]
287 pub fn new(model_type: ModelType) -> Self {
288 Self {
289 magic: MAGIC,
290 version: FORMAT_VERSION,
291 model_type,
292 metadata_size: 0,
293 payload_size: 0,
294 uncompressed_size: 0,
295 compression: Compression::default(),
296 flags: Flags::default(),
297 quality_score: 0,
298 }
299 }
300
301 #[must_use]
303 pub fn to_bytes(&self) -> [u8; HEADER_SIZE] {
304 let mut bytes = [0u8; HEADER_SIZE];
305 bytes[0..4].copy_from_slice(&self.magic);
306 bytes[4] = self.version.0;
307 bytes[5] = self.version.1;
308 let model_type = self.model_type as u16;
309 bytes[6..8].copy_from_slice(&model_type.to_le_bytes());
310 bytes[8..12].copy_from_slice(&self.metadata_size.to_le_bytes());
311 bytes[12..16].copy_from_slice(&self.payload_size.to_le_bytes());
312 bytes[16..20].copy_from_slice(&self.uncompressed_size.to_le_bytes());
313 bytes[20] = self.compression as u8;
314 bytes[21] = self.flags.bits();
315 bytes[22] = self.quality_score;
316 bytes
318 }
319
320 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
326 if bytes.len() < HEADER_SIZE {
327 return Err(AprFormatError::FormatError {
328 message: format!(
329 "Header too short: {} bytes, expected {}",
330 bytes.len(),
331 HEADER_SIZE
332 ),
333 });
334 }
335
336 let magic: [u8; 4] = bytes[0..4]
337 .try_into()
338 .map_err(|_| AprFormatError::FormatError {
339 message: "header slice too short for magic".to_string(),
340 })?;
341 if magic != MAGIC {
342 return Err(AprFormatError::FormatError {
343 message: format!(
344 "Invalid magic number: {:02X}{:02X}{:02X}{:02X}, expected APRN",
345 magic[0], magic[1], magic[2], magic[3]
346 ),
347 });
348 }
349
350 let version = (bytes[4], bytes[5]);
351 if version.0 > FORMAT_VERSION.0 {
352 return Err(AprFormatError::UnsupportedVersion {
353 found: version,
354 supported: FORMAT_VERSION,
355 });
356 }
357
358 let model_type_raw = u16::from_le_bytes([bytes[6], bytes[7]]);
359 let model_type =
360 ModelType::from_u16(model_type_raw).ok_or_else(|| AprFormatError::FormatError {
361 message: format!("Unknown model type: 0x{model_type_raw:04X}"),
362 })?;
363
364 let metadata_size = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]);
365 let payload_size = u32::from_le_bytes([bytes[12], bytes[13], bytes[14], bytes[15]]);
366 let uncompressed_size = u32::from_le_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]);
367
368 if uncompressed_size > MAX_UNCOMPRESSED_SIZE {
369 return Err(AprFormatError::FormatError {
370 message: format!(
371 "Uncompressed size {uncompressed_size} exceeds maximum {MAX_UNCOMPRESSED_SIZE} (compression bomb protection)"
372 ),
373 });
374 }
375
376 let compression =
377 Compression::from_u8(bytes[20]).ok_or_else(|| AprFormatError::FormatError {
378 message: format!("Unknown compression algorithm: 0x{:02X}", bytes[20]),
379 })?;
380 let flags = Flags::from_bits(bytes[21]);
381 let quality_score = bytes[22];
382
383 Ok(Self {
384 magic,
385 version,
386 model_type,
387 metadata_size,
388 payload_size,
389 uncompressed_size,
390 compression,
391 flags,
392 quality_score,
393 })
394 }
395}
396
397#[derive(Debug, Clone, Serialize, Deserialize)]
399pub struct TrainingInfo {
400 #[serde(default, skip_serializing_if = "Option::is_none")]
402 pub samples: Option<usize>,
403 #[serde(default, skip_serializing_if = "Option::is_none")]
405 pub duration_ms: Option<u64>,
406 #[serde(default, skip_serializing_if = "Option::is_none")]
408 pub source: Option<String>,
409}
410
411#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
417pub enum DistillMethod {
418 Standard,
420 Progressive,
422 Ensemble,
424}
425
426#[derive(Debug, Clone, Serialize, Deserialize)]
428pub struct TeacherProvenance {
429 pub hash: String,
431 #[serde(default, skip_serializing_if = "Option::is_none")]
433 pub signature: Option<String>,
434 pub model_type: ModelType,
436 pub param_count: u64,
438 #[serde(default, skip_serializing_if = "Option::is_none")]
440 pub ensemble_teachers: Option<Vec<TeacherProvenance>>,
441}
442
443#[derive(Debug, Clone, Serialize, Deserialize)]
445pub struct DistillationParams {
446 pub temperature: f32,
448 pub alpha: f32,
450 #[serde(default, skip_serializing_if = "Option::is_none")]
452 pub beta: Option<f32>,
453 pub epochs: u32,
455 #[serde(default, skip_serializing_if = "Option::is_none")]
457 pub final_loss: Option<f32>,
458}
459
460#[derive(Debug, Clone, Serialize, Deserialize)]
462pub struct LayerMapping {
463 pub student_layer: usize,
465 pub teacher_layer: usize,
467 pub weight: f32,
469}
470
471#[derive(Debug, Clone, Serialize, Deserialize)]
473pub struct DistillationInfo {
474 pub method: DistillMethod,
476 pub teacher: TeacherProvenance,
478 pub params: DistillationParams,
480 #[serde(default, skip_serializing_if = "Option::is_none")]
482 pub layer_mapping: Option<Vec<LayerMapping>>,
483}
484
485#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
491pub enum LicenseTier {
492 Personal,
494 Team,
496 Enterprise,
498 Academic,
500}
501
502#[derive(Debug, Clone, Serialize, Deserialize)]
504pub struct LicenseInfo {
505 pub uuid: String,
507 pub hash: String,
509 #[serde(default, skip_serializing_if = "Option::is_none")]
511 pub expiry: Option<String>,
512 #[serde(default, skip_serializing_if = "Option::is_none")]
514 pub seats: Option<u32>,
515 #[serde(default, skip_serializing_if = "Option::is_none")]
517 pub licensee: Option<String>,
518 pub tier: LicenseTier,
520}
521
522#[derive(Debug, Clone, Serialize, Deserialize)]
524pub struct Metadata {
525 pub created_at: String,
527 pub aprender_version: String,
529 #[serde(default, skip_serializing_if = "Option::is_none")]
531 pub model_name: Option<String>,
532 #[serde(default, skip_serializing_if = "Option::is_none")]
534 pub description: Option<String>,
535 #[serde(default, skip_serializing_if = "Option::is_none")]
537 pub training: Option<TrainingInfo>,
538 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
540 pub hyperparameters: HashMap<String, serde_json::Value>,
541 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
543 pub metrics: HashMap<String, serde_json::Value>,
544 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
546 pub custom: HashMap<String, serde_json::Value>,
547 #[serde(default, skip_serializing_if = "Option::is_none")]
549 pub distillation: Option<String>,
550 #[serde(default, skip_serializing_if = "Option::is_none")]
552 pub distillation_info: Option<DistillationInfo>,
553 #[serde(default, skip_serializing_if = "Option::is_none")]
555 pub license: Option<LicenseInfo>,
556 #[serde(default, skip_serializing_if = "Option::is_none")]
558 pub model_card: Option<ModelCard>,
559}
560
561impl Default for Metadata {
562 fn default() -> Self {
563 Self {
564 created_at: chrono_lite_now(),
565 aprender_version: env!("CARGO_PKG_VERSION").to_string(),
566 model_name: None,
567 description: None,
568 training: None,
569 hyperparameters: HashMap::new(),
570 metrics: HashMap::new(),
571 custom: HashMap::new(),
572 distillation: None,
573 distillation_info: None,
574 license: None,
575 model_card: None,
576 }
577 }
578}
579
580#[must_use]
582pub fn chrono_lite_now() -> String {
583 use std::time::{SystemTime, UNIX_EPOCH};
584 let duration = SystemTime::now()
585 .duration_since(UNIX_EPOCH)
586 .unwrap_or_default();
587 let secs = duration.as_secs();
588 format!("{secs}")
590}
591
592#[derive(Debug, Clone, Default)]
594pub struct SaveOptions {
595 pub compression: Compression,
597 pub metadata: Metadata,
599 pub quality_score: Option<u8>,
605}
606
607impl SaveOptions {
608 #[must_use]
610 pub fn new() -> Self {
611 Self::default()
612 }
613
614 #[must_use]
616 pub fn with_compression(mut self, compression: Compression) -> Self {
617 self.compression = compression;
618 self
619 }
620
621 #[must_use]
623 pub fn with_name(mut self, name: impl Into<String>) -> Self {
624 self.metadata.model_name = Some(name.into());
625 self
626 }
627
628 #[must_use]
630 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
631 self.metadata.description = Some(desc.into());
632 self
633 }
634
635 #[must_use]
637 pub fn with_distillation_info(mut self, info: DistillationInfo) -> Self {
638 self.metadata.distillation_info = Some(info);
639 self
640 }
641
642 #[must_use]
644 pub fn with_license(mut self, license: LicenseInfo) -> Self {
645 self.metadata.license = Some(license);
646 self
647 }
648
649 #[must_use]
651 pub fn with_model_card(mut self, card: ModelCard) -> Self {
652 self.metadata.model_card = Some(card);
653 self
654 }
655
656 #[must_use]
663 pub fn with_quality_score(mut self, score: u8) -> Self {
664 self.quality_score = Some(score);
665 self
666 }
667}
668
669#[derive(Debug, Clone)]
671#[allow(clippy::struct_excessive_bools)] pub struct ModelInfo {
673 pub model_type: ModelType,
675 pub format_version: (u8, u8),
677 pub metadata: Metadata,
679 pub payload_size: usize,
681 pub uncompressed_size: usize,
683 pub encrypted: bool,
685 pub signed: bool,
687 pub streaming: bool,
689 pub licensed: bool,
691 pub trueno_native: bool,
693 pub quantized: bool,
695 pub has_model_card: bool,
697}