This commit is contained in:
Timothy J. Baek 2024-09-19 21:51:45 +02:00
parent 2dad9b9432
commit cd117f5b67
2 changed files with 4 additions and 12 deletions

View File

@ -74,6 +74,10 @@ ENV RAG_EMBEDDING_MODEL="$USE_EMBEDDING_MODEL_DOCKER" \
## Hugging Face download cache ##
ENV HF_HOME="/app/backend/data/cache/embedding/models"
## Torch Extensions ##
ENV TORCH_EXTENSIONS_DIR="/app/backend/data/cache/torch_extensions"
#### Other models ##########################################################
WORKDIR /app/backend

View File

@ -206,18 +206,6 @@ def update_reranking_model(
print("ColBERT: Loading model", name)
self.device = "cuda" if torch.cuda.is_available() else "cpu"
if DOCKER:
# This is a workaround for the issue with the docker container
# where the torch extension is not loaded properly
# and the following error is thrown:
# /root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/segmented_maxsim_cpp.so: cannot open shared object file: No such file or directory
lock_file = "/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock"
if os.path.exists(lock_file):
os.remove(lock_file)
print("ColBERT: Removed lock file")
self.ckpt = Checkpoint(
name,
colbert_config=ColBERTConfig(model_name=name),