combs_formats/
litertlm.rs

1//! `.litertlm` container reader (LiteRT-LM archive format).
2//!
3//! Layout (recon + official `litertlm_header_schema.fbs` /
4//! `litertlm_read.cc`):
5//!
6//! ```text
7//! 0x00  "LITERTLM"                magic
8//! 0x08  u32 major / minor / patch version
9//! 0x14  4 bytes padding
10//! 0x18  u64 header_end_offset
11//! 0x20  LiteRTLMMetaData flatbuffer (system metadata + section directory)
12//! …     section payloads, each aligned to BLOCK_SIZE = 16 KiB
13//! ```
14//!
15//! Sections (`AnySectionDataType`): TFLiteModel (3) → handed to the
16//! TFLite block at its absolute offset; SP_Tokenizer (4) → the
17//! SentencePiece block; HF_Tokenizer_Zlib (6) → zlib-inflated HF
18//! tokenizer.json. Everything else (LlmMetadataProto, executor metadata,
19//! generic blobs) is ignored for now — config comes from the TFLite
20//! section's own LlmParameters.
21
22use std::path::Path;
23
24use crate::flatbuf::FlatBuffer;
25use crate::source::ModelSource;
26use crate::tflite::TfliteSource;
27use crate::{FormatError, Result};
28
29fn bad(what: impl Into<String>) -> FormatError {
30    FormatError::Safetensors(format!("litertlm: {}", what.into()))
31}
32
33const SECTION_TFLITE_MODEL: u8 = 3;
34const SECTION_SP_TOKENIZER: u8 = 4;
35const SECTION_HF_TOKENIZER_ZLIB: u8 = 6;
36
37/// A section directory entry.
38#[derive(Debug)]
39pub struct SectionInfo {
40    pub begin: usize,
41    pub end: usize,
42    pub data_type: u8,
43}
44
45/// Reads the section directory of a `.litertlm` file header.
46pub fn read_sections(header: &[u8]) -> Result<Vec<SectionInfo>> {
47    if header.len() < 0x20 || &header[0..8] != b"LITERTLM" {
48        return Err(bad("bad magic"));
49    }
50    let major = u32_le(header, 0x08)?;
51    if major != 1 {
52        return Err(bad(format!("unsupported major version {major}")));
53    }
54    let header_end = u64_le(header, 0x18)? as usize;
55    if header_end > header.len() || header_end < 0x20 {
56        return Err(bad("header end out of range"));
57    }
58    let fb = FlatBuffer::new(&header[0x20..header_end], None)?;
59    let root = fb.root();
60    // LiteRTLMMetaData.section_metadata (field 1) → SectionMetadata.objects (field 0)
61    let sm = fb
62        .uoffset(root, 1)
63        .ok_or_else(|| bad("no section_metadata"))?;
64    let objects = fb.table_vector(sm, 0)?;
65    let mut out = Vec::with_capacity(objects.len());
66    for obj in objects {
67        let begin = fb.scalar_u64(obj, 1).ok_or_else(|| bad("section missing begin"))? as usize;
68        let end = fb.scalar_u64(obj, 2).ok_or_else(|| bad("section missing end"))? as usize;
69        let data_type = fb.scalar_u8(obj, 3).unwrap_or(0);
70        out.push(SectionInfo { begin, end, data_type });
71    }
72    Ok(out)
73}
74
75/// Opens a `.litertlm` file as a [`ModelSource`]: finds the TFLiteModel
76/// section and delegates to the TFLite block at its absolute offset;
77/// an SP_Tokenizer section, when present, overrides the TFLite
78/// section's own tokenizer metadata. (HF_Tokenizer_Zlib override lands
79/// with the first archive that carries one.)
80pub fn open_litertlm(path: &Path) -> Result<Box<dyn ModelSource>> {
81    let head = {
82        use std::io::Read;
83        let mut f = std::fs::File::open(path)?;
84        let mut buf = vec![0u8; 1024 * 1024]; // header_end is always < 1MB
85        let n = f.read(&mut buf)?;
86        buf.truncate(n);
87        buf
88    };
89    let sections = read_sections(&head)?;
90    let tflite = sections
91        .iter()
92        .find(|s| s.data_type == SECTION_TFLITE_MODEL)
93        .ok_or_else(|| bad("no TFLiteModel section"))?;
94    // SP section bytes live at section offsets — read them eagerly
95    // (tokenizer blobs are a few MB at most).
96    let spm: Option<Vec<u8>> = sections
97        .iter()
98        .find(|s| s.data_type == SECTION_SP_TOKENIZER)
99        .map(|s| {
100            use std::io::{Read, Seek, SeekFrom};
101            let mut f = std::fs::File::open(path)?;
102            f.seek(SeekFrom::Start(s.begin as u64))?;
103            let mut buf = vec![0u8; s.end - s.begin];
104            f.read_exact(&mut buf)?;
105            Ok::<Vec<u8>, FormatError>(buf)
106        })
107        .transpose()?;
108    Ok(Box::new(TfliteSource::load_at_with_spm(
109        path,
110        tflite.begin,
111        spm.as_deref(),
112    )?))
113}
114
115fn u32_le(d: &[u8], pos: usize) -> Result<u32> {
116    let b = d.get(pos..pos + 4).ok_or_else(|| bad("u32 out of bounds"))?;
117    Ok(u32::from_le_bytes([b[0], b[1], b[2], b[3]]))
118}
119
120fn u64_le(d: &[u8], pos: usize) -> Result<u64> {
121    let b = d.get(pos..pos + 8).ok_or_else(|| bad("u64 out of bounds"))?;
122    Ok(u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]))
123}
combs_formats/litertlm.rs

combs_formats/
litertlm.rs
