src/infer.py

from typing import *
from src import utils
import whisper
  

MODEL_NAMES = ["tiny", "base", "small", "medium", "large", "large-v1","large-v2"]
MODEL_BASE = whisper.load_model("base")

def model_preloader_downloader():
    for mname in MODEL_NAMES:
        mdl = whisper.load_model(mname)
        del mdl

def model_loader(name: str):
    if name=="base":
        return MODEL_BASE
    else:
        return whisper.load_model(name)

def predict(model_name: str, language: str, mic_audio=None, audio=None):
    if mic_audio is not None:
        voice = mic_audio
    elif audio is not None:
        voice = audio
    else:
        return "(please provide audio)"
    
    voice = utils.preprocess_audio(voice)
    
    model = model_loader(model_name)
    result = model.transcribe(voice, language=language)
    return result["text"]