Merge pull request #1499 from lainedfles/whisper_auto_update

feat: introduce Whisper model auto-update control.
This commit is contained in:
Timothy Jaeryang Baek 2024-04-14 13:58:17 -07:00 committed by GitHub
commit b5d882606a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 22 additions and 6 deletions

View file

@ -28,6 +28,7 @@ from config import (
UPLOAD_DIR,
WHISPER_MODEL,
WHISPER_MODEL_DIR,
WHISPER_MODEL_AUTO_UPDATE,
DEVICE_TYPE,
)
@ -69,12 +70,24 @@ def transcribe(
f.write(contents)
f.close()
model = WhisperModel(
WHISPER_MODEL,
device=whisper_device_type,
compute_type="int8",
download_root=WHISPER_MODEL_DIR,
)
whisper_kwargs = {
"model_size_or_path": WHISPER_MODEL,
"device": whisper_device_type,
"compute_type": "int8",
"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.warning(
"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)
log.info(

View file

@ -450,6 +450,9 @@ Query: [query]"""
WHISPER_MODEL = os.getenv("WHISPER_MODEL", "base")
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"
)
####################################