Add GPU option to T-Res

README.md

@@ -16,7 +16,7 @@ The three components are used in combination in the **Pipeline** class.
 
 We also provide the code to deploy T-Res as an API, and show how to use it. Each of these elements are described in the documentation.
 
-The repository contains the code for the experiments described in our paper.
+The repository contains the code for the experiments described in [our CHR 2023 paper](https://ceur-ws.org/Vol-3558/paper4426.pdf).
 
 ## Documentation
 
@@ -130,4 +130,4 @@ In our experiments, we have used resources built from Wikidata and Wikipedia for
 
 ## Cite
 
-**[TODO]**
+Coll-Ardanuy, Mariona, Federico Nanni, Kaspar Beelen, & Luke Hare. '[The Past is a Foreign Place: Improving Toponym Linking for Historical Newspapers](https://ceur-ws.org/Vol-3558/paper4426.pdf).' Proceedings of the Computational Humanities Research Conference 2023, ed. by Artjoms Šeļa, Fotis Jannidis, and Iza Romanowska. Paris, France, December 6-8, 2023. https://ceur-ws.org/Vol-3558/ ISSN, 1613, 0073. 

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/pipeline.py

@@ -78,6 +78,7 @@ def __init__(
         mylinker: Optional[linking.Linker] = None,
         resources_path: Optional[str] = None,
         experiments_path: Optional[str] = None,
+        ner_device: Optional[str] = None,
     ):
         """
         Instantiates a Pipeline object.
@@ -92,6 +93,7 @@ def __init__(
             self.myner = recogniser.Recogniser(
                 model="Livingwithmachines/toponym-19thC-en",
                 load_from_hub=True,
+                device=ner_device,
             )
 
         # If myranker is None, instantiate the default Ranker.

t_res/geoparser/recogniser.py

@@ -94,6 +94,7 @@ def __init__(
         overwrite_training: Optional[bool] = False,
         do_test: Optional[bool] = False,
         load_from_hub: Optional[bool] = False,
+        device: Optional[str] = None,
     ):
         """
         Initialises a Recogniser object.
@@ -108,6 +109,7 @@ def __init__(
         self.overwrite_training = overwrite_training
         self.do_test = do_test
         self.load_from_hub = load_from_hub
+        self.device = device
 
         # Add "_test" to the model name if do_test is True, unless
         # the model is downloaded from Huggingface, in which case
@@ -322,11 +324,11 @@ def create_pipeline(self) -> Pipeline:
         # If the model is local (has not been obtained from the hub),
         # pre-append the model path and the extension of the model
         # to obtain the model name.
-        if self.load_from_hub == False:
+        if self.load_from_hub is False:
             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=self.device)
         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