|
import numpy as np |
|
import torch |
|
from transformers import WhisperProcessor, WhisperForConditionalGeneration |
|
|
|
|
|
model_name = "openai/whisper-base" |
|
processor = WhisperProcessor.from_pretrained(model_name) |
|
model = WhisperForConditionalGeneration.from_pretrained(model_name) |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
model.to(device) |
|
|
|
SAMPLING_RATE = 16000 |
|
|
|
|
|
def transcribe(chunk: np.ndarray) -> str: |
|
input_features = processor(chunk, sampling_rate=SAMPLING_RATE, return_tensors="pt").input_features.to(device) |
|
predicted_ids = model.generate(input_features) |
|
transcriptions = processor.batch_decode(predicted_ids, skip_special_tokens=True) |
|
print(transcriptions) |
|
return "\n".join(transcriptions) |
|
|
