run.py

from src.model import Model 
from argparse import ArgumentParser
import torch, torch.cuda
import os
from src.utils import _build_config_function

def test_model(config):
    """
    Tests the given model. It generates a sample starting with the seed string, mimicking the style
    of the text on which it was trained.

    @param config (dict): dictionary with all the configurations of the model.
    """
    model = Model(config['testing_params'], config['model_layers'], device=device)

    config = config['testing_params']
    config['dropout_prob'] = 0
    model.load(config['--from_path'])
    model.load_text_characters(config['--testing_corpus'])
    print(config['seed'] + model.sample(config['seed']))
    print(model.test(config['n_timesteps'],batch_size=16))
   
def train_model(config):
    """
    Trains the given model based on the configuration file. 
    Read about customizing the training on the README.md file.

    @param config (dict): dictionary with all the configurations of the model.
    """
    model = Model(config['training_params'], config['model_layers'], device=device)
    
    # Load text into model.train_data and model.test_data:
    if config['training_params']['character_level'] == True:
        model.load_text_characters(config['training_params']['--corpus'])
    else:
        model.load_text_words(config['training_params']['--corpus'])
    
    config = config['training_params']
    model.train(config['n_iter'],
                        config['n_timesteps'],
                        config['batch_size'],
                        config['learning_rate'],
                        config['regularization'],
                        config['patience'])

def fine_tune(config):
    """
    Fine-tunes the model cited on --from_path based on the configuration file. 
    Read about customizing the training on the README.md file.

    @param config (dict): dictionary with all the configurations of the model.
    """
    model = Model(config['fine_tuning_params'], config['model_layers'], device=device)

    model.load(config['fine_tuning_params']['--from_path'])

    # Load text into model.train_data and model.test_data:
    if config['fine_tuning_params']['character_level'] == True:
        model.load_text_characters(config['fine_tuning_params']['--corpus'])
    else:
        model.load_text_words(config['fine_tuning_params']['--corpus'])

    config = config['fine_tuning_params']
    model.train(config['n_iter'],
                        config['n_timesteps'],
                        config['batch_size'],
                        config['learning_rate'],
                        config['regularization'],
                        config['patience'])

def parse_arguments():
    """
    Parses and returns arguments from the terminal.

    @returns Namespace, arguments class.
    """
    parser = ArgumentParser(description='configuration of runtime application')
    parser.add_argument('--train', action='store_true',
                        help='train the model with provided config file and text corpus')
    parser.add_argument('--fine_tune', action='store_true',
                        help='train the model with provided config file and text corpus')
    parser.add_argument('--test', action='store_true',
                        help='test the model with provided text sample_size (default = 300) and seed')

    parser.add_argument('--config', nargs='?', type=_build_config_function, default=f"{PATH}/config.py",
                        help='path to configuration file for fine tuning/training the model')


    args = parser.parse_args()

    return args

# Get local PATH:
PATH = os.getcwd()

# Search for CUDA GPU:
if torch.cuda.is_available():
    cuda_device = torch.device("cuda:0")
    device = cuda_device
    print ("Device: cuda")
    torch.cuda.set_device(device)
else:
    device = 'cpu'
    print ("CUDA device not found, using CPU")

# Parse terminal arguments:
args = parse_arguments()

if args.train:
    config = args.config(args, device, PATH)
    train_model(config)

if args.fine_tune:
    config = args.config(args, device, PATH)
    fine_tune(config)
    
if args.test:
    config = args.config(args, device, PATH)
    test_model(config)