Update README.md

README.md

@@ -56,19 +56,35 @@ This code snippet shows how to evaluate **Wav2Vec2-Large-Tedlium** on the TEDLIU
  
 ```python
 from datasets import load_dataset
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
+from transformers import AutoProcessor, SpeechEncoderDecoderModel
 import torch
 from jiwer import wer
+
 tedlium_eval = load_dataset("LIUM/tedlium", "release3", split="test")
-model = Wav2Vec2ForCTC.from_pretrained("sanchit-gandhi/wav2vec2-large-tedlium").to("cuda")
-processor = Wav2Vec2Processor.from_pretrained("sanchit-gandhi/wav2vec2-large-tedlium")
+
+def filter_ds(text):
+    return text != "ignore_time_segment_in_scoring"
+
+# remove samples ignored from scoring
+tedlium_eval = tedlium_eval.map(filter_ds, input_columns=["text"])
+
+model = SpeechEncoderDecoderModel.from_pretrained("sanchit-gandhi/wav2vec2-2-bart-large-tedlium").to("cuda")
+processor = AutoProcessor.from_pretrained("sanchit-gandhi/wav2vec2-2-bart-large-tedlium")
+
+gen_kwargs = {
+        "max_length": 200,
+        "num_beams": 5,
+        "length_penalty": 1.2
+        }
+
 def map_to_pred(batch):
     input_values = processor(batch["audio"]["array"], return_tensors="pt", padding="longest").input_values
     with torch.no_grad():
-        logits = model(input_values.to("cuda")).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.batch_decode(predicted_ids)
-    batch["transcription"] = transcription
+        generated = model.generate(input_values.to("cuda"), **gen_kwargs)
+    decoded = processor.batch_decode(generated, skip_special_tokens=True)
+    batch["transcription"] = decoded[0]
     return batch
+
 result = tedlium_eval.map(map_to_pred, batched=True, batch_size=1, remove_columns=["speech"])
-print("WER:", wer(result["text"], result["transcription"]))
+print("WER:", wer(result["text"], result["transcription"]))
+```