Update README.md

README.md

@@ -13,7 +13,7 @@ widget:
 - example_title: Librispeech sample 2
   src: https://cdn-media.huggingface.co/speech_samples/sample2.flac
 model-index:
-- name: wav2vec2-base-960h
+- name: patrickvonplaten/wav2vec2-base-960h-4-gram
   results:
   - task:
       name: Automatic Speech Recognition
@@ -25,7 +25,7 @@ model-index:
     metrics:
     - name: Test WER
       type: wer
-      value: 3.4
+      value: 2.59
 ---
 
 # Wav2Vec2-Base-960h + 4-gram
@@ -39,33 +39,36 @@ augmented with an English 4-gram. The `4-gram.arpa.gz` of [Librispeech's officia
  
 ```python
 from datasets import load_dataset
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import AutoModelForCTC, AutoProcessor
 import torch
 from jiwer import wer
 
+model_id = "patrickvonplaten/wav2vec2-base-960h-4-gram"
 
-librispeech_eval = load_dataset("librispeech_asr", "clean", split="test")
+librispeech_eval = load_dataset("librispeech_asr", "other", split="test")
 
-model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h").to("cuda")
-processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+model = AutoModelForCTC.from_pretrained(model_id).to("cuda")
+processor = AutoProcessor.from_pretrained(model_id)
 
 def map_to_pred(batch):
-    input_values = processor(batch["audio"]["array"], return_tensors="pt", padding="longest").input_values
+    inputs = processor(batch["audio"]["array"], sampling_rate=16_000, return_tensors="pt")
+
+    inputs = {k: v.to("cuda") for k,v in inputs.items()}
+
     with torch.no_grad():
-        logits = model(input_values.to("cuda")).logits
+        logits = model(**inputs).logits
 
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.batch_decode(predicted_ids)
+    transcription = processor.batch_decode(logits.cpu().numpy()).text[0]
     batch["transcription"] = transcription
     return batch
 
-result = librispeech_eval.map(map_to_pred, batched=True, batch_size=1, remove_columns=["speech"])
+result = librispeech_eval.map(map_to_pred, remove_columns=["audio"])
 
-print("WER:", wer(result["text"], result["transcription"]))
+print(wer(result["text"], result["transcription"]))
 ```
 
 *Result (WER)*:
 
 | "clean" | "other" |
 |---|---|
-| 3.4 | 8.6 |
+| 2.59 | 6.46 |