c4181/whisper-worker
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add mutex around call transcription

Browse source
This commit is contained in:
Christopher Moyer 2024-12-17 00:32:36 -05:00
parent 4f5c1f8ac4
commit af6376e2a0
1 changed files with 38 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
src/main.rs
View file

@ -4,6 +4,7 @@ use std::env;
use std::fs::File; use std::fs::File;
use std::path::Path; use std::path::Path;
use std::process::Command; use std::process::Command;
use std::sync::{Arc, Mutex};
use chrono::{DateTime, Utc}; use chrono::{DateTime, Utc};
use lapin::{Connection, ConnectionProperties, Consumer}; use lapin::{Connection, ConnectionProperties, Consumer};
use lapin::message::{Delivery, DeliveryResult}; use lapin::message::{Delivery, DeliveryResult};
@ -40,27 +41,42 @@ _|"""""|_|"""""|_|"""""|_|"""""|_|"""""|_|"""""|_|"""""| {======|_|"""""|_|"""""
let channel = connection.create_channel().await.unwrap(); let channel = connection.create_channel().await.unwrap();
let consumer = channel.basic_consume("transcribe", "whisper-worker", BasicConsumeOptions::default(), FieldTable::default()).await.unwrap(); let consumer = channel.basic_consume("transcribe", "whisper-worker", BasicConsumeOptions::default(), FieldTable::default()).await.unwrap();
consumer.set_delegate(move |delivery: DeliveryResult| async move { let processing_lock = Arc::new(Mutex::new(()));
info!("Consuming mq message"); consumer.set_delegate({
let delivery = match delivery { let processing_lock = Arc::clone(&processing_lock);
Ok(Some(delivery)) => delivery, move |delivery: DeliveryResult| {
Ok(None) => return, let processing_lock = Arc::clone(&processing_lock);
Err(error) => { async move {
warn!("Failed to consume queue message {}", error); info!("Consuming mq message");
return; let delivery = match delivery {
Ok(Some(delivery)) => delivery,
Ok(None) => return,
Err(error) => {
warn!("Failed to consume queue message {}", error);
return;
}
};
let transcription_request: TranscriptionRequest = serde_json::from_slice(&delivery.data).unwrap();
let path = Path::new(&transcription_request.audio_file_path);
if !path.exists() {
warn!("File not found: {}", &transcription_request.audio_file_path);
return delivery.ack(BasicAckOptions::default()).await.unwrap();
}
let result = {
info!("Waiting for lock");
let _lock = processing_lock.lock();
info!("Acquired lock!");
transcribe_call(path);
};
info!("Lock released");
info!("Call transcription done");
delivery.ack(BasicAckOptions::default()).await.unwrap()
} }
};
let transcription_request: TranscriptionRequest = serde_json::from_slice(&delivery.data).unwrap();
let path = Path::new(&transcription_request.audio_file_path);
if !path.exists() {
warn!("File not found: {}", &transcription_request.audio_file_path);
return delivery.ack(BasicAckOptions::default()).await.unwrap();
} }
transcribe_call(path);
info!("Call transcription done");
delivery.ack(BasicAckOptions::default()).await.unwrap()
}); });
info!("Waiting for messages"); info!("Waiting for messages");
@ -68,8 +84,9 @@ _|"""""|_|"""""|_|"""""|_|"""""|_|"""""|_|"""""|_|"""""| {======|_|"""""|_|"""""
} }
fn transcribe_call(file_path: &Path) { fn transcribe_call(file_path: &Path) {
info!("Transcribing file {}", file_path.display());
let output_directory = file_path.parent().unwrap(); let output_directory = file_path.parent().unwrap();
Command::new("whisperx") let output = Command::new("whisperx")
.args(["--language", "en"]) .args(["--language", "en"])
.args(["--model", "largev3"]) .args(["--model", "largev3"])
.args(["--batch_size", "4"]) .args(["--batch_size", "4"])
@ -78,6 +95,7 @@ fn transcribe_call(file_path: &Path) {
.args(["--output directory", output_directory.parent().unwrap().to_str().unwrap()]) .args(["--output directory", output_directory.parent().unwrap().to_str().unwrap()])
.arg("file") .arg("file")
.output().expect("TODO: panic message"); .output().expect("TODO: panic message");
info!("Transcription done, {}", &output.status);
} }
#[derive(Serialize, Deserialize, Default, Debug)] #[derive(Serialize, Deserialize, Default, Debug)]