whisper_cpp_plus/
state.rs

1use crate::context::{ContextPtr, WhisperContext};
2use crate::error::{Result, WhisperError};
3use crate::params::FullParams;
4use std::sync::Arc;
5use whisper_cpp_plus_sys as ffi;
6
7pub struct WhisperState {
8    pub(crate) ptr: *mut ffi::whisper_state,
9    pub(crate) _context: Arc<ContextPtr>,
10}
11
12impl Drop for WhisperState {
13    fn drop(&mut self) {
14        unsafe {
15            if !self.ptr.is_null() {
16                ffi::whisper_free_state(self.ptr);
17            }
18        }
19    }
20}
21
22impl WhisperState {
23    pub fn new(context: &WhisperContext) -> Result<Self> {
24        let ptr = unsafe { ffi::whisper_init_state(context.ptr.0) };
25
26        if ptr.is_null() {
27            return Err(WhisperError::OutOfMemory);
28        }
29
30        Ok(Self {
31            ptr,
32            _context: Arc::clone(&context.ptr),
33        })
34    }
35
36    pub fn full(&mut self, params: FullParams, audio: &[f32]) -> Result<()> {
37        if audio.is_empty() {
38            return Err(WhisperError::InvalidAudioFormat);
39        }
40
41        let ret = unsafe {
42            ffi::whisper_full_with_state(
43                self._context.0,
44                self.ptr,
45                params.as_raw(),
46                audio.as_ptr(),
47                audio.len() as i32,
48            )
49        };
50
51        if ret != 0 {
52            return Err(WhisperError::TranscriptionError(format!(
53                "whisper_full returned {}",
54                ret
55            )));
56        }
57
58        Ok(())
59    }
60
61    pub fn full_parallel(
62        &mut self,
63        params: FullParams,
64        audio: &[f32],
65        n_processors: i32,
66    ) -> Result<()> {
67        if audio.is_empty() {
68            return Err(WhisperError::InvalidAudioFormat);
69        }
70
71        if n_processors < 1 {
72            return Err(WhisperError::InvalidParameter(
73                "n_processors must be at least 1".into(),
74            ));
75        }
76
77        let ret = unsafe {
78            ffi::whisper_full_parallel(
79                self._context.0,
80                params.as_raw(),
81                audio.as_ptr(),
82                audio.len() as i32,
83                n_processors,
84            )
85        };
86
87        if ret != 0 {
88            return Err(WhisperError::TranscriptionError(format!(
89                "whisper_full_parallel returned {}",
90                ret
91            )));
92        }
93
94        Ok(())
95    }
96
97    pub fn full_n_segments(&self) -> i32 {
98        unsafe { ffi::whisper_full_n_segments_from_state(self.ptr) }
99    }
100
101    pub fn full_lang_id(&self) -> i32 {
102        unsafe { ffi::whisper_full_lang_id_from_state(self.ptr) }
103    }
104
105    pub fn full_get_segment_text(&self, i_segment: i32) -> Result<String> {
106        let text_ptr = unsafe { ffi::whisper_full_get_segment_text_from_state(self.ptr, i_segment) };
107
108        if text_ptr.is_null() {
109            return Err(WhisperError::InvalidContext);
110        }
111
112        let c_str = unsafe { std::ffi::CStr::from_ptr(text_ptr) };
113        Ok(c_str.to_string_lossy().into_owned())
114    }
115
116    pub fn full_get_segment_timestamps(&self, i_segment: i32) -> (i64, i64) {
117        unsafe {
118            let t0 = ffi::whisper_full_get_segment_t0_from_state(self.ptr, i_segment);
119            let t1 = ffi::whisper_full_get_segment_t1_from_state(self.ptr, i_segment);
120            (t0, t1)
121        }
122    }
123
124    pub fn full_get_segment_speaker_turn_next(&self, i_segment: i32) -> bool {
125        unsafe {
126            ffi::whisper_full_get_segment_speaker_turn_next_from_state(self.ptr, i_segment)
127        }
128    }
129
130    pub fn full_n_tokens(&self, i_segment: i32) -> i32 {
131        unsafe { ffi::whisper_full_n_tokens_from_state(self.ptr, i_segment) }
132    }
133
134    pub fn full_get_token_text(&self, i_segment: i32, i_token: i32) -> Result<String> {
135        let text_ptr = unsafe {
136            ffi::whisper_full_get_token_text_from_state(self._context.0, self.ptr, i_segment, i_token)
137        };
138
139        if text_ptr.is_null() {
140            return Err(WhisperError::InvalidContext);
141        }
142
143        let c_str = unsafe { std::ffi::CStr::from_ptr(text_ptr) };
144        Ok(c_str.to_string_lossy().into_owned())
145    }
146
147    pub fn full_get_token_id(&self, i_segment: i32, i_token: i32) -> i32 {
148        unsafe { ffi::whisper_full_get_token_id_from_state(self.ptr, i_segment, i_token) }
149    }
150
151    pub fn full_get_token_data(
152        &self,
153        i_segment: i32,
154        i_token: i32,
155    ) -> Option<ffi::whisper_token_data> {
156        let data = unsafe {
157            ffi::whisper_full_get_token_data_from_state(self.ptr, i_segment, i_token)
158        };
159
160        if data.id == -1 {
161            None
162        } else {
163            Some(data)
164        }
165    }
166
167    pub fn full_get_token_prob(&self, i_segment: i32, i_token: i32) -> f32 {
168        unsafe { ffi::whisper_full_get_token_p_from_state(self.ptr, i_segment, i_token) }
169    }
170}
171
172unsafe impl Send for WhisperState {}
173
174#[derive(Debug, Clone)]
175pub struct TranscriptionResult {
176    pub text: String,
177    pub segments: Vec<Segment>,
178}
179
180#[derive(Debug, Clone)]
181pub struct Segment {
182    pub start_ms: i64,
183    pub end_ms: i64,
184    pub text: String,
185    pub speaker_turn_next: bool,
186}
187
188impl Segment {
189    pub fn start_seconds(&self) -> f64 {
190        self.start_ms as f64 / 1000.0
191    }
192
193    pub fn end_seconds(&self) -> f64 {
194        self.end_ms as f64 / 1000.0
195    }
196}
whisper_cpp_plus/state.rs

whisper_cpp_plus/
state.rs
