Update README.md

README.md

@@ -25,45 +25,30 @@ See this GitHub Repo [cantonese-selfish-project](https://github.com/scottykwok/c
 
 # Usage
 ```python
-import time
+import soundfile as sf
 import torch
-import torchaudio
 from datasets import load_dataset
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-import sys
 
-# inputs
-model_id = "scottykwok/wav2vec2-large-xlsr-cantonese"
-try:
-    wav_file = sys.argv[1]
-except:
-    print("Please provide an input wav filename ")
-    exit(-1)
+# load pretrained model
+processor = Wav2Vec2Processor.from_pretrained("scottykwok/wav2vec2-large-xlsr-cantonese")
+model = Wav2Vec2ForCTC.from_pretrained("scottykwok/wav2vec2-large-xlsr-cantonese")
 
-print("-"* 20)
-print("Model ID:" , model_id)
-print("Input Audio:" , wav_file)
-print("-"* 20)
+# load audio - must be 16kHz mono
+audio_input, sample_rate = sf.read('audio.wav')
 
-# load model and tokenizer
-processor = Wav2Vec2Processor.from_pretrained(model_id)
-model = Wav2Vec2ForCTC.from_pretrained(model_id)
+# pad input values and return pt tensor
+input_values = processor(audio_input, sampling_rate=sample_rate, return_tensors="pt").input_values
 
-# read audio to numpy
-def speech_to_array(path):
-    speech_array, sampling_rate = torchaudio.load(path)
-    return speech_array.squeeze().numpy()
-input_array = speech_to_array(wav_file)
+# INFERENCE
+# retrieve logits & take argmax
+logits = model(input_values).logits
+predicted_ids = torch.argmax(logits, dim=-1)
 
-# tokenize
-inputs = processor([input_array], sampling_rate=16_000, return_tensors="pt", padding=True)
+# transcribe
+transcription = processor.decode(predicted_ids[0])
+print("-" *20)
+print("Transcription:\n", transcription.lower())
+print("-" *20)
 
-# inference
-with torch.no_grad():
-    logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
-predicted_ids = torch.argmax(logits, dim=-1)
-pred = processor.batch_decode(predicted_ids)
-print("-"* 20)
-print("Prediction:", pred)
-print("-"* 20)
 ```