use std::ffi::CString;
use std::os::raw::c_void;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use transcribe_cpp_sys as sys;
use crate::cancel::{abort_trampoline, CancelToken};
use crate::error::{check, error_for_status, Error, Result};
use crate::family::{RunExtRaw, RunExtension};
use crate::model::{Model, ModelInner, SessionLimits, SessionOptions};
use crate::result::{
owned_opt_str, owned_str, Segment, SpeakerSegment, Timings, Token, Transcript, Word,
};
use crate::streaming::{build_stream_params, StreamOptions, StreamText, StreamUpdate};
use crate::types::{Diarize, Itn, Pnc, StreamState, Task, TimestampKind};
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(
feature = "serde",
derive(serde::Serialize, serde::Deserialize),
serde(default)
)]
pub struct RunOptions {
pub task: Task,
pub timestamps: TimestampKind,
pub pnc: Pnc,
pub itn: Itn,
pub diarize: Diarize,
pub language: Option<String>,
pub target_language: Option<String>,
pub keep_special_tags: bool,
pub spec_k_drafts: i32,
pub family: Option<RunExtension>,
pub vocabulary: Vec<String>,
pub prompt: Option<String>,
pub prefix: Option<String>,
}
impl Default for RunOptions {
fn default() -> Self {
RunOptions {
task: Task::Transcribe,
timestamps: TimestampKind::Auto,
pnc: Pnc::Default,
itn: Itn::Default,
diarize: Diarize::Default,
language: None,
target_language: None,
keep_special_tags: false,
spec_k_drafts: -1,
family: None,
vocabulary: Vec::new(),
prompt: None,
prefix: None,
}
}
}
pub struct Session {
ptr: *mut sys::transcribe_session,
model: Arc<ModelInner>,
cancel: Option<Arc<AtomicBool>>,
}
impl std::fmt::Debug for Session {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Session").finish_non_exhaustive()
}
}
unsafe impl Send for Session {}
impl Drop for Session {
fn drop(&mut self) {
unsafe { sys::transcribe_session_free(self.ptr) };
}
}
impl Session {
pub(crate) fn new(model: &Model, options: &SessionOptions) -> Result<Session> {
let mut params: sys::transcribe_session_params = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_session_params_init(&mut params) };
params.n_threads = options.n_threads;
params.kv_type = options.kv_type.to_raw();
params.n_ctx = options.n_ctx;
let mut out: *mut sys::transcribe_session = std::ptr::null_mut();
let status = unsafe { sys::transcribe_session_init(model.inner.ptr, ¶ms, &mut out) };
check(status, "session init")?;
debug_assert!(!out.is_null());
Ok(Session {
ptr: out,
model: Arc::clone(&model.inner),
cancel: None,
})
}
pub fn set_cancel_token(&mut self, token: &CancelToken) {
let flag = Arc::clone(&token.flag);
let userdata = Arc::as_ptr(&flag) as *mut c_void;
unsafe { sys::transcribe_set_abort_callback(self.ptr, Some(abort_trampoline), userdata) };
self.cancel = Some(flag);
}
pub fn clear_cancel_token(&mut self) {
unsafe { sys::transcribe_set_abort_callback(self.ptr, None, std::ptr::null_mut()) };
self.cancel = None;
}
pub fn was_aborted(&self) -> bool {
unsafe { sys::transcribe_was_aborted(self.ptr) }
}
pub fn was_truncated(&self) -> bool {
unsafe { sys::transcribe_was_truncated(self.ptr) }
}
pub fn run(&mut self, pcm: &[f32], options: &RunOptions) -> Result<Transcript> {
let (params, _lang, _target, _family, _prompting) = build_run_params(options)?;
let n = clamp_len(pcm.len())?;
let status = self.model.with_compute(
Some("a stream is active on this model; finish or drop it before run()"),
|_| unsafe { sys::transcribe_run(self.ptr, pcm.as_ptr(), n, ¶ms) },
)?;
match status {
s if s == sys::transcribe_status::TRANSCRIBE_OK => Ok(self.materialize_run()),
s if s == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
|| s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
|| s == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_REPETITION =>
{
let partial = Box::new(self.materialize_run());
Err(attach_partial(error_for_status(s, "run"), partial))
}
s => Err(error_for_status(s, "run")),
}
}
pub fn run_batch(
&mut self,
pcms: &[&[f32]],
options: &RunOptions,
) -> Result<Vec<Result<Transcript>>> {
let (params, _lang, _target, _family, _prompting) = build_run_params(options)?;
let ptrs: Vec<*const f32> = pcms.iter().map(|p| p.as_ptr()).collect();
let lens: Vec<i32> = pcms
.iter()
.map(|p| clamp_len(p.len()))
.collect::<Result<_>>()?;
let n = clamp_len(pcms.len())?;
let status = self.model.with_compute(
Some("a stream is active on this model; finish or drop it before run_batch()"),
|_| unsafe {
sys::transcribe_run_batch(self.ptr, ptrs.as_ptr(), lens.as_ptr(), n, ¶ms)
},
)?;
let ok = status == sys::transcribe_status::TRANSCRIBE_OK
|| status == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED;
if !ok {
return Err(error_for_status(status, "run_batch"));
}
let count = unsafe { sys::transcribe_batch_n_results(self.ptr) };
let mut results = Vec::with_capacity(count.max(0) as usize);
for i in 0..count {
let st = unsafe { sys::transcribe_batch_status(self.ptr, i) };
if st == sys::transcribe_status::TRANSCRIBE_OK {
results.push(Ok(self.materialize_batch(i)));
} else if st == sys::transcribe_status::TRANSCRIBE_ERR_ABORTED
|| st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_TRUNCATED
|| st == sys::transcribe_status::TRANSCRIBE_ERR_OUTPUT_REPETITION
{
let partial = Box::new(self.materialize_batch(i));
let err = attach_partial(error_for_status(st, "run_batch utterance"), partial);
results.push(Err(err));
} else {
results.push(Err(error_for_status(st, "run_batch utterance")));
}
}
Ok(results)
}
pub fn limits(&self) -> Result<SessionLimits> {
let mut out: sys::transcribe_session_limits = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_session_limits_init(&mut out) };
let status = unsafe { sys::transcribe_session_get_limits(self.ptr, &mut out) };
check(status, "session limits")?;
Ok(SessionLimits {
effective_n_ctx: out.effective_n_ctx,
effective_max_audio_ms: out.effective_max_audio_ms,
max_kv_bytes: out.max_kv_bytes,
})
}
pub fn model(&self) -> Model {
Model {
inner: Arc::clone(&self.model),
}
}
pub fn stream(&mut self, run: &RunOptions, stream: &StreamOptions) -> Result<Stream<'_>> {
let (run_params, _lang, _target, _family, _prompting) = build_run_params(run)?;
let (stream_params, _stream_family) = build_stream_params(stream);
self.model.with_compute(
Some("a stream is already active on this model"),
|lease| -> Result<()> {
check(
unsafe { sys::transcribe_stream_begin(self.ptr, &run_params, &stream_params) },
"stream begin",
)?;
*lease = true; Ok(())
},
)??;
Ok(Stream {
session: self,
holds_lease: true,
})
}
fn materialize_run(&self) -> Transcript {
let n_seg = unsafe { sys::transcribe_n_segments(self.ptr) };
let n_word = unsafe { sys::transcribe_n_words(self.ptr) };
let n_tok = unsafe { sys::transcribe_n_tokens(self.ptr) };
let n_speaker = unsafe { sys::transcribe_n_speaker_segments(self.ptr) };
Transcript {
text: owned_str(unsafe { sys::transcribe_full_text(self.ptr) }),
raw_text: owned_str(unsafe { sys::transcribe_raw_text(self.ptr) }),
language: owned_opt_str(unsafe { sys::transcribe_detected_language(self.ptr) }),
timestamp_kind: TimestampKind::from_raw(unsafe {
sys::transcribe_returned_timestamp_kind(self.ptr)
}),
segments: (0..n_seg).map(|i| self.segment(i)).collect(),
speaker_segments: (0..n_speaker).map(|i| self.speaker_segment(i)).collect(),
words: (0..n_word).map(|i| self.word(i)).collect(),
tokens: (0..n_tok).map(|i| self.token(i)).collect(),
timings: self.timings(),
}
}
fn segment(&self, i: i32) -> Segment {
let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_segment_init(&mut raw) };
let _ = unsafe { sys::transcribe_get_segment(self.ptr, i, &mut raw) };
Segment::from_raw(&raw)
}
fn speaker_segment(&self, i: i32) -> SpeakerSegment {
let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
let _ = unsafe { sys::transcribe_get_speaker_segment(self.ptr, i, &mut raw) };
SpeakerSegment::from_raw(&raw)
}
fn word(&self, i: i32) -> Word {
let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_word_init(&mut raw) };
let _ = unsafe { sys::transcribe_get_word(self.ptr, i, &mut raw) };
Word::from_raw(&raw)
}
fn token(&self, i: i32) -> Token {
let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_token_init(&mut raw) };
let _ = unsafe { sys::transcribe_get_token(self.ptr, i, &mut raw) };
Token::from_raw(&raw)
}
fn timings(&self) -> Timings {
let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_timings_init(&mut raw) };
let _ = unsafe { sys::transcribe_get_timings(self.ptr, &mut raw) };
Timings::from_raw(&raw)
}
fn materialize_batch(&self, i: i32) -> Transcript {
let n_seg = unsafe { sys::transcribe_batch_n_segments(self.ptr, i) };
let n_word = unsafe { sys::transcribe_batch_n_words(self.ptr, i) };
let n_tok = unsafe { sys::transcribe_batch_n_tokens(self.ptr, i) };
let n_speaker = unsafe { sys::transcribe_batch_n_speaker_segments(self.ptr, i) };
Transcript {
text: owned_str(unsafe { sys::transcribe_batch_full_text(self.ptr, i) }),
raw_text: owned_str(unsafe { sys::transcribe_batch_raw_text(self.ptr, i) }),
language: owned_opt_str(unsafe {
sys::transcribe_batch_detected_language(self.ptr, i)
}),
timestamp_kind: TimestampKind::from_raw(unsafe {
sys::transcribe_batch_returned_timestamp_kind(self.ptr, i)
}),
segments: (0..n_seg).map(|j| self.batch_segment(i, j)).collect(),
speaker_segments: (0..n_speaker)
.map(|j| self.batch_speaker_segment(i, j))
.collect(),
words: (0..n_word).map(|j| self.batch_word(i, j)).collect(),
tokens: (0..n_tok).map(|j| self.batch_token(i, j)).collect(),
timings: self.batch_timings(i),
}
}
fn batch_segment(&self, i: i32, j: i32) -> Segment {
let mut raw: sys::transcribe_segment = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_segment_init(&mut raw) };
let _ = unsafe { sys::transcribe_batch_get_segment(self.ptr, i, j, &mut raw) };
Segment::from_raw(&raw)
}
fn batch_speaker_segment(&self, i: i32, j: i32) -> SpeakerSegment {
let mut raw: sys::transcribe_speaker_segment = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_speaker_segment_init(&mut raw) };
let _ = unsafe { sys::transcribe_batch_get_speaker_segment(self.ptr, i, j, &mut raw) };
SpeakerSegment::from_raw(&raw)
}
fn batch_word(&self, i: i32, j: i32) -> Word {
let mut raw: sys::transcribe_word = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_word_init(&mut raw) };
let _ = unsafe { sys::transcribe_batch_get_word(self.ptr, i, j, &mut raw) };
Word::from_raw(&raw)
}
fn batch_token(&self, i: i32, j: i32) -> Token {
let mut raw: sys::transcribe_token = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_token_init(&mut raw) };
let _ = unsafe { sys::transcribe_batch_get_token(self.ptr, i, j, &mut raw) };
Token::from_raw(&raw)
}
fn batch_timings(&self, i: i32) -> Timings {
let mut raw: sys::transcribe_timings = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_timings_init(&mut raw) };
let _ = unsafe { sys::transcribe_batch_get_timings(self.ptr, i, &mut raw) };
Timings::from_raw(&raw)
}
}
type RunParamsBundle = (
sys::transcribe_run_params,
Option<CString>,
Option<CString>,
Option<RunExtRaw>,
PromptingKeepalive,
);
struct PromptingKeepalive {
_terms: Vec<CString>,
_term_ptrs: Vec<*const std::os::raw::c_char>,
_prompt: Option<CString>,
_prefix: Option<CString>,
}
fn build_run_params(o: &RunOptions) -> Result<RunParamsBundle> {
let mut params: sys::transcribe_run_params = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_run_params_init(&mut params) };
params.task = o.task.to_raw();
params.timestamps = o.timestamps.to_raw();
params.pnc = o.pnc.to_raw();
params.itn = o.itn.to_raw();
params.diarize = o.diarize.to_raw();
params.keep_special_tags = o.keep_special_tags;
params.spec_k_drafts = o.spec_k_drafts;
let lang = o.language.as_deref().map(CString::new).transpose()?;
let target = o.target_language.as_deref().map(CString::new).transpose()?;
params.language = lang.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
params.target_language = target.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
let family = o
.family
.as_ref()
.map(RunExtension::materialize)
.transpose()?;
params.family = family.as_ref().map_or(std::ptr::null(), |f| f.ext_ptr());
let terms = o
.vocabulary
.iter()
.map(|t| CString::new(t.as_str()))
.collect::<std::result::Result<Vec<_>, _>>()?;
let term_ptrs: Vec<*const std::os::raw::c_char> = terms.iter().map(|c| c.as_ptr()).collect();
let prompt = o.prompt.as_deref().map(CString::new).transpose()?;
let prefix = o.prefix.as_deref().map(CString::new).transpose()?;
params.vocabulary = if term_ptrs.is_empty() {
std::ptr::null()
} else {
term_ptrs.as_ptr()
};
params.n_vocabulary = clamp_len(term_ptrs.len())?;
params.prompt = prompt.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
params.prefix = prefix.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
let prompting = PromptingKeepalive {
_terms: terms,
_term_ptrs: term_ptrs,
_prompt: prompt,
_prefix: prefix,
};
Ok((params, lang, target, family, prompting))
}
pub(crate) fn clamp_len(len: usize) -> Result<i32> {
i32::try_from(len).map_err(|_| Error::InvalidArgument(format!("length {len} exceeds i32::MAX")))
}
fn attach_partial(err: Error, partial: Box<Transcript>) -> Error {
match err {
Error::Aborted { message, .. } => Error::Aborted {
message,
partial: Some(partial),
},
Error::OutputTruncated { message, .. } => Error::OutputTruncated {
message,
partial: Some(partial),
},
Error::OutputRepetition { message, .. } => Error::OutputRepetition {
message,
partial: Some(partial),
},
other => other,
}
}
pub struct Stream<'a> {
session: &'a mut Session,
holds_lease: bool,
}
impl std::fmt::Debug for Stream<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Stream")
.field("state", &self.state())
.field("revision", &self.revision())
.finish()
}
}
impl Drop for Stream<'_> {
fn drop(&mut self) {
self.reset()
}
}
impl Stream<'_> {
pub fn feed(&mut self, pcm: &[f32]) -> Result<StreamUpdate> {
let n = clamp_len(pcm.len())?;
let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_stream_update_init(&mut update) };
let ptr = self.session.ptr;
let held = self.holds_lease;
let (status, ended) = self.session.model.with_compute(None, |lease| {
let st = unsafe { sys::transcribe_stream_feed(ptr, pcm.as_ptr(), n, &mut update) };
let ended = unsafe { sys::transcribe_stream_get_state(ptr) }
== sys::transcribe_stream_state::TRANSCRIBE_STREAM_FAILED;
if held && ended {
*lease = false;
}
(st, ended)
})?;
if ended {
self.holds_lease = false;
}
check(status, "stream feed")?;
Ok(StreamUpdate::from_raw(&update))
}
pub fn finalize(&mut self) -> Result<StreamUpdate> {
let mut update: sys::transcribe_stream_update = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_stream_update_init(&mut update) };
let status = self.session.model.with_compute(None, |lease| {
let st = unsafe { sys::transcribe_stream_finalize(self.session.ptr, &mut update) };
if self.holds_lease {
*lease = false;
}
st
})?;
self.holds_lease = false;
check(status, "stream finalize")?;
Ok(StreamUpdate::from_raw(&update))
}
pub fn reset(&mut self) {
let _ = self.session.model.with_compute(None, |lease| {
unsafe { sys::transcribe_stream_reset(self.session.ptr) };
if self.holds_lease {
*lease = false;
}
});
self.holds_lease = false;
}
pub fn text(&self) -> StreamText {
let mut raw: sys::transcribe_stream_text = unsafe { std::mem::zeroed() };
unsafe { sys::transcribe_stream_text_init(&mut raw) };
let _ = unsafe { sys::transcribe_stream_get_text(self.session.ptr, &mut raw) };
StreamText::from_raw(&raw)
}
pub fn snapshot(&self) -> Transcript {
self.session.materialize_run()
}
pub fn state(&self) -> StreamState {
StreamState::from_raw(unsafe { sys::transcribe_stream_get_state(self.session.ptr) })
}
pub fn revision(&self) -> i32 {
unsafe { sys::transcribe_stream_revision(self.session.ptr) }
}
pub fn last_status(&self) -> Option<Error> {
let status = unsafe { sys::transcribe_stream_last_status(self.session.ptr) };
if status == sys::transcribe_status::TRANSCRIBE_OK {
None
} else {
Some(error_for_status(status, "stream"))
}
}
pub fn was_aborted(&self) -> bool {
self.session.was_aborted()
}
}