Skip to content

Commit

Permalink
update gpu build to use cuda
Browse files Browse the repository at this point in the history
  • Loading branch information
vemonet committed Sep 14, 2023
1 parent c4fee14 commit 8474c6a
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 8 deletions.
4 changes: 2 additions & 2 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -257,10 +257,10 @@ jobs:
with:
context: .
build-args: |
BASE_IMAGE=nvcr.io/nvidia/pytorch:23.06-py3
BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu22.04
GPU_ENABLED=true
# BASE_IMAGE=nvcr.io/nvidia/pytorch:23.06-py3
# BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu20.04
# BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu22.04
push: ${{ github.event_name != 'pull_request' }}
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
Expand Down
10 changes: 6 additions & 4 deletions Dockerfile
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
ARG BASE_IMAGE=python:3.11
# ARG BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu20.04
# ARG BASE_IMAGE=nvcr.io/nvidia/cuda:12.2.0-devel-ubuntu22.04
# 2.7GB cf. https://ngc.nvidia.com/catalog/containers/nvidia:cuda
# ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:23.06-py3
# 8.5GB cf. https://ngc.nvidia.com/catalog/containers/nvidia:pytorch

# CUDA image: https://github.com/oobabooga/text-generation-webui/blob/main/docker/Dockerfile

FROM ${BASE_IMAGE}

ARG DEBIAN_FRONTEND=noninteractive
Expand All @@ -12,14 +14,14 @@ ENV LIBRECHAT_WORKERS=1

# CUDA image required to install python
RUN apt-get update && \
apt-get install -y software-properties-common wget unzip && \
apt-get install -y software-properties-common git vim build-essential python3-dev wget unzip && \
# add-apt-repository ppa:deadsnakes/ppa && \
# apt-get install -y python3.11 && \
# ln -s /usr/bin/python3.11 /usr/bin/python && \
# wget https://bootstrap.pypa.io/get-pip.py && \
# python get-pip.py && \
# rm get-pip.py && \
pip install --upgrade pip
pip3 install --upgrade pip



Expand All @@ -37,7 +39,7 @@ WORKDIR /app

# Pre-install requirements to use cache when re-building
ADD scripts/requirements.txt .
RUN pip install -r requirements.txt && \
RUN pip3 install -r requirements.txt && \
rm requirements.txt

ADD . .
Expand Down
4 changes: 2 additions & 2 deletions src/libre_chat/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def stream(
input_text: str, memory: Any, instructions: str, temperature: float, max_new_tokens: int
) -> Iterator[Tuple[Any, str]]:
# Create a Queue
q: Queue[Any] = Queue()
q: "Queue[Any]" = Queue()
job_done = object()

# Create a function to call - this will run in a thread
Expand Down Expand Up @@ -240,7 +240,7 @@ def clear_chat() -> List[str]:
class StreamGradioCallback(BaseCallbackHandler):
"""Callback handler for streaming LLM responses to a queue."""

def __init__(self, q: Queue[Any]):
def __init__(self, q: "Queue[Any]"):
self.q = q

def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
Expand Down

0 comments on commit 8474c6a

Please sign in to comment.