浏览代码

Added preload if embedding model in CUDA mode

Jannik Streidl 1 年之前
父节点
当前提交
a602089b64
共有 1 个文件被更改,包括 2 次插入1 次删除
  1. 2 1
      Dockerfile

+ 2 - 1
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 && \