From fcf2001f24cec9961ca0dac882ad95f072ba5f97 Mon Sep 17 00:00:00 2001 From: ZHANG Hao Date: Wed, 3 Sep 2025 11:01:20 +0800 Subject: [PATCH] add support for a100 --- Dockerfile.a100 | 5 +++++ run_in_docker_a100.sh | 3 +++ test.sh | 2 ++ 3 files changed, 10 insertions(+) create mode 100644 Dockerfile.a100 create mode 100755 run_in_docker_a100.sh create mode 100755 test.sh diff --git a/Dockerfile.a100 b/Dockerfile.a100 new file mode 100644 index 0000000..9284543 --- /dev/null +++ b/Dockerfile.a100 @@ -0,0 +1,5 @@ +FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel + +WORKDIR /workspace +RUN pip install sentence-transformers -i https://nexus.4pd.io/repository/pypi-all/simple +COPY . /workspace/ diff --git a/run_in_docker_a100.sh b/run_in_docker_a100.sh new file mode 100755 index 0000000..c3f8d31 --- /dev/null +++ b/run_in_docker_a100.sh @@ -0,0 +1,3 @@ +#! /usr/bin/env bash +image=harbor-contest.4pd.io/zhanghao/feature:a100-0.1 +docker run -it -v /home/zhanghao/workspace:/workspace -v /mnt:/mnt $image bash diff --git a/test.sh b/test.sh new file mode 100755 index 0000000..772b756 --- /dev/null +++ b/test.sh @@ -0,0 +1,2 @@ +# python main.py dataset.json output.json /mnt/contest_ceph/zhanghao/models/BAAI/bge-large-zh-v1.5 --device npu +python main.py dataset.json output.json /mnt/contest_ceph/zhanghao/models/BAAI/bge-large-zh-v1.5 --device cuda