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

transcribe_cpp/
model.rs
