Introduce Whisper model auto-update control.

* Introduce WHISPER_MODEL_AUTO_UPDATE env var
* Pass local_files_only to WhisperModel()
* Handle cases where auto-update is disabled but model is non-existent
This commit is contained in:
Self Denial 2024-04-10 20:30:00 -06:00
parent 9e726a32e6
commit 429242b4d3
2 changed files with 20 additions and 6 deletions

View file

@ -28,6 +28,7 @@ from config import (
UPLOAD_DIR, UPLOAD_DIR,
WHISPER_MODEL, WHISPER_MODEL,
WHISPER_MODEL_DIR, WHISPER_MODEL_DIR,
WHISPER_MODEL_AUTO_UPDATE,
DEVICE_TYPE, DEVICE_TYPE,
) )
@ -69,12 +70,22 @@ def transcribe(
f.write(contents) f.write(contents)
f.close() f.close()
model = WhisperModel( whisper_kwargs = {
WHISPER_MODEL, "model_size_or_path": WHISPER_MODEL,
device=whisper_device_type, "device": whisper_device_type,
compute_type="int8", "compute_type": "int8",
download_root=WHISPER_MODEL_DIR, "download_root": WHISPER_MODEL_DIR,
) "local_files_only": not WHISPER_MODEL_AUTO_UPDATE,
}
log.debug(f"whisper_kwargs: {whisper_kwargs}")
try:
model = WhisperModel(**whisper_kwargs)
except:
log.debug("WhisperModel initialization failed, attempting download with local_files_only=False")
whisper_kwargs["local_files_only"] = False
model = WhisperModel(**whisper_kwargs)
segments, info = model.transcribe(file_path, beam_size=5) segments, info = model.transcribe(file_path, beam_size=5)
log.info( log.info(

View file

@ -446,6 +446,9 @@ Query: [query]"""
WHISPER_MODEL = os.getenv("WHISPER_MODEL", "base") WHISPER_MODEL = os.getenv("WHISPER_MODEL", "base")
WHISPER_MODEL_DIR = os.getenv("WHISPER_MODEL_DIR", f"{CACHE_DIR}/whisper/models") WHISPER_MODEL_DIR = os.getenv("WHISPER_MODEL_DIR", f"{CACHE_DIR}/whisper/models")
WHISPER_MODEL_AUTO_UPDATE = (
os.environ.get("WHISPER_MODEL_AUTO_UPDATE", "").lower() == "true"
)
#################################### ####################################