feat: Allow configuring device (gpu/cpu) via env vars

bdikit/config.py

@@ -0,0 +1,12 @@
+import os
+import torch
+
+
+BDIKIT_DEVICE: str = os.getenv("BDIKIT_DEVICE", default="cpu")
+
+
+def get_device() -> str:
+    if BDIKIT_DEVICE == "auto":
+        return "cuda" if torch.cuda.is_available() else "cpu"
+    else:
+        return BDIKIT_DEVICE

bdikit/mapping_algorithms/scope_reducing/_algorithms/contrastive_learning/cl_api.py

@@ -1,6 +1,6 @@
 import os
 from typing import List, Dict, Tuple, Optional
-
+from bdikit.config import get_device
 import numpy as np
 import pandas as pd
 import torch
@@ -41,7 +41,7 @@ def __init__(
 
         self.unlabeled = PretrainTableDataset()
         self.batch_size = batch_size
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = get_device()
         self.model = self.load_checkpoint()
 
     def load_checkpoint(self, lm: str = "roberta"):

bdikit/mapping_algorithms/value_mapping/algorithms.py

@@ -8,6 +8,11 @@
 from autofj import AutoFJ
 from Levenshtein import ratio
 import pandas as pd
+import flair
+import torch
+from bdikit.config import get_device
+
+flair.device = torch.device(get_device())
 
 
 class ValueMatch(NamedTuple):

examples/getting-started.ipynb

@@ -21,10 +21,7 @@
    "outputs": [],
    "source": [
     "import bdikit as bdi\n",
-    "import pandas as pd\n",
-    "\n",
-    "import flair, torch\n",
-    "flair.device = torch.device(\"cpu\") "
+    "import pandas as pd"
    ]
   },
   {