From 4da16a45a388a876b968bdfbe001cf77c832329c Mon Sep 17 00:00:00 2001
From: 0/0 <zero@imaskeleton.me>
Date: Wed, 14 Dec 2022 17:21:10 -0700
Subject: [PATCH] handle error codes more often

---
 src/error.rs       | 24 +++++++++++++--
 src/whisper_ctx.rs | 73 +++++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 87 insertions(+), 10 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 1a8200f..592e6e8 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,21 +1,39 @@
 use std::ffi::{c_int, NulError};
 use std::str::Utf8Error;
 
+/// Whisper tends to output errors to stderr, so if an error occurs, check stderr.
 #[derive(Debug, Copy, Clone)]
 pub enum WhisperError {
+    /// Failed to create a new context.
     InitError,
+    /// User didn't initialize spectrogram
     SpectrogramNotInitialized,
+    /// Encode was not called.
     EncodeNotComplete,
+    /// Decode was not called.
     DecodeNotComplete,
+    /// Failed to calculate the spectrogram for some reason.
+    UnableToCalculateSpectrogram,
+    /// Failed to evaluate model.
+    UnableToCalculateEvaluation,
+    /// Failed to run the encoder
+    FailedToEncode,
+    /// Failed to run the decoder
+    FailedToDecode,
+    /// Invalid number of mel bands.
+    InvalidMelBands,
+    /// Invalid thread count
     InvalidThreadCount,
+    /// Invalid UTF-8 detected in a string from Whisper.
     InvalidUtf8 {
         error_len: Option<usize>,
         valid_up_to: usize,
     },
-    NullByteInString {
-        idx: usize,
-    },
+    /// A null byte was detected in a user-provided string.
+    NullByteInString { idx: usize },
+    /// Whisper returned a null pointer.
     NullPointer,
+    /// Generic whisper error. Varies depending on the function.
     GenericError(c_int),
 }
 
diff --git a/src/whisper_ctx.rs b/src/whisper_ctx.rs
index 0938821..6d00f96 100644
--- a/src/whisper_ctx.rs
+++ b/src/whisper_ctx.rs
@@ -68,7 +68,9 @@ impl WhisperContext {
                 threads as c_int,
             )
         };
-        if ret == 0 {
+        if ret == -1 {
+            Err(WhisperError::UnableToCalculateSpectrogram)
+        } else if ret == 0 {
             self.spectrogram_initialized = true;
             Ok(())
         } else {
@@ -101,7 +103,9 @@ impl WhisperContext {
                 80 as c_int,
             )
         };
-        if ret == 0 {
+        if ret == -1 {
+            Err(WhisperError::InvalidMelBands)
+        } else if ret == 0 {
             self.spectrogram_initialized = true;
             Ok(())
         } else {
@@ -130,7 +134,9 @@ impl WhisperContext {
         }
         let ret =
             unsafe { whisper_rs_sys::whisper_encode(self.ctx, offset as c_int, threads as c_int) };
-        if ret == 0 {
+        if ret == -1 {
+            Err(WhisperError::UnableToCalculateEvaluation)
+        } else if ret == 0 {
             self.encode_complete = true;
             Ok(())
         } else {
@@ -174,7 +180,9 @@ impl WhisperContext {
                 threads as c_int,
             )
         };
-        if ret == 0 {
+        if ret == -1 {
+            Err(WhisperError::UnableToCalculateEvaluation)
+        } else if ret == 0 {
             self.decode_once = true;
             Ok(())
         } else {
@@ -394,10 +402,61 @@ impl WhisperContext {
         let ret = unsafe {
             whisper_rs_sys::whisper_full(self.ctx, params.fp, data.as_ptr(), data.len() as c_int)
         };
-        if ret < 0 {
-            Err(WhisperError::GenericError(ret))
+        if ret == -1 {
+            Err(WhisperError::UnableToCalculateSpectrogram)
+        } else if ret == 7 {
+            Err(WhisperError::FailedToEncode)
+        } else if ret == 8 {
+            Err(WhisperError::FailedToDecode)
+        } else if ret == 0 {
+            Ok(ret)
         } else {
-            Ok(ret as c_int)
+            Err(WhisperError::GenericError(ret))
+        }
+    }
+
+    /// Split the input audio into chunks and delegate to [WhisperContext::full].
+    ///
+    /// It seems this approach can offer some speedup in some cases,
+    /// however, the accuracy can be worse at the start and end of chunks.
+    ///
+    /// # Arguments
+    /// * params: [crate::FullParams] struct.
+    /// * pcm: PCM audio data.
+    /// * n_processors: Number of threads to use.
+    ///
+    /// # Returns
+    /// Ok(c_int) on success, Err(WhisperError) on failure.
+    ///
+    /// # C++ equivalent
+    /// `int whisper_full_parallel(struct whisper_context * ctx, struct whisper_full_params params, const float * samples, int n_samples, int n_processors)`
+    pub fn full_parallel(
+        &mut self,
+        params: FullParams,
+        data: &[f32],
+        n_processors: c_int,
+    ) -> Result<c_int, WhisperError> {
+        let ret = unsafe {
+            whisper_rs_sys::whisper_full_parallel(
+                self.ctx,
+                params.fp,
+                data.as_ptr(),
+                data.len() as c_int,
+                n_processors,
+            )
+        };
+        if ret == -1 {
+            Err(WhisperError::UnableToCalculateSpectrogram)
+        } else if ret == 7 {
+            Err(WhisperError::FailedToEncode)
+        } else if ret == 8 {
+            Err(WhisperError::FailedToDecode)
+        } else if ret == 0 {
+            // note 0 is returned on success and also when initializing other contexts fails,
+            // causing some audio to not be processed
+            Ok(ret)
+        } else {
+            Err(WhisperError::GenericError(ret))
         }
     }