From 7dbba8f7c0031d3d8d0fff1310c739eae4acb267 Mon Sep 17 00:00:00 2001 From: Lorenzo Ridolfi Date: Sat, 2 Dec 2023 11:15:07 -0300 Subject: [PATCH 1/3] Add autocast to optimize GPU memory usage --- .DS_Store | Bin 0 -> 8196 bytes docker/.DS_Store | Bin 0 -> 6148 bytes easynmt/.DS_Store | Bin 0 -> 6148 bytes easynmt/models/AutoModel.py | 46 +++++++++++++++++++++++++++++++----- models/.DS_Store | Bin 0 -> 6148 bytes 5 files changed, 40 insertions(+), 6 deletions(-) create mode 100644 .DS_Store create mode 100644 docker/.DS_Store create mode 100644 easynmt/.DS_Store create mode 100644 models/.DS_Store diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..be49690981791dcbead9713606df57c11433996f GIT binary patch literal 8196 zcmeHMPj3=I6n|4F&@?6-G^sH$n|P&KLaZ2Lf=G{g5TYj9Aj__F)7>3pK{N#2=uPod zJ^2m%23}3Pd+Nc&58%adaECUW}>p$jon<_ukBVzu7;$1pr8O;1mH?03e5r zVSX9AeT3;*n~?=MauSiCKY-&zZQrKUpRqLqh5^HXVZbn87%&X{8w}txn@MiYb6;pp zZx}EP{Fe-{`-6pzVMSsrkl#A6<0SxM8i%>yG1dXx$C6l)7z^aCIH#x{h;${g#2~^Q z^%gTntVoOn3U?sF9f(|+$O?t<)zQ!5=0GX}O>Y=549qgXYWGF3VE`WZdi|~p`<-IF zO@prRpM24s0$QutKDBW)_v+%)bN61nzm&9a4zTEMsVjV73uV9o3T@a17kYXGZ(b+$ zY7E}HC({_TIRg4Kig`9#9E2yF#r0bd!9xfkfWAG9x-#3hcGvf!{@U6uSxBW9Gg*}7 z>}Gw74(t6^e-Kq$y<6;C17C^Q)8m$X-E)V{{L#xa?6*otrdA&3qFXfJ%I&`ke8HH$i!` zgd^hnUO&O3L|W4~H_RoW(z)OHHfy1>zx(*aS9t~`=27{dZ>8-3@G~JX&;5EY^l8|^ zI9RlK*pT~k6|O)7PUC)c@Ux1ibO?8W>hYg_`8z&pMxUPLo15aCK5#8H$3OAYvyL%D zd?32;05c1qiD-QUpke5(yzOFl2y{$6{yZ0T7SC$mg?Z z6Qy*l5ZIFMyZGns@+FF6BI4=YazZpFq6{iH*u$_v#4p;Dfmy`Kvc_a~QI&Ns&JDkp z0eN;qI;Wa$imUeXJ9=Mx$|}znMYVt>J3M=QJAHmRZhMzoy!CZ^)mNn9WxAw$>{(*h z0`uGM&hc|oO#2KEA5Y(B=&e;&>(gwD>z->a=W=ldoB?Oxzc7HBEt2gk`sfTe1J1yT z0r@@zs9-c~6w{{zLuvtlBbbw5F1>{01jA_9D8d6_O$BNyTZ_S(4tub;XxJ!fIt<8 literal 0 HcmV?d00001 diff --git a/easynmt/.DS_Store b/easynmt/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..547ad617b25f6e60149fac91eccf7b2e8f6c8eee GIT binary patch literal 6148 zcmeHKy-ou$47Q=5lUO=-%p=e@2-UH$c40)xX79NY2VBsy8`5YV6qbIQ- zgpe&IpW~n8@=2G(M8vCyen~VZq6STnMVSyWPr8mQ_yovxjzqhDw@F>jX=0+kXp(O~ zrGYZ_bW3~Rj}46->2A5+{u>!;_7!v>X<*&aN?-wy))npj2SrA;au+j zOZ+mEMgBO&N6vsV@Xr|FS-Wmm*p%I^2iudoHlW?0iAY=#1pKIyzresdd>J(X literal 0 HcmV?d00001 diff --git a/easynmt/models/AutoModel.py b/easynmt/models/AutoModel.py index 7b84530..04364d8 100644 --- a/easynmt/models/AutoModel.py +++ b/easynmt/models/AutoModel.py @@ -1,15 +1,25 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer import torch +from torch.cuda.amp import autocast, GradScaler from typing import List import logging - logger = logging.getLogger(__name__) class AutoModel: def __init__(self, model_name: str, tokenizer_name: str = None, easynmt_path: str = None, lang_map=None, tokenizer_args=None): + """ + Initializes an instance of the AutoModel class. + + Args: + model_name (str): The name or path of the pre-trained model to be used for translation. + tokenizer_name (str, optional): The name or path of the tokenizer associated with the pre-trained model. Defaults to None. + easynmt_path (str, optional): The path to the EasyNMT model if the model_name or tokenizer_name is set to ".". Defaults to None. + lang_map (dict, optional): A dictionary mapping language codes to specific language codes used by the tokenizer. Defaults to None. + tokenizer_args (dict, optional): Additional arguments to be passed to the tokenizer. Defaults to None. + """ if tokenizer_args is None: tokenizer_args = {} @@ -30,10 +40,24 @@ def __init__(self, model_name: str, tokenizer_name: str = None, easynmt_path: st self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name) self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, **self.tokenizer_args) - self.max_length = None + self.max_length = 512 # Set a smaller value for low memory GPUs def translate_sentences(self, sentences: List[str], source_lang: str, target_lang: str, device: str, beam_size: int = 5, **kwargs): + """ + Translates a list of sentences from a source language to a target language. + + Args: + sentences (List[str]): The list of sentences to be translated. + source_lang (str): The source language of the sentences. + target_lang (str): The target language for translation. + device (str): The device to be used for translation (e.g. "cuda"). + beam_size (int, optional): The beam size for translation. Defaults to 5. + **kwargs: Additional keyword arguments to be passed to the translation model. + + Returns: + List[str]: A list of translated sentences. + """ self.model.to(device) if source_lang in self.lang_map: @@ -49,14 +73,24 @@ def translate_sentences(self, sentences: List[str], source_lang: str, target_lan inputs[key] = inputs[key].to(device) with torch.no_grad(): - if hasattr(self.tokenizer, 'lang_code_to_id'): - kwargs['forced_bos_token_id'] = self.tokenizer.lang_code_to_id[target_lang] - translated = self.model.generate(**inputs, num_beams=beam_size, **kwargs) - output = [self.tokenizer.decode(t, skip_special_tokens=True) for t in translated] + with autocast(): + if hasattr(self.tokenizer, 'lang_code_to_id'): + kwargs['forced_bos_token_id'] = self.tokenizer.lang_code_to_id[target_lang] + translated = self.model.generate(**inputs, num_beams=beam_size, **kwargs) + output = [self.tokenizer.decode(t, skip_special_tokens=True) for t in translated] return output def save(self, output_path): + """ + Saves the model and tokenizer to the specified output path. + + Args: + output_path (str): The path to save the model and tokenizer. + + Returns: + dict: A dictionary containing the saved model and tokenizer information. + """ self.model.save_pretrained(output_path) self.tokenizer.save_pretrained(output_path) return { diff --git a/models/.DS_Store b/models/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..517979944906e994401721bfbf09294e22a6c9ed GIT binary patch literal 6148 zcmeHK%}&BV5S~>MiFhEMJej_Lgf^NOPlePc;G!p>kZ4S6iNL*g=HGJs19eWZKGLXe4V{y8+K$DsPr5zwejhb zORW-H21Ef-;I}EjYj=lY^5DE~v-RsyIG&E8Fi9dT?rKtW;r2#JJkgb@zO#-SCqWQK zNeE8;LX|yYC=-rx$MSrOsX-G?${s$Hy|c0> z6s7Nu`5jFslk@K(GB*M@5jcYwad)F4J+@<%|+AdM*SqYAtOT{5Se literal 0 HcmV?d00001 From 15d3b0f00525bddaef0cf76d6dad21b0f1dbd6ad Mon Sep 17 00:00:00 2001 From: Lorenzo Ridolfi Date: Sat, 2 Dec 2023 11:19:54 -0300 Subject: [PATCH 2/3] Update .gitignore to exclude .DS_Store --- .gitignore | 1 + docker/.DS_Store | Bin 6148 -> 0 bytes easynmt/.DS_Store | Bin 6148 -> 0 bytes models/.DS_Store | Bin 6148 -> 0 bytes 4 files changed, 1 insertion(+) delete mode 100644 docker/.DS_Store delete mode 100644 easynmt/.DS_Store delete mode 100644 models/.DS_Store diff --git a/.gitignore b/.gitignore index 5b613a7..001f58a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ *.pt *.model tmp/ +.DS_Store # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/docker/.DS_Store b/docker/.DS_Store deleted file mode 100644 index 61dedc4aefc22c0a1cfdd0d7a10afa402d90d588..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6148 zcmeHKy-veG47S@2K`b2^?-lw0h@%Qm&;>D-QUpke5(yzOFl2y{$6{yZ0T7SC$mg?Z z6Qy*l5ZIFMyZGns@+FF6BI4=YazZpFq6{iH*u$_v#4p;Dfmy`Kvc_a~QI&Ns&JDkp z0eN;qI;Wa$imUeXJ9=Mx$|}znMYVt>J3M=QJAHmRZhMzoy!CZ^)mNn9WxAw$>{(*h z0`uGM&hc|oO#2KEA5Y(B=&e;&>(gwD>z->a=W=ldoB?Oxzc7HBEt2gk`sfTe1J1yT z0r@@zs9-c~6w{{zLuvtlBbbw5F1>{01jA_9D8d6_O$BNyTZ_S(4tub;XxJ!fIt<8 diff --git a/easynmt/.DS_Store b/easynmt/.DS_Store deleted file mode 100644 index 547ad617b25f6e60149fac91eccf7b2e8f6c8eee..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6148 zcmeHKy-ou$47Q=5lUO=-%p=e@2-UH$c40)xX79NY2VBsy8`5YV6qbIQ- zgpe&IpW~n8@=2G(M8vCyen~VZq6STnMVSyWPr8mQ_yovxjzqhDw@F>jX=0+kXp(O~ zrGYZ_bW3~Rj}46->2A5+{u>!;_7!v>X<*&aN?-wy))npj2SrA;au+j zOZ+mEMgBO&N6vsV@Xr|FS-Wmm*p%I^2iudoHlW?0iAY=#1pKIyzresdd>J(X diff --git a/models/.DS_Store b/models/.DS_Store deleted file mode 100644 index 517979944906e994401721bfbf09294e22a6c9ed..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6148 zcmeHK%}&BV5S~>MiFhEMJej_Lgf^NOPlePc;G!p>kZ4S6iNL*g=HGJs19eWZKGLXe4V{y8+K$DsPr5zwejhb zORW-H21Ef-;I}EjYj=lY^5DE~v-RsyIG&E8Fi9dT?rKtW;r2#JJkgb@zO#-SCqWQK zNeE8;LX|yYC=-rx$MSrOsX-G?${s$Hy|c0> z6s7Nu`5jFslk@K(GB*M@5jcYwad)F4J+@<%|+AdM*SqYAtOT{5Se From c99ddec3fc9440cf1889bf4387730fb2e9a8b435 Mon Sep 17 00:00:00 2001 From: lorenzoridolfi Date: Sat, 2 Dec 2023 12:04:32 -0300 Subject: [PATCH 3/3] Adding a parameter to control the optional usage of autocast --- .DS_Store | Bin 8196 -> 8196 bytes easynmt/models/AutoModel.py | 15 +++++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/.DS_Store b/.DS_Store index be49690981791dcbead9713606df57c11433996f..2e777f3ad3e4bb6ea3c3222bc7c90cb5d225d606 100644 GIT binary patch delta 21 dcmZp1XmQwJE5tN)=VW`K9L5Km*9b-M0svS}2j2hy delta 21 dcmZp1XmQwJE5tPU-(-8C9L5Wq*9b-M0svbO2oeAQ diff --git a/easynmt/models/AutoModel.py b/easynmt/models/AutoModel.py index 04364d8..a05c72f 100644 --- a/easynmt/models/AutoModel.py +++ b/easynmt/models/AutoModel.py @@ -1,6 +1,6 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer import torch -from torch.cuda.amp import autocast, GradScaler +from torch.cuda.amp import autocast from typing import List import logging @@ -43,7 +43,7 @@ def __init__(self, model_name: str, tokenizer_name: str = None, easynmt_path: st self.max_length = 512 # Set a smaller value for low memory GPUs - def translate_sentences(self, sentences: List[str], source_lang: str, target_lang: str, device: str, beam_size: int = 5, **kwargs): + def translate_sentences(self, sentences: List[str], source_lang: str, target_lang: str, device: str, beam_size: int = 5, with_autocast: bool = False, **kwargs): """ Translates a list of sentences from a source language to a target language. @@ -53,6 +53,7 @@ def translate_sentences(self, sentences: List[str], source_lang: str, target_lan target_lang (str): The target language for translation. device (str): The device to be used for translation (e.g. "cuda"). beam_size (int, optional): The beam size for translation. Defaults to 5. + with_autocast (bool, optional): Whether to use autocast for translation. Defaults to False. **kwargs: Additional keyword arguments to be passed to the translation model. Returns: @@ -73,11 +74,17 @@ def translate_sentences(self, sentences: List[str], source_lang: str, target_lan inputs[key] = inputs[key].to(device) with torch.no_grad(): - with autocast(): + if with_autocast: + with autocast(): + if hasattr(self.tokenizer, 'lang_code_to_id'): + kwargs['forced_bos_token_id'] = self.tokenizer.lang_code_to_id[target_lang] + translated = self.model.generate(**inputs, num_beams=beam_size, **kwargs) + else: if hasattr(self.tokenizer, 'lang_code_to_id'): kwargs['forced_bos_token_id'] = self.tokenizer.lang_code_to_id[target_lang] translated = self.model.generate(**inputs, num_beams=beam_size, **kwargs) - output = [self.tokenizer.decode(t, skip_special_tokens=True) for t in translated] + + output = [self.tokenizer.decode(t, skip_special_tokens=True) for t in translated] return output