| #!/usr/bin/env python3 | |
| from transformers import WhisperForCausalLM, WhisperForConditionalGeneration, WhisperProcessor | |
| import torch | |
| from datasets import load_dataset | |
| processor = WhisperProcessor.from_pretrained("openai/whisper-large-v2") | |
| model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v2") | |
| assistant_model = WhisperForCausalLM.from_pretrained("patrickvonplaten/whisper-large-v2-32-2") | |
| ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") | |
| sample = ds[0]["audio"] | |
| input_features = processor(sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt").input_features | |
| predicted_ids = model.generate(input_features, assistant_model=assistant_model) | |
| # decode token ids to text | |
| transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True) | |
| print(transcription) | |