superfluid_daemon/
dylib_tokenizer.rs

1//! A runtime's tokenizer, loaded from its package at run time.
2
3use std::ffi::{c_char, c_void, CStr, CString};
4use std::path::{Path, PathBuf};
5
6use superfluid_engine::tokenizer_abi::{self as abi, symbol};
7use superfluid_engine::Tokenizer;
8
9struct Fns {
10    close: abi::CloseFn,
11    encode: abi::EncodeFn,
12    token_bytes: abi::TokenBytesFn,
13    vocab_size: abi::VocabSizeFn,
14    special_tokens: abi::SpecialTokensFn,
15    bos: abi::BosFn,
16    eos: abi::EosFn,
17    chat_template: abi::ChatTemplateFn,
18}
19
20pub struct DylibTokenizer {
21    tok: *mut c_void,
22    f: Fns,
23    library: PathBuf,
24}
25
26// SAFETY: the interface's tokenizers are immutable once open — every call
27// but `close` is a read, safe from any thread concurrently (the contract
28// `superfluid_engine::Tokenizer: Send + Sync` states for every tokenizer).
29unsafe impl Send for DylibTokenizer {}
30// SAFETY: as above.
31unsafe impl Sync for DylibTokenizer {}
32
33fn dl_error() -> String {
34    // SAFETY: dlerror returns NULL or a NUL-terminated message owned by the
35    // loader, valid until the next dl* call on this thread.
36    let e = unsafe { libc::dlerror() };
37    if e.is_null() {
38        "unknown loader error".into()
39    } else {
40        // SAFETY: as above.
41        unsafe { CStr::from_ptr(e) }.to_string_lossy().into_owned()
42    }
43}
44
45impl DylibTokenizer {
46    pub fn open(library: &Path, model: &Path) -> Result<DylibTokenizer, String> {
47        let lib_c = CString::new(library.to_string_lossy().as_bytes()).map_err(|_| "NUL in library path".to_string())?;
48        // SAFETY: a NUL-terminated path; RTLD_LOCAL keeps the library's
49        // symbols (its own llama.cpp) out of the daemon's namespace.
50        let lib = unsafe { libc::dlopen(lib_c.as_ptr(), libc::RTLD_NOW | libc::RTLD_LOCAL) };
51        if lib.is_null() {
52            return Err(format!("could not load the tokenizer library {}: {}", library.display(), dl_error()));
53        }
54        let sym = |name: &[u8]| -> Result<*mut c_void, String> {
55            // SAFETY: `lib` is a live handle; `name` is NUL-terminated.
56            let p = unsafe { libc::dlsym(lib, name.as_ptr().cast::<c_char>()) };
57            if p.is_null() {
58                let name = String::from_utf8_lossy(&name[..name.len() - 1]).into_owned();
59                return Err(format!("{} is not a superfluid tokenizer library (no {name})", library.display()));
60            }
61            Ok(p)
62        };
63        // SAFETY: each symbol is the exported function of that name, whose
64        // signature the interface fixes (so for every transmute below).
65        let version: abi::AbiVersionFn = unsafe { std::mem::transmute(sym(symbol::ABI_VERSION)?) };
66        // SAFETY: as above.
67        let have = unsafe { version() };
68        if have != abi::ABI_VERSION {
69            return Err(format!(
70                "the tokenizer library {} speaks interface {have}; this superfluid speaks {}",
71                library.display(),
72                abi::ABI_VERSION
73            ));
74        }
75        // SAFETY: as above.
76        let open: abi::OpenFn = unsafe { std::mem::transmute(sym(symbol::OPEN)?) };
77        // SAFETY: as above, for each field.
78        let f = unsafe {
79            Fns {
80                close: std::mem::transmute::<*mut c_void, abi::CloseFn>(sym(symbol::CLOSE)?),
81                encode: std::mem::transmute::<*mut c_void, abi::EncodeFn>(sym(symbol::ENCODE)?),
82                token_bytes: std::mem::transmute::<*mut c_void, abi::TokenBytesFn>(sym(symbol::TOKEN_BYTES)?),
83                vocab_size: std::mem::transmute::<*mut c_void, abi::VocabSizeFn>(sym(symbol::VOCAB_SIZE)?),
84                special_tokens: std::mem::transmute::<*mut c_void, abi::SpecialTokensFn>(sym(symbol::SPECIAL_TOKENS)?),
85                bos: std::mem::transmute::<*mut c_void, abi::BosFn>(sym(symbol::BOS)?),
86                eos: std::mem::transmute::<*mut c_void, abi::EosFn>(sym(symbol::EOS)?),
87                chat_template: std::mem::transmute::<*mut c_void, abi::ChatTemplateFn>(sym(symbol::CHAT_TEMPLATE)?),
88            }
89        };
90        let model_c = CString::new(model.to_string_lossy().as_bytes()).map_err(|_| "NUL in model path".to_string())?;
91        let mut err = vec![0 as c_char; 1024];
92        // SAFETY: a NUL-terminated path and an error buffer of its length.
93        let tok = unsafe { open(model_c.as_ptr(), err.as_mut_ptr(), err.len()) };
94        if tok.is_null() {
95            // SAFETY: the library NUL-terminates what it writes; the buffer
96            // was zeroed, so an unwritten one reads as empty.
97            let why = unsafe { CStr::from_ptr(err.as_ptr()) }.to_string_lossy().into_owned();
98            return Err(format!("tokenizer load failed for {}: {why}", model.display()));
99        }
100        Ok(DylibTokenizer { tok, f, library: library.to_path_buf() })
101    }
102
103    pub fn library(&self) -> &Path {
104        &self.library
105    }
106
107    fn filled<T: Copy + Default>(guess: usize, mut call: impl FnMut(*mut T, usize) -> usize) -> Option<Vec<T>> {
108        let mut buf = vec![T::default(); guess];
109        let n = call(buf.as_mut_ptr(), buf.len());
110        if n == usize::MAX {
111            return None;
112        }
113        if n > buf.len() {
114            buf = vec![T::default(); n];
115            let again = call(buf.as_mut_ptr(), buf.len());
116            buf.truncate(again.min(n));
117            return Some(buf);
118        }
119        buf.truncate(n);
120        Some(buf)
121    }
122}
123
124impl Drop for DylibTokenizer {
125    fn drop(&mut self) {
126        // SAFETY: the handle is live and released once; the library stays
127        // loaded (see `open`).
128        unsafe { (self.f.close)(self.tok) };
129    }
130}
131
132impl Tokenizer for DylibTokenizer {
133    fn encode(&self, text: &str) -> Vec<u32> {
134        Self::filled(text.len() + 8, |out, cap| {
135            // SAFETY: `text` is a live slice of `len` bytes; `out` holds `cap` ids.
136            unsafe { (self.f.encode)(self.tok, text.as_ptr(), text.len(), out, cap) }
137        })
138        .unwrap_or_default()
139    }
140
141    fn token_bytes(&self, token: u32) -> Vec<u8> {
142        Self::filled(64, |out, cap| {
143            // SAFETY: `out` holds `cap` bytes.
144            unsafe { (self.f.token_bytes)(self.tok, token, out, cap) }
145        })
146        .unwrap_or_default()
147    }
148
149    fn vocab_size(&self) -> u32 {
150        // SAFETY: the handle is live.
151        unsafe { (self.f.vocab_size)(self.tok) }
152    }
153
154    fn special_tokens(&self) -> Vec<(String, u32)> {
155        let packed = Self::filled(16 << 10, |out, cap| {
156            // SAFETY: `out` holds `cap` bytes.
157            unsafe { (self.f.special_tokens)(self.tok, out, cap) }
158        })
159        .unwrap_or_default();
160        abi::unpack_special(&packed)
161    }
162
163    fn bos_token(&self) -> Option<u32> {
164        // SAFETY: the handle is live.
165        let b = unsafe { (self.f.bos)(self.tok) };
166        u32::try_from(b).ok()
167    }
168
169    fn eos_token(&self) -> u32 {
170        // SAFETY: the handle is live.
171        unsafe { (self.f.eos)(self.tok) }
172    }
173
174    fn chat_template_jinja(&self) -> String {
175        let bytes = Self::filled(16 << 10, |out, cap| {
176            // SAFETY: a NULL name asks for the default; `out` holds `cap` bytes.
177            unsafe { (self.f.chat_template)(self.tok, std::ptr::null(), out, cap) }
178        });
179        bytes.map(|b| String::from_utf8_lossy(&b).into_owned()).unwrap_or_default()
180    }
181
182    fn chat_template_named(&self, name: &str) -> Option<String> {
183        let name = CString::new(name).ok()?;
184        let bytes = Self::filled(16 << 10, |out, cap| {
185            // SAFETY: a NUL-terminated name; `out` holds `cap` bytes.
186            unsafe { (self.f.chat_template)(self.tok, name.as_ptr(), out, cap) }
187        })?;
188        Some(String::from_utf8_lossy(&bytes).into_owned())
189    }
190}
191
192pub fn library_file(stem: &str) -> String {
193    if cfg!(target_os = "macos") {
194        format!("lib{stem}.dylib")
195    } else {
196        format!("lib{stem}.so")
197    }
198}
199
200pub fn find_library(worker: &Path, stem: &str) -> Option<PathBuf> {
201    let file = library_file(stem);
202    let dir = worker.parent()?;
203    [dir.join(&file), dir.join("..").join("lib").join(&file)].into_iter().find(|p| p.is_file())
204}
superfluid_daemon/dylib_tokenizer.rs

superfluid_daemon/
dylib_tokenizer.rs
