tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Releases 2 Packages Activity Actions

Merge pull request #1499 from lainedfles/whisper_auto_update

Browse source 
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
Show stats Download patch file Download diff file Expand all files Collapse all files

25
backend/apps/audio/main.py
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,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(

3
backend/config.py
View file

@ -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"
)
#################################### ####################################