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

transcribe_cpp/
model.rs
