Update utlis.py

utlis.py

@@ -127,17 +127,15 @@ 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,  length_penalty:float = 0.6):
+def translate(model: torch.nn.Module, strategy:str = 'greedy' , src_sentence: str, lenght_extend :int = 5, 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()
     # Tokenize the source sentence
     src = source_tokenizer(src_sentence, **token_config)['input_ids']
     num_tokens = src.shape[1]
     # Create a source mask
     src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
     if strategy == 'greedy':
-       tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + 5, start_symbol=target_tokenizer.bos_token_id).flatten()
+       tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + lenght_extend, start_symbol=target_tokenizer.bos_token_id).flatten()
     # Generate the target tokens using beam search decoding
     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()
@@ -163,4 +161,8 @@ token_config = {
                 "return_tensors": True,
              }
 
-load_model()
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+model = load_model()
+model.to(device)
+model.eval()