transcribe_cpp/
model.rs

1//! [`Model`] — a loaded GGUF, shareable across threads.
2//!
3//! A `Model` is `Arc`-backed and `Send + Sync`: cloning it is cheap and hands
4//! out another handle to the same native model. The native model is freed when
5//! the last handle AND every [`Session`](crate::Session) derived from it have
6//! been dropped — Rust ownership gives the C "model must outlive its sessions"
7//! contract (and close-ordering safety) for free, in any drop order.
8//!
9//! The `Arc` also carries the per-model compute lock that enforces the C
10//! library's 0.x concurrency limitation: at most one `run` / `run_batch` /
11//! active stream may be in flight across ALL sessions of a model. One-shot
12//! runs/batches queue on the lock (serialized); an *active stream* spans many
13//! native calls (`begin`→`feed`*→`finalize`), so the lock also carries a flag
14//! marking a stream in flight — any run/batch/stream that would overlap it is
15//! refused with [`Error::Busy`](crate::Error) rather than allowed to race.
16
17use std::ffi::CString;
18use std::path::Path;
19use std::sync::{Arc, Mutex};
20
21use transcribe_cpp_sys as sys;
22
23use crate::backend::Device;
24use crate::error::{check, Result};
25use crate::result::owned_str;
26use crate::session::Session;
27use crate::types::{Backend, ExtSlot, Feature, TimestampKind};
28use crate::version;
29
30/// Options for loading a model.
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct ModelOptions {
33    /// Which backend to request. Default [`Backend::Auto`]. An explicit device
34    /// must match a non-Auto backend request.
35    pub backend: Backend,
36    /// Exact process-local compute device. `None` applies the backend's
37    /// automatic policy; `Some` selects that device or fails without fallback.
38    pub device: Option<crate::Device>,
39}
40
41impl Default for ModelOptions {
42    fn default() -> Self {
43        ModelOptions {
44            backend: Backend::Auto,
45            device: None,
46        }
47    }
48}
49
50/// Immutable, model-level capabilities read from GGUF metadata.
51#[derive(Debug, Clone, PartialEq)]
52pub struct Capabilities {
53    pub native_sample_rate: i32,
54    /// Supported language codes (empty if the model is language-agnostic).
55    pub languages: Vec<String>,
56    /// Supported translation target language codes (empty if not advertised).
57    pub translate_target_languages: Vec<String>,
58    /// The finest timestamp granularity the model can produce.
59    pub max_timestamp_kind: TimestampKind,
60    pub supports_language_detect: bool,
61    pub supports_translate: bool,
62    pub supports_streaming: bool,
63    pub supports_spec_decode: bool,
64    /// Longest accepted audio in ms (0 = no practical limit).
65    pub max_audio_ms: i64,
66}
67
68/// The native model handle plus the per-model compute lock, shared via `Arc`.
69pub(crate) struct ModelInner {
70    pub(crate) ptr: *mut sys::transcribe_model,
71    /// Serializes the compute path across all sessions of this model. The bool
72    /// is `true` while a stream is ACTIVE (from `begin` until the earliest of
73    /// `finalize` / `reset` / the `Stream` being dropped); a held lock plus a
74    /// `true` flag is how `run`/`run_batch`/`stream` detect and refuse an
75    /// overlapping compute. See the module docs.
76    pub(crate) compute_lock: Mutex<bool>,
77}
78
79// SAFETY: the native model is documented thread-safe for query + session
80// creation; the compute path is serialized by `compute_lock`. The raw pointer
81// is only ever used behind `&ModelInner`, never mutated through aliasing.
82unsafe impl Send for ModelInner {}
83unsafe impl Sync for ModelInner {}
84
85impl Drop for ModelInner {
86    fn drop(&mut self) {
87        // Only reached once every Session (each holding an Arc clone) is gone.
88        unsafe { sys::transcribe_model_free(self.ptr) };
89    }
90}
91
92/// A loaded transcription model.
93#[derive(Clone)]
94pub struct Model {
95    pub(crate) inner: Arc<ModelInner>,
96}
97
98impl std::fmt::Debug for Model {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        f.debug_struct("Model")
101            .field("arch", &self.arch())
102            .field("variant", &self.variant())
103            .field("backend", &self.backend())
104            .finish()
105    }
106}
107
108impl Model {
109    /// Load a GGUF model from disk with default options.
110    pub fn load(path: impl AsRef<Path>) -> Result<Model> {
111        Model::load_with(path, &ModelOptions::default())
112    }
113
114    /// Load a GGUF model from disk with explicit options.
115    pub fn load_with(path: impl AsRef<Path>, options: &ModelOptions) -> Result<Model> {
116        // Pre-1.0 base-version lock against the loaded library (once).
117        version::ensure_compatible()?;
118
119        let path = path.as_ref();
120        let c_path = CString::new(path_bytes(path)?)?;
121
122        let mut params: sys::transcribe_model_load_params = unsafe { std::mem::zeroed() };
123        unsafe { sys::transcribe_model_load_params_init(&mut params) };
124        params.backend = options.backend.to_raw();
125        params.device = options
126            .device
127            .as_ref()
128            .map_or(std::ptr::null_mut(), |device| device.handle);
129
130        let mut out: *mut sys::transcribe_model = std::ptr::null_mut();
131        let status = unsafe { sys::transcribe_model_load_file(c_path.as_ptr(), &params, &mut out) };
132        check_model_load(status, &format!("load {}", path.display()))?;
133        debug_assert!(!out.is_null());
134
135        Ok(Model {
136            inner: Arc::new(ModelInner {
137                ptr: out,
138                compute_lock: Mutex::new(false),
139            }),
140        })
141    }
142
143    /// Open a session bound to this model with default session options.
144    pub fn session(&self) -> Result<Session> {
145        self.session_with(&SessionOptions::default())
146    }
147
148    /// Open a session bound to this model with explicit options.
149    pub fn session_with(&self, options: &SessionOptions) -> Result<Session> {
150        Session::new(self, options)
151    }
152
153    /// The model's immutable capabilities.
154    pub fn capabilities(&self) -> Capabilities {
155        let mut caps: sys::transcribe_capabilities = unsafe { std::mem::zeroed() };
156        unsafe { sys::transcribe_capabilities_init(&mut caps) };
157        // A NULL/struct-size fault cannot happen here (we own a valid model
158        // and an init'd struct), so a non-OK status leaves zeroed defaults.
159        let _ = unsafe { sys::transcribe_model_get_capabilities(self.inner.ptr, &mut caps) };
160
161        let mut languages = Vec::new();
162        if !caps.languages.is_null() && caps.n_languages > 0 {
163            let slice =
164                unsafe { std::slice::from_raw_parts(caps.languages, caps.n_languages as usize) };
165            for &lang in slice {
166                languages.push(owned_str(lang));
167            }
168        }
169        let mut translate_target_languages = Vec::new();
170        if !caps.translate_target_languages.is_null() && caps.n_translate_target_languages > 0 {
171            let slice = unsafe {
172                std::slice::from_raw_parts(
173                    caps.translate_target_languages,
174                    caps.n_translate_target_languages as usize,
175                )
176            };
177            for &lang in slice {
178                translate_target_languages.push(owned_str(lang));
179            }
180        }
181
182        Capabilities {
183            native_sample_rate: caps.native_sample_rate,
184            languages,
185            translate_target_languages,
186            max_timestamp_kind: TimestampKind::from_raw(caps.max_timestamp_kind),
187            supports_language_detect: caps.supports_language_detect,
188            supports_translate: caps.supports_translate,
189            supports_streaming: caps.supports_streaming,
190            supports_spec_decode: caps.supports_spec_decode,
191            max_audio_ms: caps.max_audio_ms,
192        }
193    }
194
195    /// Probe a yes/no feature.
196    pub fn supports(&self, feature: Feature) -> bool {
197        unsafe { sys::transcribe_model_supports(self.inner.ptr, feature.to_raw()) }
198    }
199
200    /// Whether this model accepts the given family-extension `kind` in `slot`.
201    pub fn accepts_ext(&self, slot: ExtSlot, kind: u32) -> bool {
202        unsafe { sys::transcribe_model_accepts_ext_kind(self.inner.ptr, slot.to_raw(), kind) }
203    }
204
205    /// The `general.architecture` string, e.g. "parakeet".
206    pub fn arch(&self) -> String {
207        owned_str(unsafe { sys::transcribe_model_arch_string(self.inner.ptr) })
208    }
209
210    /// The `stt.variant` string, e.g. "tdt-0.6b-v2" (may be empty).
211    pub fn variant(&self) -> String {
212        owned_str(unsafe { sys::transcribe_model_variant_string(self.inner.ptr) })
213    }
214
215    /// The runtime backend bound to this model, e.g. "cpu", "metal" — the way
216    /// to detect CPU fallback after requesting a GPU.
217    pub fn backend(&self) -> String {
218        owned_str(unsafe { sys::transcribe_model_backend(self.inner.ptr) })
219    }
220
221    /// The compute [`Device`] this model is running on — the one that owns its
222    /// weights. Its `memory_free` is a live snapshot, so re-call to poll how
223    /// much memory is left on the device after the model loaded. Errors with
224    /// [`Error::Backend`](crate::Error) if the model has no resolved device.
225    pub fn device(&self) -> Result<Device> {
226        let handle = unsafe { sys::transcribe_model_device(self.inner.ptr) };
227        if handle.is_null() {
228            return Err(crate::Error::Backend(
229                "model has no resolved device".to_string(),
230            ));
231        }
232        let mut raw: sys::transcribe_device_info = unsafe { std::mem::zeroed() };
233        unsafe { sys::transcribe_device_info_init(&mut raw) };
234        let status = unsafe { sys::transcribe_device_get_info(handle, &mut raw) };
235        check(status, "device_get_info")?;
236        Ok(Device::from_raw(&raw, handle, None))
237    }
238
239    /// Tokenize plain UTF-8 text into the model's vocabulary (no BOS/EOS, no
240    /// special tags). Errors with [`Error::NotImplemented`](crate::Error) for
241    /// families whose tokenizer has no encode path (e.g. SentencePiece today).
242    pub fn tokenize(&self, text: &str) -> Result<Vec<i32>> {
243        let c_text = CString::new(text)?;
244        // Grow-and-retry on the negative-N "buffer too small" signal.
245        let mut buf = vec![0i32; text.len().max(16)];
246        loop {
247            let n = unsafe {
248                sys::transcribe_tokenize(
249                    self.inner.ptr,
250                    c_text.as_ptr(),
251                    buf.as_mut_ptr(),
252                    buf.len(),
253                )
254            };
255            if n == i32::MIN {
256                return Err(crate::error::Error::NotImplemented(
257                    "model tokenizer has no encode path".to_string(),
258                ));
259            }
260            if n < 0 {
261                buf.resize((-n) as usize, 0);
262                continue;
263            }
264            buf.truncate(n as usize);
265            return Ok(buf);
266        }
267    }
268}
269
270fn check_model_load(status: sys::transcribe_status, context: &str) -> Result<()> {
271    #[cfg(feature = "dynamic-backends")]
272    {
273        if status == sys::transcribe_status::TRANSCRIBE_ERR_BACKEND
274            && crate::backend::device_count() == 0
275        {
276            let code = status.0 as i32;
277            let base = crate::error::status_string(code);
278            return Err(crate::error::Error::Backend(format!(
279                "{context}: {base} (status {code}); dynamic-backends builds \
280                 require backend modules to be loaded before model load. Call \
281                 transcribe_cpp::init_backends_default() when modules are \
282                 bundled next to libtranscribe, or init_backends(dir) for a \
283                 custom provider directory."
284            )));
285        }
286    }
287    check(status, context)
288}
289
290/// Options for creating a session.
291#[derive(Debug, Clone, PartialEq, Eq)]
292pub struct SessionOptions {
293    /// CPU threads for ops that run on CPU; 0 = library default.
294    pub n_threads: i32,
295    /// K/V activation precision.
296    pub kv_type: crate::types::KvType,
297    /// Optional decoder context cap in tokens; 0 = model maximum.
298    pub n_ctx: i32,
299}
300
301impl Default for SessionOptions {
302    fn default() -> Self {
303        SessionOptions {
304            n_threads: 0,
305            kv_type: crate::types::KvType::Auto,
306            n_ctx: 0,
307        }
308    }
309}
310
311/// Per-session effective limits.
312#[derive(Debug, Clone, Copy, PartialEq, Eq)]
313pub struct SessionLimits {
314    /// Decoder context cap in force for this session (0 = unbounded family).
315    pub effective_n_ctx: i32,
316    /// Longest accepted audio in ms (0 = no practical limit).
317    pub effective_max_audio_ms: i64,
318    /// Worst-case decoder KV bytes for one utterance (0 = no decoder KV).
319    pub max_kv_bytes: i64,
320}
321
322/// Convert a filesystem path to the bytes the C API's `const char *` expects.
323/// On Unix paths are arbitrary bytes, passed through verbatim; elsewhere the
324/// API takes UTF-8, so a non-UTF-8 path is rejected rather than lossily mangled
325/// (`to_string_lossy` would silently substitute U+FFFD). Shared with
326/// [`crate::backend::init_backends`].
327#[cfg(unix)]
328pub(crate) fn path_bytes(path: &Path) -> Result<Vec<u8>> {
329    use std::os::unix::ffi::OsStrExt;
330    Ok(path.as_os_str().as_bytes().to_vec())
331}
332
333#[cfg(not(unix))]
334pub(crate) fn path_bytes(path: &Path) -> Result<Vec<u8>> {
335    // The C API takes a UTF-8 `const char *`; require a UTF-8-representable path.
336    path.to_str().map(|s| s.as_bytes().to_vec()).ok_or_else(|| {
337        crate::error::Error::InvalidArgument(format!("non-UTF-8 path: {}", path.display()))
338    })
339}
transcribe_cpp/model.rs

transcribe_cpp/
model.rs
