Update utlis.py

utlis.py

@@ -1,6 +1,17 @@
 import yaml 
+import torch
+from .model import Seq2SeqTransformer
+from transformers import AutoTokenizer
+from transformers import PreTrainedTokenizerFast
+from tokenizers.processors import TemplateProcessing
 
-def load_checkpoint(model_checkpoint_dir='model.pt',config_dir='config.yaml'):
+
+def addPreprocessing(tokenizer):
+      tokenizer._tokenizer.post_processor = TemplateProcessing(
+          single=tokenizer.bos_token + " $A " + tokenizer.eos_token,
+          special_tokens=[(tokenizer.eos_token, tokenizer.eos_token_id), (tokenizer.bos_token, tokenizer.bos_token_id)])
+
+def load_model(model_checkpoint_dir='model.pt',config_dir='config.yaml'):
     
     with open(config_dir, 'r') as yaml_file:
         loaded_model_params = yaml.safe_load(yaml_file)
@@ -116,7 +127,7 @@ def beam_search_decode(model, src, src_mask, max_len, start_symbol, beam_size ,l
     best_beam = beams[0][0]
     return best_beam
 
-def translate(model: torch.nn.Module, strategy:str, src_sentence: str, lenght_extend :int = 0, beam_size: int = 5, raw: bool = False, length_penalty:float = 0.6):
+def translate(model: torch.nn.Module, strategy:str, src_sentence: str, lenght_extend :int = 0, beam_size: int = 5,  length_penalty:float = 0.6):
     assert strategy in ['greedy','beam search'], 'the strategy for decoding has to be either greedy or beam search'
     model.to(device)
     model.eval()
@@ -131,4 +142,25 @@ def translate(model: torch.nn.Module, strategy:str, src_sentence: str, lenght_ex
     else:
         tgt_tokens = beam_search_decode(model, src, src_mask, max_len=num_tokens + lenght_extend, start_symbol=target_tokenizer.bos_token_id, beam_size=beam_size,length_penalty=length_penalty).flatten()
     # Decode the target tokens and clean up the result
-    return target_tokenizer.decode(tgt_tokens, clean_up_tokenization_spaces=True, skip_special_tokens=True)
+    return target_tokenizer.decode(tgt_tokens, clean_up_tokenization_spaces=True, skip_special_tokens=True)
+
+special_tokens = {'unk_token':"[UNK]",
+                  'cls_token':"[CLS]",
+                  'eos_token': '[EOS]',
+                  'sep_token':"[SEP]",
+                  'pad_token':"[PAD]",
+                  'mask_token':"[MASK]",
+                  'bos_token':"[BOS]"}
+
+source_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased", **special_tokens)
+target_tokenizer = PreTrainedTokenizerFast.from_pretrained('Sifal/E2KT')
+
+addPreprocessing(source_tokenizer)
+addPreprocessing(target_tokenizer)
+
+token_config = {
+                "add_special_tokens": True,
+                "return_tensors": True,
+             }
+
+load_model()