Merge pull request #4 from atrifat/feat-enable-cuda-support

feat: enable cuda support

Author: atrifat

.env.example

@@ -10,4 +10,7 @@ DETOXIFY_MODEL=unbiased-small
 ENABLE_CACHE=false
 
 # (Optional. Default: 60. Options: integer value) Duration of cache in seconds. Parameter will be used when the cache is enabled
-CACHE_DURATION_SECONDS=60
+CACHE_DURATION_SECONDS=60
+
+# (Optional. Default: auto. Options: auto,cpu,cuda) Set torch default device for detoxify library. Automatically detect if cuda/gpu device is available
+TORCH_DEVICE=auto

app.py

@@ -4,6 +4,7 @@
 from flask import Flask, request, jsonify
 import functools
 import datetime
+import torch
 from detoxify import Detoxify
 import logging
 from flask_caching import Cache
@@ -18,6 +19,7 @@
 DETOXIFY_MODEL = os.getenv("DETOXIFY_MODEL", "unbiased-small")
 CACHE_DURATION_SECONDS = int(os.getenv("CACHE_DURATION_SECONDS", 60))
 ENABLE_CACHE = os.getenv("ENABLE_CACHE", "false") == "true"
+TORCH_DEVICE = os.getenv("TORCH_DEVICE", "auto")
 
 APP_VERSION = "0.1.0"
 
@@ -40,7 +42,12 @@
 if ENABLE_API_TOKEN and API_TOKEN == "":
     raise Exception("API_TOKEN is required if ENABLE_API_TOKEN is enabled")
 
-model = Detoxify(DETOXIFY_MODEL)
+if TORCH_DEVICE == "auto":
+    torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+else:
+    torch_device = TORCH_DEVICE
+
+model = Detoxify(DETOXIFY_MODEL, device=torch_device)
 
 app = Flask(__name__)