whisper.py

import numpy as np
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration

# Whisperモデルとプロセッサのロード
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)