From 78bbcd305ea511d4e4591048a11cc2f7a1b41b84 Mon Sep 17 00:00:00 2001 From: Rosie Wood Date: Tue, 18 Jun 2024 15:07:51 +0100 Subject: [PATCH 1/3] implement on gpu --- t_res/geoparser/linking.py | 6 ++++++ t_res/geoparser/recogniser.py | 2 +- t_res/utils/REL/entity_disambiguation.py | 2 +- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/t_res/geoparser/linking.py b/t_res/geoparser/linking.py index 1731a903..bb21749b 100644 --- a/t_res/geoparser/linking.py +++ b/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"], diff --git a/t_res/geoparser/recogniser.py b/t_res/geoparser/recogniser.py index 975f0010..18787e58 100644 --- a/t_res/geoparser/recogniser.py +++ b/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 # ------------------------------------------------------------- diff --git a/t_res/utils/REL/entity_disambiguation.py b/t_res/utils/REL/entity_disambiguation.py index dd3c0dfa..68a1c2cf 100644 --- a/t_res/utils/REL/entity_disambiguation.py +++ b/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 From 9933152dde1a0ff59525582186812e93469eb662 Mon Sep 17 00:00:00 2001 From: Rosie Wood Date: Tue, 18 Jun 2024 15:29:40 +0100 Subject: [PATCH 2/3] use device instead of device_map in NER --- t_res/geoparser/recogniser.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/t_res/geoparser/recogniser.py b/t_res/geoparser/recogniser.py index 18787e58..214a4f0a 100644 --- a/t_res/geoparser/recogniser.py +++ b/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=[], device_map="auto") + self.pipe = pipeline("ner", model=model_name, ignore_labels=[], device=self.device) return self.pipe # ------------------------------------------------------------- From 0d5fd011b2c659245a0dd7435ac8e9f6176a91dc Mon Sep 17 00:00:00 2001 From: Rosie Wood Date: Wed, 19 Jun 2024 12:04:08 +0100 Subject: [PATCH 3/3] add ner device arg to pipeline --- t_res/geoparser/pipeline.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/t_res/geoparser/pipeline.py b/t_res/geoparser/pipeline.py index 371c8e05..2d3230cf 100644 --- a/t_res/geoparser/pipeline.py +++ b/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.