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 = self.model.with_compute(
176 Some("a stream is active on this model; finish or drop it before run()"),
177 |_| unsafe { sys::transcribe_run(self.ptr, pcm.as_ptr(), n, ¶ms) },
178 )?;
179
180 match status {
181 s if s == sys::transcribe_status::TRANSCRIBE_OK => Ok(self.materialize_run()),
182 s if s == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
183 || s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
184 || s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_REPETITION =>
185 {
186 let partial = Box::new(self.materialize_run());
188 Err(attach_partial(error_for_status(s, "run"), partial))
189 }
190 s => Err(error_for_status(s, "run")),
191 }
192 }
193
194 pub fn run_batch(
200 &mut self,
201 pcms: &[&[f32]],
202 options: &RunOptions,
203 ) -> Result<Vec<Result<Transcript>>> {
204 let (params, _lang, _target, _family, _prompting) = build_run_params(options)?;
205 let ptrs: Vec<*const f32> = pcms.iter().map(|p| p.as_ptr()).collect();
206 let lens: Vec<i32> = pcms
207 .iter()
208 .map(|p| clamp_len(p.len()))
209 .collect::<Result<_>>()?;
210 let n = clamp_len(pcms.len())?;
211
212 let status = self.model.with_compute(
213 Some("a stream is active on this model; finish or drop it before run_batch()"),
214 |_| unsafe {
215 sys::transcribe_run_batch(self.ptr, ptrs.as_ptr(), lens.as_ptr(), n, ¶ms)
216 },
217 )?;
218
219 let ok = status == sys::transcribe_status::TRANSCRIBE_OK
222 || status == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED;
223 if !ok {
224 return Err(error_for_status(status, "run_batch"));
225 }
226
227 let count = unsafe { sys::transcribe_batch_n_results(self.ptr) };
228 let mut results = Vec::with_capacity(count.max(0) as usize);
229 for i in 0..count {
230 let st = unsafe { sys::transcribe_batch_status(self.ptr, i) };
231 if st == sys::transcribe_status::TRANSCRIBE_OK {
232 results.push(Ok(self.materialize_batch(i)));
233 } else if st == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
234 || st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
235 || st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_REPETITION
236 {
237 let partial = Box::new(self.materialize_batch(i));
238 let err = attach_partial(error_for_status(st, "run_batch utterance"), partial);
239 results.push(Err(err));
240 } else {
241 results.push(Err(error_for_status(st, "run_batch utterance")));
242 }
243 }
244 Ok(results)
245 }
246
247 pub fn limits(&self) -> Result<SessionLimits> {
249 let mut out: sys::transcribe_session_limits = unsafe { std::mem::zeroed() };
250 unsafe { sys::transcribe_session_limits_init(&mut out) };
251 let status = unsafe { sys::transcribe_session_get_limits(self.ptr, &mut out) };
252 check(status, "session limits")?;
253 Ok(SessionLimits {
254 effective_n_ctx: out.effective_n_ctx,
255 effective_max_audio_ms: out.effective_max_audio_ms,
256 max_kv_bytes: out.max_kv_bytes,
257 })
258 }
259
260 pub fn model(&self) -> Model {
262 Model {
263 inner: Arc::clone(&self.model),
264 }
265 }
266
267 pub fn stream(&mut self, run: &RunOptions, stream: &StreamOptions) -> Result<Stream<'_>> {
275 let (run_params, _lang, _target, _family, _prompting) = build_run_params(run)?;
276 let (stream_params, _stream_family) = build_stream_params(stream);
277 self.model.with_compute(
281 Some("a stream is already active on this model"),
282 |lease| -> Result<()> {
283 check(
284 unsafe { sys::transcribe_stream_begin(self.ptr, &run_params, &stream_params) },
285 "stream begin",
286 )?;
287 *lease = true; Ok(())
289 },
290 )??;
291 Ok(Stream {
292 session: self,
293 holds_lease: true,
294 })
295 }
296
297 fn materialize_run(&self) -> Transcript {
300 let n_seg = unsafe { sys::transcribe_n_segments(self.ptr) };
301 let n_word = unsafe { sys::transcribe_n_words(self.ptr) };
302 let n_tok = unsafe { sys::transcribe_n_tokens(self.ptr) };
303 let n_speaker = unsafe { sys::transcribe_n_speaker_segments(self.ptr) };
304
305 Transcript {
306 text: owned_str(unsafe { sys::transcribe_full_text(self.ptr) }),
307 raw_text: owned_str(unsafe { sys::transcribe_raw_text(self.ptr) }),
308 language: owned_opt_str(unsafe { sys::transcribe_detected_language(self.ptr) }),
309 timestamp_kind: TimestampKind::from_raw(unsafe {
310 sys::transcribe_returned_timestamp_kind(self.ptr)
311 }),
312 segments: (0..n_seg).map(|i| self.segment(i)).collect(),
313 speaker_segments: (0..n_speaker).map(|i| self.speaker_segment(i)).collect(),
314 words: (0..n_word).map(|i| self.word(i)).collect(),
315 tokens: (0..n_tok).map(|i| self.token(i)).collect(),
316 timings: self.timings(),
317 }
318 }
319
320 fn segment(&self, i: i32) -> Segment {
321 let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
322 unsafe { sys::transcribe_segment_init(&mut raw) };
323 let _ = unsafe { sys::transcribe_get_segment(self.ptr, i, &mut raw) };
324 Segment::from_raw(&raw)
325 }
326
327 fn speaker_segment(&self, i: i32) -> SpeakerSegment {
328 let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
329 unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
330 let _ = unsafe { sys::transcribe_get_speaker_segment(self.ptr, i, &mut raw) };
331 SpeakerSegment::from_raw(&raw)
332 }
333
334 fn word(&self, i: i32) -> Word {
335 let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
336 unsafe { sys::transcribe_word_init(&mut raw) };
337 let _ = unsafe { sys::transcribe_get_word(self.ptr, i, &mut raw) };
338 Word::from_raw(&raw)
339 }
340
341 fn token(&self, i: i32) -> Token {
342 let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
343 unsafe { sys::transcribe_token_init(&mut raw) };
344 let _ = unsafe { sys::transcribe_get_token(self.ptr, i, &mut raw) };
345 Token::from_raw(&raw)
346 }
347
348 fn timings(&self) -> Timings {
349 let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
350 unsafe { sys::transcribe_timings_init(&mut raw) };
351 let _ = unsafe { sys::transcribe_get_timings(self.ptr, &mut raw) };
352 Timings::from_raw(&raw)
353 }
354
355 fn materialize_batch(&self, i: i32) -> Transcript {
358 let n_seg = unsafe { sys::transcribe_batch_n_segments(self.ptr, i) };
359 let n_word = unsafe { sys::transcribe_batch_n_words(self.ptr, i) };
360 let n_tok = unsafe { sys::transcribe_batch_n_tokens(self.ptr, i) };
361 let n_speaker = unsafe { sys::transcribe_batch_n_speaker_segments(self.ptr, i) };
362
363 Transcript {
364 text: owned_str(unsafe { sys::transcribe_batch_full_text(self.ptr, i) }),
365 raw_text: owned_str(unsafe { sys::transcribe_batch_raw_text(self.ptr, i) }),
366 language: owned_opt_str(unsafe {
367 sys::transcribe_batch_detected_language(self.ptr, i)
368 }),
369 timestamp_kind: TimestampKind::from_raw(unsafe {
370 sys::transcribe_batch_returned_timestamp_kind(self.ptr, i)
371 }),
372 segments: (0..n_seg).map(|j| self.batch_segment(i, j)).collect(),
373 speaker_segments: (0..n_speaker)
374 .map(|j| self.batch_speaker_segment(i, j))
375 .collect(),
376 words: (0..n_word).map(|j| self.batch_word(i, j)).collect(),
377 tokens: (0..n_tok).map(|j| self.batch_token(i, j)).collect(),
378 timings: self.batch_timings(i),
379 }
380 }
381
382 fn batch_segment(&self, i: i32, j: i32) -> Segment {
383 let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
384 unsafe { sys::transcribe_segment_init(&mut raw) };
385 let _ = unsafe { sys::transcribe_batch_get_segment(self.ptr, i, j, &mut raw) };
386 Segment::from_raw(&raw)
387 }
388
389 fn batch_speaker_segment(&self, i: i32, j: i32) -> SpeakerSegment {
390 let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
391 unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
392 let _ = unsafe { sys::transcribe_batch_get_speaker_segment(self.ptr, i, j, &mut raw) };
393 SpeakerSegment::from_raw(&raw)
394 }
395
396 fn batch_word(&self, i: i32, j: i32) -> Word {
397 let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
398 unsafe { sys::transcribe_word_init(&mut raw) };
399 let _ = unsafe { sys::transcribe_batch_get_word(self.ptr, i, j, &mut raw) };
400 Word::from_raw(&raw)
401 }
402
403 fn batch_token(&self, i: i32, j: i32) -> Token {
404 let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
405 unsafe { sys::transcribe_token_init(&mut raw) };
406 let _ = unsafe { sys::transcribe_batch_get_token(self.ptr, i, j, &mut raw) };
407 Token::from_raw(&raw)
408 }
409
410 fn batch_timings(&self, i: i32) -> Timings {
411 let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
412 unsafe { sys::transcribe_timings_init(&mut raw) };
413 let _ = unsafe { sys::transcribe_batch_get_timings(self.ptr, i, &mut raw) };
414 Timings::from_raw(&raw)
415 }
416}
417
418type RunParamsBundle = (
421 sys::transcribe_run_params,
422 Option<CString>,
423 Option<CString>,
424 Option<RunExtRaw>,
425 PromptingKeepalive,
426);
427
428struct PromptingKeepalive {
430 _terms: Vec<CString>,
431 _term_ptrs: Vec<*const std::os::raw::c_char>,
432 _prompt: Option<CString>,
433 _prefix: Option<CString>,
434}
435
436fn build_run_params(o: &RunOptions) -> Result<RunParamsBundle> {
440 let mut params: sys::transcribe_run_params = unsafe { std::mem::zeroed() };
441 unsafe { sys::transcribe_run_params_init(&mut params) };
442
443 params.task = o.task.to_raw();
444 params.timestamps = o.timestamps.to_raw();
445 params.pnc = o.pnc.to_raw();
446 params.itn = o.itn.to_raw();
447 params.diarize = o.diarize.to_raw();
448 params.keep_special_tags = o.keep_special_tags;
449 params.spec_k_drafts = o.spec_k_drafts;
450
451 let lang = o.language.as_deref().map(CString::new).transpose()?;
452 let target = o.target_language.as_deref().map(CString::new).transpose()?;
453 params.language = lang.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
454 params.target_language = target.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
455
456 let family = o
457 .family
458 .as_ref()
459 .map(RunExtension::materialize)
460 .transpose()?;
461 params.family = family.as_ref().map_or(std::ptr::null(), |f| f.ext_ptr());
462
463 let terms = o
464 .vocabulary
465 .iter()
466 .map(|t| CString::new(t.as_str()))
467 .collect::<std::result::Result<Vec<_>, _>>()?;
468 let term_ptrs: Vec<*const std::os::raw::c_char> = terms.iter().map(|c| c.as_ptr()).collect();
469 let prompt = o.prompt.as_deref().map(CString::new).transpose()?;
470 let prefix = o.prefix.as_deref().map(CString::new).transpose()?;
471 params.vocabulary = if term_ptrs.is_empty() {
472 std::ptr::null()
473 } else {
474 term_ptrs.as_ptr()
475 };
476 params.n_vocabulary = clamp_len(term_ptrs.len())?;
477 params.prompt = prompt.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
478 params.prefix = prefix.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
479 let prompting = PromptingKeepalive {
480 _terms: terms,
481 _term_ptrs: term_ptrs,
482 _prompt: prompt,
483 _prefix: prefix,
484 };
485
486 Ok((params, lang, target, family, prompting))
487}
488
489pub(crate) fn clamp_len(len: usize) -> Result<i32> {
491 i32::try_from(len).map_err(|_| Error::InvalidArgument(format!("length {len} exceeds i32::MAX")))
492}
493
494fn attach_partial(err: Error, partial: Box<Transcript>) -> Error {
497 match err {
498 Error::Aborted { message, .. } => Error::Aborted {
499 message,
500 partial: Some(partial),
501 },
502 Error::OutputTruncated { message, .. } => Error::OutputTruncated {
503 message,
504 partial: Some(partial),
505 },
506 Error::OutputRepetition { message, .. } => Error::OutputRepetition {
507 message,
508 partial: Some(partial),
509 },
510 other => other,
511 }
512}
513
514pub struct Stream<'a> {
535 session: &'a mut Session,
536 holds_lease: bool,
541}
542
543impl std::fmt::Debug for Stream<'_> {
544 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
545 f.debug_struct("Stream")
546 .field("state", &self.state())
547 .field("revision", &self.revision())
548 .finish()
549 }
550}
551
552impl Drop for Stream<'_> {
553 fn drop(&mut self) {
554 self.reset()
555 }
556}
557
558impl Stream<'_> {
559 pub fn feed(&mut self, pcm: &[f32]) -> Result<StreamUpdate> {
561 let n = clamp_len(pcm.len())?;
562 let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
563 unsafe { sys::transcribe_stream_update_init(&mut update) };
564 let ptr = self.session.ptr;
569 let held = self.holds_lease;
570 let (status, ended) = self.session.model.with_compute(None, |lease| {
571 let st = unsafe { sys::transcribe_stream_feed(ptr, pcm.as_ptr(), n, &mut update) };
572 let ended = unsafe { sys::transcribe_stream_get_state(ptr) }
573 == sys::transcribe_stream_state::TRANSCRIBE_STREAM_FAILED;
574 if held && ended {
575 *lease = false;
576 }
577 (st, ended)
578 })?;
579 if ended {
580 self.holds_lease = false;
581 }
582 check(status, "stream feed")?;
583 Ok(StreamUpdate::from_raw(&update))
584 }
585
586 pub fn finalize(&mut self) -> Result<StreamUpdate> {
589 let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
590 unsafe { sys::transcribe_stream_update_init(&mut update) };
591 let status = self.session.model.with_compute(None, |lease| {
592 let st = unsafe { sys::transcribe_stream_finalize(self.session.ptr, &mut update) };
593 if self.holds_lease {
597 *lease = false;
598 }
599 st
600 })?;
601 self.holds_lease = false;
602 check(status, "stream finalize")?;
603 Ok(StreamUpdate::from_raw(&update))
604 }
605
606 pub fn reset(&mut self) {
611 let _ = self.session.model.with_compute(None, |lease| {
613 unsafe { sys::transcribe_stream_reset(self.session.ptr) };
614 if self.holds_lease {
615 *lease = false;
616 }
617 });
618 self.holds_lease = false;
619 }
620
621 pub fn text(&self) -> StreamText {
623 let mut raw: sys::transcribe_stream_text = unsafe { std::mem::zeroed() };
624 unsafe { sys::transcribe_stream_text_init(&mut raw) };
625 let _ = unsafe { sys::transcribe_stream_get_text(self.session.ptr, &mut raw) };
626 StreamText::from_raw(&raw)
627 }
628
629 pub fn snapshot(&self) -> Transcript {
632 self.session.materialize_run()
633 }
634
635 pub fn state(&self) -> StreamState {
637 StreamState::from_raw(unsafe { sys::transcribe_stream_get_state(self.session.ptr) })
638 }
639
640 pub fn revision(&self) -> i32 {
642 unsafe { sys::transcribe_stream_revision(self.session.ptr) }
643 }
644
645 pub fn last_status(&self) -> Option<Error> {
651 let status = unsafe { sys::transcribe_stream_last_status(self.session.ptr) };
652 if status == sys::transcribe_status::TRANSCRIBE_OK {
653 None
654 } else {
655 Some(error_for_status(status, "stream"))
656 }
657 }
658
659 pub fn was_aborted(&self) -> bool {
661 self.session.was_aborted()
662 }
663}