add torchaudio

app.py

@@ -3,9 +3,9 @@ from transformers.file_utils import cached_path, hf_bucket_url
 import os, zipfile
 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 from datasets import load_dataset
-import soundfile as sf
 import torch
 import kenlm
+import torchaudio
 from pyctcdecode import Alphabet, BeamSearchDecoderCTC, LanguageModel
 
 cache_dir = './cache/'
@@ -17,7 +17,6 @@ with zipfile.ZipFile(lm_file, 'r') as zip_ref:
     zip_ref.extractall(cache_dir)
 lm_file = cache_dir + 'vi_lm_4grams.bin'
 
-
 def get_decoder_ngram_model(tokenizer, ngram_lm_path):
     vocab_dict = tokenizer.get_vocab()
     sort_vocab = sorted((value, key) for (key, value) in vocab_dict.items())
@@ -41,17 +40,25 @@ def get_decoder_ngram_model(tokenizer, ngram_lm_path):
 ngram_lm_model = get_decoder_ngram_model(processor.tokenizer, lm_file)
 
 # define function to read in sound file
-def map_to_array(batch):
-    speech, sampling_rate = sf.read(batch["file"])
-    batch["speech"] = speech
-    batch["sampling_rate"] = sampling_rate
+def speech_file_to_array_fn(path, max_seconds=10):
+    batch = {"file": path}
+    speech_array, sampling_rate = torchaudio.load(batch["file"])
+    if sampling_rate != 16000:
+      transform = torchaudio.transforms.Resample(orig_freq=sampling_rate,
+                                                 new_freq=16000)
+      speech_array = transform(speech_array)
+    speech_array = speech_array[0]
+    if max_seconds > 0:
+      speech_array = speech_array[:max_seconds*16000]
+    batch["speech"] = speech_array.numpy()
+    batch["sampling_rate"] = 16000
     return batch
 
 # tokenize
 def inference(audio):
    # read in sound file
     # load dummy dataset and read soundfiles
-    ds = map_to_array({"file": audio})
+    ds = speech_file_to_array_fn({"file": audio})
     # infer model
     input_values = processor(
           ds["speech"], 
@@ -59,7 +66,6 @@ def inference(audio):
           return_tensors="pt"
     ).input_values
     # decode ctc output
-    logits = model(input_values).logits[0]
     pred_ids = torch.argmax(logits, dim=-1)
     greedy_search_output = processor.decode(pred_ids)
     beam_search_output = ngram_lm_model.decode(logits.cpu().detach().numpy(), beam_width=500)

requirements.txt

@@ -1,7 +1,6 @@
 torch==1.9.0
+torchaudio==0.9.0
 transformers==4.9.2 
-soundfile 
-librosa
 datasets==1.11.0 
 pyctcdecode==v0.1.0
 https://github.com/kpu/kenlm/archive/master.zip