Add clap flag translate (#5)

Co-authored-by: Miguel Piedrafita <[email protected]>
m1guelpf · Jan 24, 2023 · dd706bf · dd706bf
1 parent 417e03f
commit dd706bf
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 84 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/main.rs b/src/main.rs
@@ -26,6 +26,10 @@ struct Args {
 
     /// Path to the audio file to transcribe
     audio: String,
+
+    /// Toggle translation
+    #[clap(short, long, default_value = "false")]
+    translate: bool,
 }
 
 #[tokio::main]
@@ -46,7 +50,7 @@ async fn main() {
     );
 
     let mut whisper = Whisper::new(Model::new(args.model), args.lang).await;
-    let transcript = whisper.transcribe(audio).unwrap();
+    let transcript = whisper.transcribe(audio, args.translate).unwrap();
 
     write_to(
         audio.with_file_name(format!("{file_name}.txt")),

diff --git a/src/utils.rs b/src/utils.rs
@@ -1,3 +1,4 @@
+#![allow(dead_code)]
 use futures_util::StreamExt;
 use indicatif::{ProgressBar, ProgressStyle};
 use num::integer::div_floor;

diff --git a/src/whisper.rs b/src/whisper.rs
@@ -335,9 +335,9 @@ impl Whisper {
         }
     }
 
-    pub fn transcribe<P: AsRef<Path>>(&mut self, audio: P) -> Result<Transcript> {
+    pub fn transcribe<P: AsRef<Path>>(&mut self, audio: P, translate: bool) -> Result<Transcript> {
         let mut params = FullParams::new(SamplingStrategy::Greedy { n_past: 0 });
-        params.set_translate(true);
+        params.set_translate(translate);
         params.set_print_special(false);
         params.set_print_progress(false);
         params.set_print_realtime(false);