Update model card example.

README.md

@@ -20,30 +20,38 @@ Training of the acoustic model is the work of KBLab. See [VoxRex-C](https://hugg
 VoxRex-C is extended with a 4-gram language model estimated from a subset extracted from [The Swedish Culturomics Gigaword Corpus](https://spraakbanken.gu.se/resurser/gigaword) from Språkbanken. The subset contains 40M words from the social media genre between 2010 and 2015.
 
 ## How to use
-Audio should be downsampled to 16kHz.
+Example of transcribing 1% of the Common Voice test split, using GPU if available. The model expects 16kHz audio.
 
 ```python
-import torch
-import torchaudio
+from transformers import Wav2Vec2ForCTC, Wav2Vec2ProcessorWithLM
 from datasets import load_dataset
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-test_dataset = load_dataset("common_voice", "sv-SE", split="test[:2%]").
-processor = Wav2Vec2Processor.from_pretrained("KBLab/wav2vec2-large-voxrex-swedish")
-model = Wav2Vec2ForCTC.from_pretrained("KBLab/wav2vec2-large-voxrex-swedish")
-resampler = torchaudio.transforms.Resample(48_000, 16_000)
-# Preprocessing the datasets.
-# We need to read the aduio files as arrays
-def speech_file_to_array_fn(batch):
-    speech_array, sampling_rate = torchaudio.load(batch["path"])
-    batch["speech"] = resampler(speech_array).squeeze().numpy()
-    return batch
-test_dataset = test_dataset.map(speech_file_to_array_fn)
-inputs = processor(test_dataset["speech"][:2], sampling_rate=16_000, return_tensors="pt", padding=True)
+import torch
+import torchaudio.functional as F
+
+# Import model and processor
+model_name = 'viktor-enzell/wav2vec2-large-voxrex-swedish-4gram'
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = Wav2Vec2ForCTC.from_pretrained(model_name).to(device);
+processor = Wav2Vec2ProcessorWithLM.from_pretrained(model_name)
+
+# Import and process speech data 
+common_voice = load_dataset('common_voice', 'sv-SE', split='test[:1%]')
+
+def speech_file_to_array(sample):
+    # Convert speech file to array and downsample to 16 kHz
+    sampling_rate = sample['audio']['sampling_rate']
+    sample['speech'] = F.resample(torch.tensor(sample['audio']['array']), sampling_rate, 16_000)
+    return sample
+
+common_voice = common_voice.map(speech_file_to_array)
+
+# Run inference
+inputs = processor(common_voice['speech'], sampling_rate=16_000, return_tensors='pt', padding=True).to(device)
+
 with torch.no_grad():
-    logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
-predicted_ids = torch.argmax(logits, dim=-1)
-print("Prediction:", processor.batch_decode(predicted_ids))
-print("Reference:", test_dataset["sentence"][:2])
+    logits = model(**inputs).logits
+
+transcripts = processor.batch_decode(logits.cpu().numpy()).text
 ```
 
 ## Training procedure