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, language: str = "en") -> str:
    # 言語設定用のトークナイズオプションを設定
    forced_decoder_ids = processor.tokenizer.get_decoder_prompt_ids(language=language, task="transcribe")

    input_features = processor(chunk, sampling_rate=SAMPLING_RATE, return_tensors="pt").input_features.to(device)
    predicted_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
    transcriptions = processor.batch_decode(predicted_ids, skip_special_tokens=True)
    print(transcriptions)
    return "\n".join(transcriptions)