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