Unverified Commit b5d88260 authored by Timothy Jaeryang Baek's avatar Timothy Jaeryang Baek Committed by GitHub
Browse files

Merge pull request #1499 from lainedfles/whisper_auto_update

feat: introduce Whisper model auto-update control.
parents d9ce1d3e faa58841
...@@ -28,6 +28,7 @@ from config import ( ...@@ -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,24 @@ def transcribe( ...@@ -69,12 +70,24 @@ 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.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) segments, info = model.transcribe(file_path, beam_size=5)
log.info( log.info(
......
...@@ -450,6 +450,9 @@ Query: [query]""" ...@@ -450,6 +450,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"
)
#################################### ####################################
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment