1use std::ffi::CString;
12use std::os::raw::c_void;
13use std::sync::atomic::AtomicBool;
14use std::sync::Arc;
15
16use transcribe_cpp_sys as sys;
17
18use crate::cancel::{abort_trampoline, CancelToken};
19use crate::error::{check, error_for_status, Error, Result};
20use crate::family::{RunExtRaw, RunExtension};
21use crate::model::{Model, ModelInner, SessionLimits, SessionOptions};
22use crate::result::{
23 owned_opt_str, owned_str, Segment, SpeakerSegment, Timings, Token, Transcript, Word,
24};
25use crate::streaming::{build_stream_params, StreamOptions, StreamText, StreamUpdate};
26use crate::types::{Diarize, Itn, Pnc, StreamState, Task, TimestampKind};
27
28#[derive(Debug, Clone, PartialEq)]
30pub struct RunOptions {
31 pub task: Task,
32 pub timestamps: TimestampKind,
33 pub pnc: Pnc,
34 pub itn: Itn,
35 pub diarize: Diarize,
36 pub language: Option<String>,
38 pub target_language: Option<String>,
40 pub keep_special_tags: bool,
42 pub spec_k_drafts: i32,
44 pub family: Option<RunExtension>,
46}
47
48impl Default for RunOptions {
49 fn default() -> Self {
50 RunOptions {
51 task: Task::Transcribe,
52 timestamps: TimestampKind::Auto,
53 pnc: Pnc::Default,
54 itn: Itn::Default,
55 diarize: Diarize::Default,
56 language: None,
57 target_language: None,
58 keep_special_tags: false,
59 spec_k_drafts: -1,
60 family: None,
61 }
62 }
63}
64
65pub struct Session {
67 ptr: *mut sys::transcribe_session,
68 model: Arc<ModelInner>,
71 cancel: Option<Arc<AtomicBool>>,
74}
75
76impl std::fmt::Debug for Session {
77 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78 f.debug_struct("Session").finish_non_exhaustive()
79 }
80}
81
82unsafe impl Send for Session {}
86
87impl Drop for Session {
88 fn drop(&mut self) {
89 unsafe { sys::transcribe_session_free(self.ptr) };
90 }
92}
93
94impl Session {
95 pub(crate) fn new(model: &Model, options: &SessionOptions) -> Result<Session> {
96 let mut params: sys::transcribe_session_params = unsafe { std::mem::zeroed() };
97 unsafe { sys::transcribe_session_params_init(&mut params) };
98 params.n_threads = options.n_threads;
99 params.kv_type = options.kv_type.to_raw();
100 params.n_ctx = options.n_ctx;
101
102 let mut out: *mut sys::transcribe_session = std::ptr::null_mut();
103 let status = unsafe { sys::transcribe_session_init(model.inner.ptr, ¶ms, &mut out) };
104 check(status, "session init")?;
105 debug_assert!(!out.is_null());
106
107 Ok(Session {
108 ptr: out,
109 model: Arc::clone(&model.inner),
110 cancel: None,
111 })
112 }
113
114 pub fn set_cancel_token(&mut self, token: &CancelToken) {
120 let flag = Arc::clone(&token.flag);
121 let userdata = Arc::as_ptr(&flag) as *mut c_void;
122 unsafe { sys::transcribe_set_abort_callback(self.ptr, Some(abort_trampoline), userdata) };
126 self.cancel = Some(flag);
127 }
128
129 pub fn clear_cancel_token(&mut self) {
131 unsafe { sys::transcribe_set_abort_callback(self.ptr, None, std::ptr::null_mut()) };
132 self.cancel = None;
133 }
134
135 pub fn was_aborted(&self) -> bool {
137 unsafe { sys::transcribe_was_aborted(self.ptr) }
138 }
139
140 pub fn was_truncated(&self) -> bool {
143 unsafe { sys::transcribe_was_truncated(self.ptr) }
144 }
145
146 pub fn run(&mut self, pcm: &[f32], options: &RunOptions) -> Result<Transcript> {
151 let (params, _lang, _target, _family) = build_run_params(options)?;
152 let n = clamp_len(pcm.len())?;
153
154 let status = {
158 let guard = self
159 .model
160 .compute_lock
161 .lock()
162 .unwrap_or_else(|e| e.into_inner());
163 if *guard {
164 return Err(Error::Busy(
165 "a stream is active on this model; finish or drop it before run()".into(),
166 ));
167 }
168 unsafe { sys::transcribe_run(self.ptr, pcm.as_ptr(), n, ¶ms) }
169 };
170
171 match status {
172 s if s == sys::transcribe_status::TRANSCRIBE_OK => Ok(self.materialize_run()),
173 s if s == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
174 || s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED =>
175 {
176 let partial = Box::new(self.materialize_run());
178 Err(attach_partial(error_for_status(s, "run"), partial))
179 }
180 s => Err(error_for_status(s, "run")),
181 }
182 }
183
184 pub fn run_batch(
190 &mut self,
191 pcms: &[&[f32]],
192 options: &RunOptions,
193 ) -> Result<Vec<Result<Transcript>>> {
194 let (params, _lang, _target, _family) = build_run_params(options)?;
195 let ptrs: Vec<*const f32> = pcms.iter().map(|p| p.as_ptr()).collect();
196 let lens: Vec<i32> = pcms
197 .iter()
198 .map(|p| clamp_len(p.len()))
199 .collect::<Result<_>>()?;
200 let n = clamp_len(pcms.len())?;
201
202 let status = {
203 let guard = self
204 .model
205 .compute_lock
206 .lock()
207 .unwrap_or_else(|e| e.into_inner());
208 if *guard {
209 return Err(Error::Busy(
210 "a stream is active on this model; finish or drop it before run_batch()".into(),
211 ));
212 }
213 unsafe { sys::transcribe_run_batch(self.ptr, ptrs.as_ptr(), lens.as_ptr(), n, ¶ms) }
214 };
215
216 let ok = status == sys::transcribe_status::TRANSCRIBE_OK
219 || status == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED;
220 if !ok {
221 return Err(error_for_status(status, "run_batch"));
222 }
223
224 let count = unsafe { sys::transcribe_batch_n_results(self.ptr) };
225 let mut results = Vec::with_capacity(count.max(0) as usize);
226 for i in 0..count {
227 let st = unsafe { sys::transcribe_batch_status(self.ptr, i) };
228 if st == sys::transcribe_status::TRANSCRIBE_OK {
229 results.push(Ok(self.materialize_batch(i)));
230 } else if st == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
231 || st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
232 {
233 let partial = Box::new(self.materialize_batch(i));
234 let err = attach_partial(error_for_status(st, "run_batch utterance"), partial);
235 results.push(Err(err));
236 } else {
237 results.push(Err(error_for_status(st, "run_batch utterance")));
238 }
239 }
240 Ok(results)
241 }
242
243 pub fn limits(&self) -> Result<SessionLimits> {
245 let mut out: sys::transcribe_session_limits = unsafe { std::mem::zeroed() };
246 unsafe { sys::transcribe_session_limits_init(&mut out) };
247 let status = unsafe { sys::transcribe_session_get_limits(self.ptr, &mut out) };
248 check(status, "session limits")?;
249 Ok(SessionLimits {
250 effective_n_ctx: out.effective_n_ctx,
251 effective_max_audio_ms: out.effective_max_audio_ms,
252 max_kv_bytes: out.max_kv_bytes,
253 })
254 }
255
256 pub fn model(&self) -> Model {
258 Model {
259 inner: Arc::clone(&self.model),
260 }
261 }
262
263 pub fn stream(&mut self, run: &RunOptions, stream: &StreamOptions) -> Result<Stream<'_>> {
270 let (run_params, _lang, _target, _family) = build_run_params(run)?;
271 let (stream_params, _stream_family) = build_stream_params(stream);
272 {
273 let mut guard = self
277 .model
278 .compute_lock
279 .lock()
280 .unwrap_or_else(|e| e.into_inner());
281 if *guard {
282 return Err(Error::Busy(
283 "a stream is already active on this model".into(),
284 ));
285 }
286 check(
287 unsafe { sys::transcribe_stream_begin(self.ptr, &run_params, &stream_params) },
288 "stream begin",
289 )?;
290 *guard = true; }
292 Ok(Stream {
293 session: self,
294 holds_lease: true,
295 })
296 }
297
298 fn materialize_run(&self) -> Transcript {
301 let n_seg = unsafe { sys::transcribe_n_segments(self.ptr) };
302 let n_word = unsafe { sys::transcribe_n_words(self.ptr) };
303 let n_tok = unsafe { sys::transcribe_n_tokens(self.ptr) };
304 let n_speaker = unsafe { sys::transcribe_n_speaker_segments(self.ptr) };
305
306 Transcript {
307 text: owned_str(unsafe { sys::transcribe_full_text(self.ptr) }),
308 raw_text: owned_str(unsafe { sys::transcribe_raw_text(self.ptr) }),
309 language: owned_opt_str(unsafe { sys::transcribe_detected_language(self.ptr) }),
310 timestamp_kind: TimestampKind::from_raw(unsafe {
311 sys::transcribe_returned_timestamp_kind(self.ptr)
312 }),
313 segments: (0..n_seg).map(|i| self.segment(i)).collect(),
314 speaker_segments: (0..n_speaker).map(|i| self.speaker_segment(i)).collect(),
315 words: (0..n_word).map(|i| self.word(i)).collect(),
316 tokens: (0..n_tok).map(|i| self.token(i)).collect(),
317 timings: self.timings(),
318 }
319 }
320
321 fn segment(&self, i: i32) -> Segment {
322 let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
323 unsafe { sys::transcribe_segment_init(&mut raw) };
324 let _ = unsafe { sys::transcribe_get_segment(self.ptr, i, &mut raw) };
325 Segment::from_raw(&raw)
326 }
327
328 fn speaker_segment(&self, i: i32) -> SpeakerSegment {
329 let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
330 unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
331 let _ = unsafe { sys::transcribe_get_speaker_segment(self.ptr, i, &mut raw) };
332 SpeakerSegment::from_raw(&raw)
333 }
334
335 fn word(&self, i: i32) -> Word {
336 let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
337 unsafe { sys::transcribe_word_init(&mut raw) };
338 let _ = unsafe { sys::transcribe_get_word(self.ptr, i, &mut raw) };
339 Word::from_raw(&raw)
340 }
341
342 fn token(&self, i: i32) -> Token {
343 let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
344 unsafe { sys::transcribe_token_init(&mut raw) };
345 let _ = unsafe { sys::transcribe_get_token(self.ptr, i, &mut raw) };
346 Token::from_raw(&raw)
347 }
348
349 fn timings(&self) -> Timings {
350 let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
351 unsafe { sys::transcribe_timings_init(&mut raw) };
352 let _ = unsafe { sys::transcribe_get_timings(self.ptr, &mut raw) };
353 Timings::from_raw(&raw)
354 }
355
356 fn materialize_batch(&self, i: i32) -> Transcript {
359 let n_seg = unsafe { sys::transcribe_batch_n_segments(self.ptr, i) };
360 let n_word = unsafe { sys::transcribe_batch_n_words(self.ptr, i) };
361 let n_tok = unsafe { sys::transcribe_batch_n_tokens(self.ptr, i) };
362 let n_speaker = unsafe { sys::transcribe_batch_n_speaker_segments(self.ptr, i) };
363
364 Transcript {
365 text: owned_str(unsafe { sys::transcribe_batch_full_text(self.ptr, i) }),
366 raw_text: owned_str(unsafe { sys::transcribe_batch_raw_text(self.ptr, i) }),
367 language: owned_opt_str(unsafe {
368 sys::transcribe_batch_detected_language(self.ptr, i)
369 }),
370 timestamp_kind: TimestampKind::from_raw(unsafe {
371 sys::transcribe_batch_returned_timestamp_kind(self.ptr, i)
372 }),
373 segments: (0..n_seg).map(|j| self.batch_segment(i, j)).collect(),
374 speaker_segments: (0..n_speaker)
375 .map(|j| self.batch_speaker_segment(i, j))
376 .collect(),
377 words: (0..n_word).map(|j| self.batch_word(i, j)).collect(),
378 tokens: (0..n_tok).map(|j| self.batch_token(i, j)).collect(),
379 timings: self.batch_timings(i),
380 }
381 }
382
383 fn batch_segment(&self, i: i32, j: i32) -> Segment {
384 let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
385 unsafe { sys::transcribe_segment_init(&mut raw) };
386 let _ = unsafe { sys::transcribe_batch_get_segment(self.ptr, i, j, &mut raw) };
387 Segment::from_raw(&raw)
388 }
389
390 fn batch_speaker_segment(&self, i: i32, j: i32) -> SpeakerSegment {
391 let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
392 unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
393 let _ = unsafe { sys::transcribe_batch_get_speaker_segment(self.ptr, i, j, &mut raw) };
394 SpeakerSegment::from_raw(&raw)
395 }
396
397 fn batch_word(&self, i: i32, j: i32) -> Word {
398 let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
399 unsafe { sys::transcribe_word_init(&mut raw) };
400 let _ = unsafe { sys::transcribe_batch_get_word(self.ptr, i, j, &mut raw) };
401 Word::from_raw(&raw)
402 }
403
404 fn batch_token(&self, i: i32, j: i32) -> Token {
405 let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
406 unsafe { sys::transcribe_token_init(&mut raw) };
407 let _ = unsafe { sys::transcribe_batch_get_token(self.ptr, i, j, &mut raw) };
408 Token::from_raw(&raw)
409 }
410
411 fn batch_timings(&self, i: i32) -> Timings {
412 let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
413 unsafe { sys::transcribe_timings_init(&mut raw) };
414 let _ = unsafe { sys::transcribe_batch_get_timings(self.ptr, i, &mut raw) };
415 Timings::from_raw(&raw)
416 }
417}
418
419type RunParamsBundle = (
422 sys::transcribe_run_params,
423 Option<CString>,
424 Option<CString>,
425 Option<RunExtRaw>,
426);
427
428fn build_run_params(o: &RunOptions) -> Result<RunParamsBundle> {
432 let mut params: sys::transcribe_run_params = unsafe { std::mem::zeroed() };
433 unsafe { sys::transcribe_run_params_init(&mut params) };
434
435 params.task = o.task.to_raw();
436 params.timestamps = o.timestamps.to_raw();
437 params.pnc = o.pnc.to_raw();
438 params.itn = o.itn.to_raw();
439 params.diarize = o.diarize.to_raw();
440 params.keep_special_tags = o.keep_special_tags;
441 params.spec_k_drafts = o.spec_k_drafts;
442
443 let lang = o.language.as_deref().map(CString::new).transpose()?;
444 let target = o.target_language.as_deref().map(CString::new).transpose()?;
445 params.language = lang.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
446 params.target_language = target.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
447
448 let family = o
449 .family
450 .as_ref()
451 .map(RunExtension::materialize)
452 .transpose()?;
453 params.family = family.as_ref().map_or(std::ptr::null(), |f| f.ext_ptr());
454
455 Ok((params, lang, target, family))
456}
457
458fn clamp_len(len: usize) -> Result<i32> {
460 i32::try_from(len).map_err(|_| Error::InvalidArgument(format!("length {len} exceeds i32::MAX")))
461}
462
463fn attach_partial(err: Error, partial: Box<Transcript>) -> Error {
466 match err {
467 Error::Aborted { message, .. } => Error::Aborted {
468 message,
469 partial: Some(partial),
470 },
471 Error::OutputTruncated { message, .. } => Error::OutputTruncated {
472 message,
473 partial: Some(partial),
474 },
475 other => other,
476 }
477}
478
479pub struct Stream<'a> {
486 session: &'a mut Session,
487 holds_lease: bool,
492}
493
494impl std::fmt::Debug for Stream<'_> {
495 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
496 f.debug_struct("Stream")
497 .field("state", &self.state())
498 .field("revision", &self.revision())
499 .finish()
500 }
501}
502
503impl Drop for Stream<'_> {
504 fn drop(&mut self) {
505 let mut guard = self
510 .session
511 .model
512 .compute_lock
513 .lock()
514 .unwrap_or_else(|e| e.into_inner());
515 unsafe { sys::transcribe_stream_reset(self.session.ptr) };
516 if self.holds_lease {
517 *guard = false;
518 }
519 }
520}
521
522impl Stream<'_> {
523 pub fn feed(&mut self, pcm: &[f32]) -> Result<StreamUpdate> {
525 let n = clamp_len(pcm.len())?;
526 let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
527 unsafe { sys::transcribe_stream_update_init(&mut update) };
528 let status = {
529 let _guard = self
532 .session
533 .model
534 .compute_lock
535 .lock()
536 .unwrap_or_else(|e| e.into_inner());
537 unsafe { sys::transcribe_stream_feed(self.session.ptr, pcm.as_ptr(), n, &mut update) }
538 };
539 check(status, "stream feed")?;
540 Ok(StreamUpdate::from_raw(&update))
541 }
542
543 pub fn finalize(&mut self) -> Result<StreamUpdate> {
546 let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
547 unsafe { sys::transcribe_stream_update_init(&mut update) };
548 let status = {
549 let mut guard = self
550 .session
551 .model
552 .compute_lock
553 .lock()
554 .unwrap_or_else(|e| e.into_inner());
555 let st = unsafe { sys::transcribe_stream_finalize(self.session.ptr, &mut update) };
556 if self.holds_lease {
560 *guard = false;
561 }
562 st
563 };
564 self.holds_lease = false;
565 check(status, "stream finalize")?;
566 Ok(StreamUpdate::from_raw(&update))
567 }
568
569 pub fn reset(&mut self) {
574 let mut guard = self
575 .session
576 .model
577 .compute_lock
578 .lock()
579 .unwrap_or_else(|e| e.into_inner());
580 unsafe { sys::transcribe_stream_reset(self.session.ptr) };
581 if self.holds_lease {
582 *guard = false;
583 }
584 self.holds_lease = false;
585 }
586
587 pub fn text(&self) -> StreamText {
589 let mut raw: sys::transcribe_stream_text = unsafe { std::mem::zeroed() };
590 unsafe { sys::transcribe_stream_text_init(&mut raw) };
591 let _ = unsafe { sys::transcribe_stream_get_text(self.session.ptr, &mut raw) };
592 StreamText::from_raw(&raw)
593 }
594
595 pub fn snapshot(&self) -> Transcript {
598 self.session.materialize_run()
599 }
600
601 pub fn state(&self) -> StreamState {
603 StreamState::from_raw(unsafe { sys::transcribe_stream_get_state(self.session.ptr) })
604 }
605
606 pub fn revision(&self) -> i32 {
608 unsafe { sys::transcribe_stream_revision(self.session.ptr) }
609 }
610
611 pub fn last_status(&self) -> Option<Error> {
617 let status = unsafe { sys::transcribe_stream_last_status(self.session.ptr) };
618 if status == sys::transcribe_status::TRANSCRIBE_OK {
619 None
620 } else {
621 Some(error_for_status(status, "stream"))
622 }
623 }
624
625 pub fn was_aborted(&self) -> bool {
627 self.session.was_aborted()
628 }
629}