34 lines
949 B
Python
34 lines
949 B
Python
import torch
|
|
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
|
|
from datasets import load_dataset
|
|
|
|
|
|
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
|
|
|
model_id = "openai/whisper-large-v3"
|
|
|
|
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
|
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True
|
|
)
|
|
model.to(device)
|
|
|
|
processor = AutoProcessor.from_pretrained(model_id)
|
|
|
|
pipe = pipeline(
|
|
"automatic-speech-recognition",
|
|
model=model,
|
|
tokenizer=processor.tokenizer,
|
|
feature_extractor=processor.feature_extractor,
|
|
chunk_length_s=30,
|
|
batch_size=16, # batch size for inference - set based on your device
|
|
torch_dtype=torch_dtype,
|
|
device=device,
|
|
)
|
|
|
|
dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
|
|
sample = dataset[0]["audio"]
|
|
|
|
result = pipe(sample)
|
|
print(result["text"])
|