From 7ec77f7fac584edc553e54680a2bdaa897d99519 Mon Sep 17 00:00:00 2001
From: Andrey Tkachenko <andreytkachenko64@gmail.com>
Date: Wed, 19 Aug 2020 17:54:09 +0400
Subject: [PATCH] add transcription structs

---
 Cargo.toml           |  8 +++++---
 examples/demo.rs     |  6 +++---
 src/lib.rs           |  4 +++-
 src/model.rs         | 12 ++++++------
 src/transcription.rs | 21 +++++++++++++++++++++
 5 files changed, 38 insertions(+), 13 deletions(-)
 create mode 100644 src/transcription.rs

diff --git a/Cargo.toml b/Cargo.toml
index 53c3853..8baea42 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,10 +10,12 @@ build = "build.rs"
 [dependencies]
 cblas = "0.2.0"
 openblas-src = { version = "0.9.0", features = ["cache", "static"] }
+serde = {version = "1.0", features = ["derive"]}
+serde_json = "1.0"
 
 [build-dependencies]
-bindgen = "0.54.1"
-cc = { version = "1.0.58", features = ["parallel"] }
+bindgen = "0.54"
+cc = { version = "1.0", features = ["parallel"] }
 
 [dev-dependencies]
-audrey = "0.2.0"
+audrey = "0.2"
diff --git a/examples/demo.rs b/examples/demo.rs
index 94fcbdd..58cc7c0 100644
--- a/examples/demo.rs
+++ b/examples/demo.rs
@@ -40,13 +40,13 @@ pub fn main() {
         println!("feed {}", buff.len());
 
         if model.feed(&mut sess, buff.as_slice()) {
-            println!("{}", model.get_result(&mut sess));
+            println!("{:?}", model.get_result(&mut sess));
         } else {
-            println!("{}", model.get_partial_result(&mut sess));
+            println!("{:?}", model.get_partial_result(&mut sess));
         }
     }
 
-    println!("{}", model.get_final_result(&mut sess));
+    println!("{:?}", model.get_final_result(&mut sess));
 
     // let audio_buf :Vec<_> = if desc.sample_rate() == SAMPLE_RATE {
     //     .map(|s| s.unwrap()).collect()
diff --git a/src/lib.rs b/src/lib.rs
index 85bf53e..9086cba 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -17,7 +17,9 @@ mod ffi {
 mod model;
 mod session;
 mod speaker;
+mod transcription;
 
 pub use model::VoskModel;
 pub use session::{VoskSession, VoskSessionConfig, VoskSessionConfigBuilder};
-pub use speaker::SpeakerModel;
\ No newline at end of file
+pub use speaker::SpeakerModel;
+pub use transcription::{TranscriptionResult, TranscriptionPartialResult, TranscriptionWord};
\ No newline at end of file
diff --git a/src/model.rs b/src/model.rs
index 443a7b0..41dec87 100644
--- a/src/model.rs
+++ b/src/model.rs
@@ -28,24 +28,24 @@ impl VoskModel {
     }
 
     #[inline]
-    pub fn get_result(&self, sess: &mut VoskSession) -> String {
+    pub fn get_result(&self, sess: &mut VoskSession) -> crate::TranscriptionResult {
         let cstr = unsafe { CStr::from_ptr(ffi::vosk_recognizer_result(sess.inner)) };
 
-        cstr.to_string_lossy().to_owned().to_string()
+        serde_json::from_str(cstr.to_str().unwrap()).unwrap()
     }
 
     #[inline]
-    pub fn get_partial_result(&self, sess: &mut VoskSession) -> String {
+    pub fn get_partial_result(&self, sess: &mut VoskSession) -> crate::TranscriptionPartialResult {
         let cstr = unsafe { CStr::from_ptr(ffi::vosk_recognizer_partial_result(sess.inner)) };
 
-        cstr.to_string_lossy().to_owned().to_string()
+        serde_json::from_str(cstr.to_str().unwrap()).unwrap()
     }
 
     #[inline]
-    pub fn get_final_result(&self, sess: &mut VoskSession) -> String {
+    pub fn get_final_result(&self, sess: &mut VoskSession) -> crate::TranscriptionResult {
         let cstr = unsafe { CStr::from_ptr(ffi::vosk_recognizer_final_result(sess.inner)) };
 
-        cstr.to_string_lossy().to_owned().to_string()
+        serde_json::from_str(cstr.to_str().unwrap()).unwrap()
     }
 }
 
diff --git a/src/transcription.rs b/src/transcription.rs
new file mode 100644
index 0000000..29e1c30
--- /dev/null
+++ b/src/transcription.rs
@@ -0,0 +1,21 @@
+use serde::{Serialize, Deserialize};
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct TranscriptionWord {
+    pub conf: f32,
+    pub end: f32,
+    pub start: f32,
+    pub word: String,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct TranscriptionResult {
+    pub text: String,
+    #[serde(default = "Vec::new")]
+    pub result: Vec<TranscriptionWord>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct TranscriptionPartialResult {
+    pub partial: String,
+}
\ No newline at end of file