Upload load_model.py

load_model.py

@@ -0,0 +1,50 @@
+from typing import Tuple
+
+import torch
+from transformer import get_model, Transformer
+from config import load_config, get_weights_file_path
+from train import get_local_dataset_tokenizer
+from tokenizer import get_or_build_local_tokenizer
+
+from tokenizers import Tokenizer
+   
+
+def load_model_tokenizer(
+    config,
+    device = torch.device('cpu'),
+) -> Tuple[Transformer, Tokenizer, Tokenizer]:
+    """
+    Loads a local model and tokenizer from a given config
+    """
+    if config['model']['preload'] is None:
+        raise ValueError('Unspecified preload model')
+    
+    src_tokenizer = get_or_build_local_tokenizer(
+        config=config, 
+        ds=None, 
+        lang=config['dataset']['src_lang'], 
+        tokenizer_type=config['dataset']['src_tokenizer']
+    )
+    tgt_tokenizer = get_or_build_local_tokenizer(
+        config=config, 
+        ds=None, 
+        lang=config['dataset']['tgt_lang'], 
+        tokenizer_type=config['dataset']['tgt_tokenizer']
+    )
+
+    model = get_model(
+        config,
+        src_tokenizer.get_vocab_size(), 
+        tgt_tokenizer.get_vocab_size(),
+    ).to(device)
+
+    model_filename = get_weights_file_path(config, config['model']['preload'])
+    state = torch.load(model_filename, map_location=device)
+    model.load_state_dict(state['model_state_dict'])
+
+    print(f'Finish loading model and tokenizers')
+    return (model, src_tokenizer, tgt_tokenizer)
+
+if __name__ == '__main__':
+    config = load_config(file_name='config/config_final.yaml')
+    model, src_tokenizer, tgt_tokenizer = load_model_tokenizer(config)