refactor(state): create_state from whisper context wrapper

This commit is contained in:
jiahua 2024-04-25 15:33:06 +08:00
parent 74e83185bf
commit 9d978add9d
3 changed files with 40 additions and 19 deletions

View file

@ -12,6 +12,7 @@ mod whisper_state;
mod whisper_sys_log;
#[cfg(feature = "whisper-cpp-tracing")]
mod whisper_sys_tracing;
mod whisper_ctx_wrapper;
#[cfg(any(feature = "whisper-cpp-log", feature = "whisper-cpp-tracing"))]
static LOG_TRAMPOLINE_INSTALL: Once = Once::new();
@ -22,6 +23,7 @@ pub use standalone::*;
use std::sync::Once;
pub use utilities::*;
pub use whisper_ctx::WhisperContext;
pub use whisper_ctx_wrapper::WhisperContextWrapper;
pub use whisper_ctx::WhisperContextParameters;
pub use whisper_grammar::{WhisperGrammarElement, WhisperGrammarElementType};
pub use whisper_params::{FullParams, SamplingStrategy};

View file

@ -10,7 +10,7 @@ use std::ffi::{c_int, CStr, CString};
/// then run a full transcription with [WhisperState::full].
#[derive(Debug)]
pub struct WhisperContext {
ctx: *mut whisper_rs_sys::whisper_context,
pub(crate) ctx: *mut whisper_rs_sys::whisper_context,
}
impl WhisperContext {
@ -114,24 +114,6 @@ impl WhisperContext {
}
}
// we don't implement `whisper_init()` here since i have zero clue what `whisper_model_loader` does
/// Create a new state object, ready for use.
///
/// # Returns
/// Ok(WhisperState) on success, Err(WhisperError) on failure.
///
/// # C++ equivalent
/// `struct whisper_state * whisper_init_state(struct whisper_context * ctx);`
pub fn create_state(&self) -> Result<WhisperState, WhisperError> {
let state = unsafe { whisper_rs_sys::whisper_init_state(self.ctx) };
if state.is_null() {
Err(WhisperError::InitError)
} else {
// SAFETY: this is known to be a valid pointer to a `whisper_state` struct
Ok(WhisperState::new(self.ctx, state))
}
}
/// Convert the provided text into tokens.
///

View file

@ -0,0 +1,37 @@
use std::sync::Arc;
use crate::{WhisperContext, WhisperContextParameters, WhisperError, WhisperState};
pub struct WhisperContextWrapper {
ctx: Arc<WhisperContext>,
}
impl WhisperContextWrapper {
/// wrapper of WhisperContext::new_with_params.
pub fn new_with_params(
path: &str,
parameters: WhisperContextParameters,
) -> Result<Self, WhisperError> {
let ctx = WhisperContext::new_with_params(path, parameters)?;
Ok(Self { ctx: Arc::new(ctx) })
}
// we don't implement `whisper_init()` here since i have zero clue what `whisper_model_loader` does
/// Create a new state object, ready for use.
///
/// # Returns
/// Ok(WhisperState) on success, Err(WhisperError) on failure.
///
/// # C++ equivalent
/// `struct whisper_state * whisper_init_state(struct whisper_context * ctx);`
pub fn create_state(&self) -> Result<WhisperState, WhisperError> {
let state = unsafe { whisper_rs_sys::whisper_init_state(self.ctx.ctx) };
if state.is_null() {
Err(WhisperError::InitError)
} else {
// SAFETY: this is known to be a valid pointer to a `whisper_state` struct
Ok(WhisperState::new(self.ctx.clone(), state))
}
}
}