Kr08 commited on
Commit
f0f9802
1 Parent(s): abfb57d

Create audio_processing.py, include module for language detection, modifying sampling rate and processing audios longer than 30 seconds

Browse files
Files changed (1) hide show
  1. audio_processing.py +48 -0
audio_processing.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import whisper
3
+ import torchaudio as ta
4
+ from model_utils import get_processor, get_model, get_whisper_model_small, get_device
5
+ from config import SAMPLING_RATE, CHUNK_LENGTH_S
6
+
7
+ def detect_language(audio_file):
8
+ whisper_model = get_whisper_model_small()
9
+ trimmed_audio = whisper.pad_or_trim(audio_file.squeeze())
10
+ mel = whisper.log_mel_spectrogram(trimmed_audio).to(whisper_model.device)
11
+ _, probs = whisper_model.detect_language(mel)
12
+ detected_lang = max(probs[0], key=probs[0].get)
13
+ print(f"Detected language: {detected_lang}")
14
+ return detected_lang
15
+
16
+ def process_long_audio(waveform, sampling_rate, task="transcribe", language=None):
17
+ processor = get_processor()
18
+ model = get_model()
19
+ device = get_device()
20
+
21
+ input_length = waveform.shape[1]
22
+ chunk_length = int(CHUNK_LENGTH_S * sampling_rate)
23
+ chunks = [waveform[:, i:i + chunk_length] for i in range(0, input_length, chunk_length)]
24
+
25
+ results = []
26
+ for chunk in chunks:
27
+ input_features = processor(chunk[0], sampling_rate=sampling_rate, return_tensors="pt").input_features.to(device)
28
+
29
+ with torch.no_grad():
30
+ if task == "translate":
31
+ forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task="translate")
32
+ generated_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
33
+ else:
34
+ generated_ids = model.generate(input_features)
35
+
36
+ transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)
37
+ results.extend(transcription)
38
+
39
+ # Clear GPU cache
40
+ torch.cuda.empty_cache()
41
+
42
+ return " ".join(results)
43
+
44
+ def load_and_resample_audio(file):
45
+ waveform, sampling_rate = ta.load(file)
46
+ if sampling_rate != SAMPLING_RATE:
47
+ waveform = ta.functional.resample(waveform, orig_freq=sampling_rate, new_freq=SAMPLING_RATE)
48
+ return waveform