From a602089b64cfd5abf2b6a447cbfeb9bdbe04de9a Mon Sep 17 00:00:00 2001 From: Jannik Streidl Date: Thu, 4 Apr 2024 10:06:29 +0200 Subject: [PATCH] Added preload if embedding model in CUDA mode --- Dockerfile | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index 737d478eca..e189152224 100644 --- a/Dockerfile +++ b/Dockerfile @@ -69,7 +69,8 @@ RUN if [ "$USE_CUDA" = "true" ]; then \ # If you use CUDA the whisper and embedding modell will be downloaded on first use pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/$USE_CUDA_DOCKER_VER --no-cache-dir && \ pip3 install -r requirements.txt --no-cache-dir && \ - 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'])"; \ + 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'])" && \ + python -c "import os; from chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device='cpu')"; \ else \ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --no-cache-dir && \ pip3 install -r requirements.txt --no-cache-dir && \