From 66b104d26a97d7a232f0018ada5adee07d1c78ef Mon Sep 17 00:00:00 2001
From: cocomany <124849750+cocomany@users.noreply.github.com>
Date: Fri, 28 Apr 2023 11:24:59 +0800
Subject: [PATCH] supports GPU usage in docker (#196)

Co-authored-by: Huaiyong <huaiyong.sun@hpe.com>
---
 Dockerfile-cuda | 13 +++++++++++++
 README.md       | 20 +++++++++++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)
 create mode 100644 Dockerfile-cuda

diff --git a/Dockerfile-cuda b/Dockerfile-cuda
new file mode 100644
index 0000000..b6408db
--- /dev/null
+++ b/Dockerfile-cuda
@@ -0,0 +1,13 @@
+FROM  nvidia/cuda:12.1.0-runtime-ubuntu20.04
+LABEL MAINTAINER="chatGLM"
+
+COPY . /chatGLM/
+
+WORKDIR /chatGLM
+
+RUN apt-get update -y && apt-get install python3 python3-pip curl -y  && apt-get clean
+RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py 
+
+RUN pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn && rm -rf `pip3 cache dir`
+
+CMD ["python3","-u", "webui.py"]
\ No newline at end of file
diff --git a/README.md b/README.md
index 3c389d7..7c4a8a9 100644
--- a/README.md
+++ b/README.md
@@ -37,12 +37,22 @@
     本项目中默认选用的 Embedding 模型 [GanymedeNil/text2vec-large-chinese](https://huggingface.co/GanymedeNil/text2vec-large-chinese/tree/main) 约占用显存 3GB，也可修改为在 CPU 中运行。
 
 ## Docker 部署
-
-```commandline
-$ docker build -t chatglm:v1.0 .
-
-$ docker run -d --restart=always --name chatglm -p 7860:7860 -v /www/wwwroot/code/langchain-ChatGLM:/chatGLM  chatglm
+为了能让容器使用主机GPU资源，需要在主机上安装 [NVIDIA Container Toolkit](https://github.com/NVIDIA/nvidia-container-toolkit)。具体安装步骤如下：
+```shell
+sudo apt-get update
+sudo apt-get install -y nvidia-container-toolkit-base
+sudo systemctl daemon-reload 
+sudo systemctl restart docker
 ```
+安装完成后，可以使用以下命令编译镜像和启动容器：
+```
+docker build -t chatglm-cuda:latest .
+docker run --gpus all -d --name chatglm -p 7860:7860  chatglm-cuda:latest
+
+#若要使用离线模型，请配置好模型路径，然后此repo挂载到Container
+docker run --gpus all -d --name chatglm -p 7860:7860 -v ~/github/langchain-ChatGLM:/chatGLM  chatglm-cuda:latest
+```
+
 
 ## 开发部署