implement on gpu

t_res/geoparser/linking.py

@@ -112,6 +112,7 @@ def __init__(
         linking_resources: Optional[dict] = dict(),
         overwrite_training: Optional[bool] = False,
         rel_params: Optional[dict] = None,
+        rel_device: Optional[str] = None, 
     ):
         """
         Initialises a Linker object.
@@ -136,6 +137,7 @@ def __init__(
             }
 
         self.rel_params = rel_params
+        self.rel_device = rel_device
 
     def __str__(self) -> str:
         """
@@ -455,6 +457,8 @@ def train_load_model(
                     "mode": "train",
                     "model_path": os.path.join(linker_name, "model"),
                 }
+                if self.rel_device is not None:
+                    config_rel["device"] = self.rel_device
 
                 # Instantiate the entity disambiguation model:
                 model = entity_disambiguation.EntityDisambiguation(
@@ -476,6 +480,8 @@ def train_load_model(
                     "mode": "eval",
                     "model_path": os.path.join(linker_name, "model"),
                 }
+                if self.rel_device is not None:
+                    config_rel["device"] = self.rel_device
 
                 model = entity_disambiguation.EntityDisambiguation(
                     self.rel_params["db_embeddings"],

t_res/geoparser/recogniser.py

@@ -326,7 +326,7 @@ def create_pipeline(self) -> Pipeline:
             model_name = os.path.join(self.model_path, f"{self.model}.model")
 
         # Load a NER pipeline:
-        self.pipe = pipeline("ner", model=model_name, ignore_labels=[])
+        self.pipe = pipeline("ner", model=model_name, ignore_labels=[], device_map="auto")
         return self.pipe
 
     # -------------------------------------------------------------

t_res/utils/REL/entity_disambiguation.py

@@ -67,7 +67,7 @@ def __init__(self, db_embs, user_config, reset_embeddings=False):
         self.config = self.__get_config(user_config)
 
         # Use CPU if cuda is not available:
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = self.config.get("device", "cuda" if torch.cuda.is_available() else "cpu")
         self.prerank_model = None
         self.model = None
         self.reset_embeddings = reset_embeddings