CUDA対応

Closes #18

.devcontainer/devcontainer.json

@@ -2,8 +2,12 @@
   "name": "search-research",
   "dockerComposeFile": "../docker-compose.yml",
   "service": "backend",
+  "runServices": ["backend"],
   "workspaceFolder": "/workspace/search-research",
   "postCreateCommand": "/bin/sh ./.devcontainer/postCreateCommand.sh",
+  "features": {
+    "ghcr.io/devcontainers/features/nvidia-cuda:1": {}
+  },
   "customizations": {
     "vscode": {
       "settings": {

backend/es/processors.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Any, Dict
 
 import japanese_clip as ja_clip
@@ -21,10 +22,13 @@ class JaClipEncodeProcessor(Processor):
     設定先は、target_fieldで指定
     """
 
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.INFO)
     _MODEL_NAME = "rinna/japanese-clip-vit-b-16"
 
     def __init__(self, target_field: str):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.logger.info(f"device is {self.device}")
         self.model, preprocess = ja_clip.load(self._MODEL_NAME, device=self.device)
         self.tokenizer = ja_clip.load_tokenizer()
         self.target_field = target_field

docker-compose.yml

@@ -1,4 +1,4 @@
-version: "2.2"
+version: "3.8"
 
 services:
   backend:
@@ -9,6 +9,13 @@ services:
     environment:
       - TZ=Asia/Tokyo
     command: sleep infinity
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
     volumes:
       - .:/workspace/search-research:cached
       - venv-search-research-backend:/workspace/search-research/.venv