This commit is contained in:
Timothy J. Baek 2024-02-14 23:32:54 -08:00
parent 8ab0338e71
commit 63e6f39b31
3 changed files with 8 additions and 9 deletions

View file

@ -31,8 +31,8 @@ ENV SCARF_NO_ANALYTICS true
ENV DO_NOT_TRACK true ENV DO_NOT_TRACK true
#Whisper TTS Settings #Whisper TTS Settings
ENV WHISPER_DIR="/app/backend/data/cache/whisper/models"
ENV WHISPER_MODEL="base" ENV WHISPER_MODEL="base"
ENV WHISPER_MODEL_DIR="/app/backend/data/cache/whisper/models"
WORKDIR /app/backend WORKDIR /app/backend
@ -49,7 +49,7 @@ RUN apt-get update \
&& rm -rf /var/lib/apt/lists/* && rm -rf /var/lib/apt/lists/*
# RUN python -c "from sentence_transformers import SentenceTransformer; model = SentenceTransformer('all-MiniLM-L6-v2')" # RUN python -c "from sentence_transformers import SentenceTransformer; model = SentenceTransformer('all-MiniLM-L6-v2')"
RUN python -c "import os; from faster_whisper import WhisperModel; WhisperModel(os.environ['WHISPER_MODEL'], device='cpu', compute_type='int8', download_root=os.environ['WHISPER_DIR'])" RUN python -c "import os; from faster_whisper import WhisperModel; WhisperModel(os.environ['WHISPER_MODEL'], device='cpu', compute_type='int8', download_root=os.environ['WHISPER_MODEL_DIR'])"
# copy embedding weight from build # copy embedding weight from build

View file

@ -21,7 +21,7 @@ from utils.utils import (
) )
from utils.misc import calculate_sha256 from utils.misc import calculate_sha256
from config import CACHE_DIR, UPLOAD_DIR, WHISPER_MODEL_NAME from config import CACHE_DIR, UPLOAD_DIR, WHISPER_MODEL, WHISPER_MODEL_DIR
app = FastAPI() app = FastAPI()
app.add_middleware( app.add_middleware(
@ -54,14 +54,11 @@ def transcribe(
f.write(contents) f.write(contents)
f.close() f.close()
model_name = os.getenv('WHISPER_MODEL', WHISPER_MODEL_NAME)
download_root = os.getenv('WHISPER_DIR', f"{CACHE_DIR}/whisper/models")
model = WhisperModel( model = WhisperModel(
model_name, WHISPER_MODEL,
device="cpu", device="cpu",
compute_type="int8", compute_type="int8",
download_root=download_root, download_root=WHISPER_MODEL_DIR,
) )
segments, info = model.transcribe(file_path, beam_size=5) segments, info = model.transcribe(file_path, beam_size=5)

View file

@ -136,4 +136,6 @@ CHUNK_OVERLAP = 100
#################################### ####################################
# Transcribe # Transcribe
#################################### ####################################
WHISPER_MODEL_NAME = "base"
WHISPER_MODEL = os.getenv("WHISPER_MODEL", "base")
WHISPER_MODEL_DIR = os.getenv("WHISPER_MODEL_DIR", f"{CACHE_DIR}/whisper/models")