diff --git a/docker_sample/Dockerfile b/docker_sample/Dockerfile index d857bc3..47eb57f 100644 --- a/docker_sample/Dockerfile +++ b/docker_sample/Dockerfile @@ -11,7 +11,9 @@ RUN apt-get update && \ && rm -rf /var/lib/apt/lists/* # vLLMのインストール -RUN pip install --no-cache-dir vllm +RUN pip install --no-cache-dir vllm nvidia-ml-py && \ + pip uninstall --yes pynvml && \ + pip install -U --no-cache-dir opencv-python-headless==4.5.5.62 # モデルファイルをイメージに取り込む (FIXME: モデルファイルのパスを変更) COPY ./models/llm-jp/llm-jp-3-1.8b-instruct /workspace/model diff --git a/docker_sample/README.md b/docker_sample/README.md index 11c8dac..15bcd96 100644 --- a/docker_sample/README.md +++ b/docker_sample/README.md @@ -72,7 +72,6 @@ sudo docker load < [チームID].tar ```bash docker run --rm \ - --runtime=nvidia \ --gpus all \ --network none \ -v "$(pwd)/data:/data" \ diff --git a/docker_sample/src/main.py b/docker_sample/src/main.py index 78c5d20..effbc40 100644 --- a/docker_sample/src/main.py +++ b/docker_sample/src/main.py @@ -3,6 +3,8 @@ import argparse from vllm import LLM +# see more examples at https://github.com/vllm-project/vllm/tree/main/examples +chat_template = "{{ (messages|selectattr('role', 'equalto', 'system')|list|last).content|trim if (messages|selectattr('role', 'equalto', 'system')|list) else '' }}\n\n{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ message['content']|trim -}}\n{% if not loop.last %}\n\n\n{% endif %}\n{% elif message['role'] == 'assistant' %}\n{{ message['content']|trim -}}\n{% if not loop.last %}\n\n\n{% endif %}\n{% elif message['role'] == 'user_context' %}\n{{ message['content']|trim -}}\n{% if not loop.last %}\n\n\n{% endif %}\n{% endif %}\n{% endfor %}\n" def main(): parser = argparse.ArgumentParser() @@ -26,9 +28,9 @@ def main(): for d in data: messages_list.append({"role": "user", "content": d["text"]}) - outputs = llm.chat(messages_list) - for i, output in enumerate(outputs): - data[i]["response"] = output.outputs[0].text + for i, m in enumerate(messages_list): + output = llm.chat([m], chat_template=chat_template) + data[i]["response"] = output[0].outputs[0].text.strip() os.makedirs(os.path.dirname(args.output_path), exist_ok=True) with open(args.output_path, "w", encoding="utf-8") as f: