Update README.md

README.md

@@ -85,16 +85,24 @@ To transcribe audio files the model can be used as a standalone acoustic model a
  transcription = processor.batch_decode(predicted_ids)
  ```
  
-  ## Evaluation
- 
- This code snippet shows how to evaluate **facebook/wav2vec2-large-960h-lv60-self** on LibriSpeech's "clean" and "other" test data.
+## Evaluation
  
+First, ensure the required Python packages are installed. We'll require `transformers` for running the Wav2Vec2 model,
+`datasets` for loading the LibriSpeech dataset, and `evaluate` plus `jiwer` for computing the word-error rate (WER):
+
+```
+pip install --upgrade pip
+pip install --upgrade transformers datasets evaluate jiwer
+```
+
+The following code snippet shows how to evaluate **facebook/wav2vec2-large-960h-lv60-self** on LibriSpeech's "clean" and "other" test data.
+The batch size can be set according to your device, and is set to `8` by default:
+
 ```python
+import torch
 from datasets import load_dataset
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-import torch
-from jiwer import wer
-
+from evaluate import load
 
 librispeech_eval = load_dataset("librispeech_asr", "clean", split="test")
 
@@ -102,21 +110,21 @@ model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h-lv60-self")
 processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h-lv60-self")
 
 def map_to_pred(batch):
-    inputs = processor(batch["audio"]["array"], return_tensors="pt", padding="longest")
-    input_values = inputs.input_values.to("cuda")
-    attention_mask = inputs.attention_mask.to("cuda")
-    
+    audios = [audio["array"] for audio in batch["audio"]]
+    sampling_rate = batch["audio"][0]["sampling_rate"]
+    input_values = processor(audios, sampling_rate=sampling_rate, return_tensors="pt", padding="longest").input_values
     with torch.no_grad():
-        logits = model(input_values, attention_mask=attention_mask).logits
+        logits = model(input_values.to("cuda")).logits
 
     predicted_ids = torch.argmax(logits, dim=-1)
     transcription = processor.batch_decode(predicted_ids)
-    batch["transcription"] = transcription
+    batch["transcription"] = [t for t in transcription]
     return batch
 
-result = librispeech_eval.map(map_to_pred, remove_columns=["audio"])
+result = librispeech_eval.map(map_to_pred, batched=True, batch_size=8, remove_columns=["audio"])
+wer = load("wer")
 
-print("WER:", wer(result["text"], result["transcription"]))
+print("WER:", wer.compute(references=result["text"], predictions=result["transcription"]))
 ```
 
 *Result (WER)*: