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

transcribe_cpp/
model.rs
