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)]
30#[cfg_attr(
31 feature = "serde",
32 derive(serde::Serialize, serde::Deserialize),
33 serde(default)
34)]
35pub struct RunOptions {
36 pub task: Task,
37 pub timestamps: TimestampKind,
38 pub pnc: Pnc,
39 pub itn: Itn,
40 pub diarize: Diarize,
41 pub language: Option<String>,
43 pub target_language: Option<String>,
45 pub keep_special_tags: bool,
47 pub spec_k_drafts: i32,
49 pub family: Option<RunExtension>,
51 pub vocabulary: Vec<String>,
54 pub prompt: Option<String>,
58 pub prefix: Option<String>,
62}
63
64impl Default for RunOptions {
65 fn default() -> Self {
66 RunOptions {
67 task: Task::Transcribe,
68 timestamps: TimestampKind::Auto,
69 pnc: Pnc::Default,
70 itn: Itn::Default,
71 diarize: Diarize::Default,
72 language: None,
73 target_language: None,
74 keep_special_tags: false,
75 spec_k_drafts: -1,
76 family: None,
77 vocabulary: Vec::new(),
78 prompt: None,
79 prefix: None,
80 }
81 }
82}
83
84pub struct Session {
86 ptr: *mut sys::transcribe_session,
87 model: Arc<ModelInner>,
90 cancel: Option<Arc<AtomicBool>>,
93}
94
95impl std::fmt::Debug for Session {
96 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97 f.debug_struct("Session").finish_non_exhaustive()
98 }
99}
100
101unsafe impl Send for Session {}
105
106impl Drop for Session {
107 fn drop(&mut self) {
108 unsafe { sys::transcribe_session_free(self.ptr) };
109 }
111}
112
113impl Session {
114 pub(crate) fn new(model: &Model, options: &SessionOptions) -> Result<Session> {
115 let mut params: sys::transcribe_session_params = unsafe { std::mem::zeroed() };
116 unsafe { sys::transcribe_session_params_init(&mut params) };
117 params.n_threads = options.n_threads;
118 params.kv_type = options.kv_type.to_raw();
119 params.n_ctx = options.n_ctx;
120
121 let mut out: *mut sys::transcribe_session = std::ptr::null_mut();
122 let status = unsafe { sys::transcribe_session_init(model.inner.ptr, ¶ms, &mut out) };
123 check(status, "session init")?;
124 debug_assert!(!out.is_null());
125
126 Ok(Session {
127 ptr: out,
128 model: Arc::clone(&model.inner),
129 cancel: None,
130 })
131 }
132
133 pub fn set_cancel_token(&mut self, token: &CancelToken) {
139 let flag = Arc::clone(&token.flag);
140 let userdata = Arc::as_ptr(&flag) as *mut c_void;
141 unsafe { sys::transcribe_set_abort_callback(self.ptr, Some(abort_trampoline), userdata) };
145 self.cancel = Some(flag);
146 }
147
148 pub fn clear_cancel_token(&mut self) {
150 unsafe { sys::transcribe_set_abort_callback(self.ptr, None, std::ptr::null_mut()) };
151 self.cancel = None;
152 }
153
154 pub fn was_aborted(&self) -> bool {
156 unsafe { sys::transcribe_was_aborted(self.ptr) }
157 }
158
159 pub fn was_truncated(&self) -> bool {
163 unsafe { sys::transcribe_was_truncated(self.ptr) }
164 }
165
166 pub fn run(&mut self, pcm: &[f32], options: &RunOptions) -> Result<Transcript> {
172 let (params, _lang, _target, _family, _prompting) = build_run_params(options)?;
173 let n = clamp_len(pcm.len())?;
174
175 let status = {
179 let guard = self
180 .model
181 .compute_lock
182 .lock()
183 .unwrap_or_else(|e| e.into_inner());
184 if *guard {
185 return Err(Error::Busy(
186 "a stream is active on this model; finish or drop it before run()".into(),
187 ));
188 }
189 unsafe { sys::transcribe_run(self.ptr, pcm.as_ptr(), n, ¶ms) }
190 };
191
192 match status {
193 s if s == sys::transcribe_status::TRANSCRIBE_OK => Ok(self.materialize_run()),
194 s if s == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
195 || s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
196 || s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_REPETITION =>
197 {
198 let partial = Box::new(self.materialize_run());
200 Err(attach_partial(error_for_status(s, "run"), partial))
201 }
202 s => Err(error_for_status(s, "run")),
203 }
204 }
205
206 pub fn run_batch(
212 &mut self,
213 pcms: &[&[f32]],
214 options: &RunOptions,
215 ) -> Result<Vec<Result<Transcript>>> {
216 let (params, _lang, _target, _family, _prompting) = build_run_params(options)?;
217 let ptrs: Vec<*const f32> = pcms.iter().map(|p| p.as_ptr()).collect();
218 let lens: Vec<i32> = pcms
219 .iter()
220 .map(|p| clamp_len(p.len()))
221 .collect::<Result<_>>()?;
222 let n = clamp_len(pcms.len())?;
223
224 let status = {
225 let guard = self
226 .model
227 .compute_lock
228 .lock()
229 .unwrap_or_else(|e| e.into_inner());
230 if *guard {
231 return Err(Error::Busy(
232 "a stream is active on this model; finish or drop it before run_batch()".into(),
233 ));
234 }
235 unsafe { sys::transcribe_run_batch(self.ptr, ptrs.as_ptr(), lens.as_ptr(), n, ¶ms) }
236 };
237
238 let ok = status == sys::transcribe_status::TRANSCRIBE_OK
241 || status == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED;
242 if !ok {
243 return Err(error_for_status(status, "run_batch"));
244 }
245
246 let count = unsafe { sys::transcribe_batch_n_results(self.ptr) };
247 let mut results = Vec::with_capacity(count.max(0) as usize);
248 for i in 0..count {
249 let st = unsafe { sys::transcribe_batch_status(self.ptr, i) };
250 if st == sys::transcribe_status::TRANSCRIBE_OK {
251 results.push(Ok(self.materialize_batch(i)));
252 } else if st == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
253 || st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
254 || st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_REPETITION
255 {
256 let partial = Box::new(self.materialize_batch(i));
257 let err = attach_partial(error_for_status(st, "run_batch utterance"), partial);
258 results.push(Err(err));
259 } else {
260 results.push(Err(error_for_status(st, "run_batch utterance")));
261 }
262 }
263 Ok(results)
264 }
265
266 pub fn limits(&self) -> Result<SessionLimits> {
268 let mut out: sys::transcribe_session_limits = unsafe { std::mem::zeroed() };
269 unsafe { sys::transcribe_session_limits_init(&mut out) };
270 let status = unsafe { sys::transcribe_session_get_limits(self.ptr, &mut out) };
271 check(status, "session limits")?;
272 Ok(SessionLimits {
273 effective_n_ctx: out.effective_n_ctx,
274 effective_max_audio_ms: out.effective_max_audio_ms,
275 max_kv_bytes: out.max_kv_bytes,
276 })
277 }
278
279 pub fn model(&self) -> Model {
281 Model {
282 inner: Arc::clone(&self.model),
283 }
284 }
285
286 pub fn stream(&mut self, run: &RunOptions, stream: &StreamOptions) -> Result<Stream<'_>> {
294 let (run_params, _lang, _target, _family, _prompting) = build_run_params(run)?;
295 let (stream_params, _stream_family) = build_stream_params(stream);
296 {
297 let mut guard = self
301 .model
302 .compute_lock
303 .lock()
304 .unwrap_or_else(|e| e.into_inner());
305 if *guard {
306 return Err(Error::Busy(
307 "a stream is already active on this model".into(),
308 ));
309 }
310 check(
311 unsafe { sys::transcribe_stream_begin(self.ptr, &run_params, &stream_params) },
312 "stream begin",
313 )?;
314 *guard = true; }
316 Ok(Stream {
317 session: self,
318 holds_lease: true,
319 })
320 }
321
322 fn materialize_run(&self) -> Transcript {
325 let n_seg = unsafe { sys::transcribe_n_segments(self.ptr) };
326 let n_word = unsafe { sys::transcribe_n_words(self.ptr) };
327 let n_tok = unsafe { sys::transcribe_n_tokens(self.ptr) };
328 let n_speaker = unsafe { sys::transcribe_n_speaker_segments(self.ptr) };
329
330 Transcript {
331 text: owned_str(unsafe { sys::transcribe_full_text(self.ptr) }),
332 raw_text: owned_str(unsafe { sys::transcribe_raw_text(self.ptr) }),
333 language: owned_opt_str(unsafe { sys::transcribe_detected_language(self.ptr) }),
334 timestamp_kind: TimestampKind::from_raw(unsafe {
335 sys::transcribe_returned_timestamp_kind(self.ptr)
336 }),
337 segments: (0..n_seg).map(|i| self.segment(i)).collect(),
338 speaker_segments: (0..n_speaker).map(|i| self.speaker_segment(i)).collect(),
339 words: (0..n_word).map(|i| self.word(i)).collect(),
340 tokens: (0..n_tok).map(|i| self.token(i)).collect(),
341 timings: self.timings(),
342 }
343 }
344
345 fn segment(&self, i: i32) -> Segment {
346 let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
347 unsafe { sys::transcribe_segment_init(&mut raw) };
348 let _ = unsafe { sys::transcribe_get_segment(self.ptr, i, &mut raw) };
349 Segment::from_raw(&raw)
350 }
351
352 fn speaker_segment(&self, i: i32) -> SpeakerSegment {
353 let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
354 unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
355 let _ = unsafe { sys::transcribe_get_speaker_segment(self.ptr, i, &mut raw) };
356 SpeakerSegment::from_raw(&raw)
357 }
358
359 fn word(&self, i: i32) -> Word {
360 let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
361 unsafe { sys::transcribe_word_init(&mut raw) };
362 let _ = unsafe { sys::transcribe_get_word(self.ptr, i, &mut raw) };
363 Word::from_raw(&raw)
364 }
365
366 fn token(&self, i: i32) -> Token {
367 let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
368 unsafe { sys::transcribe_token_init(&mut raw) };
369 let _ = unsafe { sys::transcribe_get_token(self.ptr, i, &mut raw) };
370 Token::from_raw(&raw)
371 }
372
373 fn timings(&self) -> Timings {
374 let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
375 unsafe { sys::transcribe_timings_init(&mut raw) };
376 let _ = unsafe { sys::transcribe_get_timings(self.ptr, &mut raw) };
377 Timings::from_raw(&raw)
378 }
379
380 fn materialize_batch(&self, i: i32) -> Transcript {
383 let n_seg = unsafe { sys::transcribe_batch_n_segments(self.ptr, i) };
384 let n_word = unsafe { sys::transcribe_batch_n_words(self.ptr, i) };
385 let n_tok = unsafe { sys::transcribe_batch_n_tokens(self.ptr, i) };
386 let n_speaker = unsafe { sys::transcribe_batch_n_speaker_segments(self.ptr, i) };
387
388 Transcript {
389 text: owned_str(unsafe { sys::transcribe_batch_full_text(self.ptr, i) }),
390 raw_text: owned_str(unsafe { sys::transcribe_batch_raw_text(self.ptr, i) }),
391 language: owned_opt_str(unsafe {
392 sys::transcribe_batch_detected_language(self.ptr, i)
393 }),
394 timestamp_kind: TimestampKind::from_raw(unsafe {
395 sys::transcribe_batch_returned_timestamp_kind(self.ptr, i)
396 }),
397 segments: (0..n_seg).map(|j| self.batch_segment(i, j)).collect(),
398 speaker_segments: (0..n_speaker)
399 .map(|j| self.batch_speaker_segment(i, j))
400 .collect(),
401 words: (0..n_word).map(|j| self.batch_word(i, j)).collect(),
402 tokens: (0..n_tok).map(|j| self.batch_token(i, j)).collect(),
403 timings: self.batch_timings(i),
404 }
405 }
406
407 fn batch_segment(&self, i: i32, j: i32) -> Segment {
408 let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
409 unsafe { sys::transcribe_segment_init(&mut raw) };
410 let _ = unsafe { sys::transcribe_batch_get_segment(self.ptr, i, j, &mut raw) };
411 Segment::from_raw(&raw)
412 }
413
414 fn batch_speaker_segment(&self, i: i32, j: i32) -> SpeakerSegment {
415 let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
416 unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
417 let _ = unsafe { sys::transcribe_batch_get_speaker_segment(self.ptr, i, j, &mut raw) };
418 SpeakerSegment::from_raw(&raw)
419 }
420
421 fn batch_word(&self, i: i32, j: i32) -> Word {
422 let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
423 unsafe { sys::transcribe_word_init(&mut raw) };
424 let _ = unsafe { sys::transcribe_batch_get_word(self.ptr, i, j, &mut raw) };
425 Word::from_raw(&raw)
426 }
427
428 fn batch_token(&self, i: i32, j: i32) -> Token {
429 let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
430 unsafe { sys::transcribe_token_init(&mut raw) };
431 let _ = unsafe { sys::transcribe_batch_get_token(self.ptr, i, j, &mut raw) };
432 Token::from_raw(&raw)
433 }
434
435 fn batch_timings(&self, i: i32) -> Timings {
436 let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
437 unsafe { sys::transcribe_timings_init(&mut raw) };
438 let _ = unsafe { sys::transcribe_batch_get_timings(self.ptr, i, &mut raw) };
439 Timings::from_raw(&raw)
440 }
441}
442
443type RunParamsBundle = (
446 sys::transcribe_run_params,
447 Option<CString>,
448 Option<CString>,
449 Option<RunExtRaw>,
450 PromptingKeepalive,
451);
452
453struct PromptingKeepalive {
455 _terms: Vec<CString>,
456 _term_ptrs: Vec<*const std::os::raw::c_char>,
457 _prompt: Option<CString>,
458 _prefix: Option<CString>,
459}
460
461fn build_run_params(o: &RunOptions) -> Result<RunParamsBundle> {
465 let mut params: sys::transcribe_run_params = unsafe { std::mem::zeroed() };
466 unsafe { sys::transcribe_run_params_init(&mut params) };
467
468 params.task = o.task.to_raw();
469 params.timestamps = o.timestamps.to_raw();
470 params.pnc = o.pnc.to_raw();
471 params.itn = o.itn.to_raw();
472 params.diarize = o.diarize.to_raw();
473 params.keep_special_tags = o.keep_special_tags;
474 params.spec_k_drafts = o.spec_k_drafts;
475
476 let lang = o.language.as_deref().map(CString::new).transpose()?;
477 let target = o.target_language.as_deref().map(CString::new).transpose()?;
478 params.language = lang.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
479 params.target_language = target.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
480
481 let family = o
482 .family
483 .as_ref()
484 .map(RunExtension::materialize)
485 .transpose()?;
486 params.family = family.as_ref().map_or(std::ptr::null(), |f| f.ext_ptr());
487
488 let terms = o
489 .vocabulary
490 .iter()
491 .map(|t| CString::new(t.as_str()))
492 .collect::<std::result::Result<Vec<_>, _>>()?;
493 let term_ptrs: Vec<*const std::os::raw::c_char> = terms.iter().map(|c| c.as_ptr()).collect();
494 let prompt = o.prompt.as_deref().map(CString::new).transpose()?;
495 let prefix = o.prefix.as_deref().map(CString::new).transpose()?;
496 params.vocabulary = if term_ptrs.is_empty() {
497 std::ptr::null()
498 } else {
499 term_ptrs.as_ptr()
500 };
501 params.n_vocabulary = clamp_len(term_ptrs.len())?;
502 params.prompt = prompt.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
503 params.prefix = prefix.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
504 let prompting = PromptingKeepalive {
505 _terms: terms,
506 _term_ptrs: term_ptrs,
507 _prompt: prompt,
508 _prefix: prefix,
509 };
510
511 Ok((params, lang, target, family, prompting))
512}
513
514fn clamp_len(len: usize) -> Result<i32> {
516 i32::try_from(len).map_err(|_| Error::InvalidArgument(format!("length {len} exceeds i32::MAX")))
517}
518
519fn attach_partial(err: Error, partial: Box<Transcript>) -> Error {
522 match err {
523 Error::Aborted { message, .. } => Error::Aborted {
524 message,
525 partial: Some(partial),
526 },
527 Error::OutputTruncated { message, .. } => Error::OutputTruncated {
528 message,
529 partial: Some(partial),
530 },
531 Error::OutputRepetition { message, .. } => Error::OutputRepetition {
532 message,
533 partial: Some(partial),
534 },
535 other => other,
536 }
537}
538
539pub struct Stream<'a> {
546 session: &'a mut Session,
547 holds_lease: bool,
552}
553
554impl std::fmt::Debug for Stream<'_> {
555 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
556 f.debug_struct("Stream")
557 .field("state", &self.state())
558 .field("revision", &self.revision())
559 .finish()
560 }
561}
562
563impl Drop for Stream<'_> {
564 fn drop(&mut self) {
565 let mut guard = self
570 .session
571 .model
572 .compute_lock
573 .lock()
574 .unwrap_or_else(|e| e.into_inner());
575 unsafe { sys::transcribe_stream_reset(self.session.ptr) };
576 if self.holds_lease {
577 *guard = false;
578 }
579 }
580}
581
582impl Stream<'_> {
583 pub fn feed(&mut self, pcm: &[f32]) -> Result<StreamUpdate> {
585 let n = clamp_len(pcm.len())?;
586 let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
587 unsafe { sys::transcribe_stream_update_init(&mut update) };
588 let status = {
589 let _guard = self
592 .session
593 .model
594 .compute_lock
595 .lock()
596 .unwrap_or_else(|e| e.into_inner());
597 unsafe { sys::transcribe_stream_feed(self.session.ptr, pcm.as_ptr(), n, &mut update) }
598 };
599 check(status, "stream feed")?;
600 Ok(StreamUpdate::from_raw(&update))
601 }
602
603 pub fn finalize(&mut self) -> Result<StreamUpdate> {
606 let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
607 unsafe { sys::transcribe_stream_update_init(&mut update) };
608 let status = {
609 let mut guard = self
610 .session
611 .model
612 .compute_lock
613 .lock()
614 .unwrap_or_else(|e| e.into_inner());
615 let st = unsafe { sys::transcribe_stream_finalize(self.session.ptr, &mut update) };
616 if self.holds_lease {
620 *guard = false;
621 }
622 st
623 };
624 self.holds_lease = false;
625 check(status, "stream finalize")?;
626 Ok(StreamUpdate::from_raw(&update))
627 }
628
629 pub fn reset(&mut self) {
634 let mut guard = self
635 .session
636 .model
637 .compute_lock
638 .lock()
639 .unwrap_or_else(|e| e.into_inner());
640 unsafe { sys::transcribe_stream_reset(self.session.ptr) };
641 if self.holds_lease {
642 *guard = false;
643 }
644 self.holds_lease = false;
645 }
646
647 pub fn text(&self) -> StreamText {
649 let mut raw: sys::transcribe_stream_text = unsafe { std::mem::zeroed() };
650 unsafe { sys::transcribe_stream_text_init(&mut raw) };
651 let _ = unsafe { sys::transcribe_stream_get_text(self.session.ptr, &mut raw) };
652 StreamText::from_raw(&raw)
653 }
654
655 pub fn snapshot(&self) -> Transcript {
658 self.session.materialize_run()
659 }
660
661 pub fn state(&self) -> StreamState {
663 StreamState::from_raw(unsafe { sys::transcribe_stream_get_state(self.session.ptr) })
664 }
665
666 pub fn revision(&self) -> i32 {
668 unsafe { sys::transcribe_stream_revision(self.session.ptr) }
669 }
670
671 pub fn last_status(&self) -> Option<Error> {
677 let status = unsafe { sys::transcribe_stream_last_status(self.session.ptr) };
678 if status == sys::transcribe_status::TRANSCRIBE_OK {
679 None
680 } else {
681 Some(error_for_status(status, "stream"))
682 }
683 }
684
685 pub fn was_aborted(&self) -> bool {
687 self.session.was_aborted()
688 }
689}