File size: 583 Bytes
745e35f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from transformers import AutoModelForCTC, AutoProcessor
from datasets import load_dataset
import torch
dummy_dataset = load_dataset("common_voice", "ab", split="test")
model = AutoModelForCTC.from_pretrained("hf-internal-testing/tiny-random-wav2vec2")
model.to("cuda")
processor = AutoProcessor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2")
input_values = processor(dummy_dataset[0]["audio"]["array"], return_tensors="pt", sampling_rate=16_000).input_values
input_values = input_values.to("cuda")
logits = model(input_values).logits
assert logits.shape[-1] == 32 |