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

cuda fix

Browse source
This commit is contained in:
Jannik Streidl 2024-03-22 12:48:48 +01:00
parent fc4e762b05
commit fdef2abdfb
4 changed files with 25 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
Dockerfile
View file

@ -36,7 +36,10 @@ ARG INCLUDE_OLLAMA
## Basis ##
ENV ENV=prod \
PORT=8080 \
INCLUDE_OLLAMA_ENV=${INCLUDE_OLLAMA}
# pass build args to the build
INCLUDE_OLLAMA_DOCKER=${INCLUDE_OLLAMA} \
USE_MPS_DOCKER=${USE_MPS} \
USE_CUDA_DOCKER=${USE_CUDA}
## Basis URL Config ##
ENV OLLAMA_BASE_URL="/ollama" \
@ -65,7 +68,7 @@ ENV RAG_EMBEDDING_MODEL="all-MiniLM-L6-v2" \
# Important:
# If you want to use CUDA you need to install the nvidia-container-toolkit (https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
# you can set this to "cuda" but its recomended to use --build-arg CUDA_ENABLED=true flag when building the image
RAG_EMBEDDING_MODEL_DEVICE_TYPE="cpu" \
# RAG_EMBEDDING_MODEL_DEVICE_TYPE="cpu" \
DEVICE_COMPUTE_TYPE="int8"
# device type for whisper tts and embbeding models - "cpu" (default), "cuda" (nvidia gpu and CUDA required) or "mps" (apple silicon) - choosing this right can lead to better performance
#### Preloaded models ##########################################################
@ -75,21 +78,18 @@ WORKDIR /app/backend
COPY ./backend/requirements.txt ./requirements.txt
RUN if [ "$USE_CUDA" = "true" ]; then \
export DEVICE_TYPE="cuda" && \
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117 --no-cache-dir && \
pip3 install -r requirements.txt --no-cache-dir; \
elif [ "$USE_MPS" = "true" ]; then \
export DEVICE_TYPE="mps" && \
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --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 chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device=os.environ['DEVICE_TYPE'])"; \
python -c "import os; from chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device='mps')"; \
else \
export DEVICE_TYPE="cpu" && \
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --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 chromadb.utils import embedding_functions; sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=os.environ['RAG_EMBEDDING_MODEL'], device=os.environ['DEVICE_TYPE'])"; \
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')"; \
fi

3
backend/apps/rag/main.py
View file

@ -71,7 +71,7 @@ from constants import ERROR_MESSAGES
# sentence_transformer_ef = SentenceTransformer(
# model_name_or_path=RAG_EMBEDDING_MODEL,
# cache_folder=RAG_EMBEDDING_MODEL_DIR,
# device=RAG_EMBEDDING_MODEL_DEVICE_TYPE,
# device=DEVICE_TYPE,
# )
@ -178,7 +178,6 @@ async def update_embedding_model(
device=DEVICE_TYPE,
)
)
return {
"status": True,
"embedding_model": app.state.RAG_EMBEDDING_MODEL,

21
backend/config.py
View file

@ -208,7 +208,7 @@ OLLAMA_API_BASE_URL = os.environ.get(
)
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "")
INCLUDE_OLLAMA = os.environ.get("INCLUDE_OLLAMA", "false")
INCLUDE_OLLAMA = os.environ.get("INCLUDE_OLLAMA_ENV", "false")
if OLLAMA_BASE_URL == "" and OLLAMA_API_BASE_URL != "":
@ -220,7 +220,7 @@ if OLLAMA_BASE_URL == "" and OLLAMA_API_BASE_URL != "":
if ENV == "prod":
if OLLAMA_BASE_URL == "/ollama":
if INCLUDE_OLLAMA == "true":
if INCLUDE_OLLAMA.lower() == "true":
# if you use all-in-one docker container (Open WebUI + Ollama)
# with the docker build arg INCLUDE_OLLAMA=true (--build-arg="INCLUDE_OLLAMA=true") this only works with http://localhost:11434
OLLAMA_BASE_URL = "http://localhost:11434"
@ -336,9 +336,20 @@ CHROMA_DATA_PATH = f"{DATA_DIR}/vector_db"
# this uses the model defined in the Dockerfile ENV variable. If you dont use docker or docker based deployments such as k8s, the default embedding model will be used (all-MiniLM-L6-v2)
RAG_EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
# device type ebbeding models - "cpu" (default), "cuda" (nvidia gpu required) or "mps" (apple silicon) - choosing this right can lead to better performance
DEVICE_TYPE = os.environ.get(
"DEVICE_TYPE", "cpu"
)
USE_CUDA = os.environ.get("USE_CUDA_DOCKER", "false")
USE_MPS = os.environ.get("USE_MPS_DOCKER", "false")
if USE_CUDA.lower() == "true" and USE_MPS.lower() == "true":
print("Both USE_CUDA and USE_MPS cannot be set to true. Defaulting to CPU.")
DEVICE_TYPE = "cpu"
elif USE_CUDA.lower() == "true":
DEVICE_TYPE = "cuda"
elif USE_MPS.lower() == "true":
DEVICE_TYPE = "mps"
else:
DEVICE_TYPE = "cpu"
CHROMA_CLIENT = chromadb.PersistentClient(
path=CHROMA_DATA_PATH,
settings=Settings(allow_reset=True, anonymized_telemetry=False),

2
backend/start.sh
View file

@ -2,7 +2,7 @@
# Get the INCLUDE_OLLAMA_ENV environment variable which is set in the Dockerfile
# This includes the ollama in the image
INCLUDE_OLLAMA=${INCLUDE_OLLAMA_ENV:-false}
INCLUDE_OLLAMA=${INCLUDE_OLLAMA_DOCKER}
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
cd "$SCRIPT_DIR" || exit