sample_from_model.py

#!/usr/bin/env python
#  coding=utf-8

"""
Samples an existing RNN model.
"""

import argparse
import gzip
import functools

import tqdm

import models.model as mm
import models.actions as ma
import utils.torch as ut
import utils.log as ul


def parse_args():
    """Parses input arguments."""
    parser = argparse.ArgumentParser(description="Samples a model.")
    parser.add_argument("--model-path", "-m", help="Path to the model.", type=str, required=True)
    parser.add_argument("--output-smiles-path", "-o",
                        help="Path to the output file (if none given it will use stdout).", type=str)
    parser.add_argument("--num", "-n", help="Number of SMILES to sample [DEFAULT: 1024]", type=int, default=1024)
    parser.add_argument("--with-nll", help="Store the NLL in a column after the SMILES.",
                        action="store_true", default=False)
    parser.add_argument("--batch-size", "-b",
                        help="Batch size (beware GPU memory usage) [DEFAULT: 128]", type=int, default=128)
    parser.add_argument("--use-gzip", help="Compress the output file (if set).", action="store_true", default=False)

    return parser.parse_args()


def main():
    """Main function."""
    args = parse_args()

    ut.set_default_device("cuda")

    model = mm.Model.load_from_file(args.model_path, mode="eval")

    open_func = open
    if args.use_gzip:
        open_func = gzip.open
        args.output_smiles_path += ".gz"

    if args.output_smiles_path:
        csv_file = open_func(args.output_smiles_path, "wt+")
        write_func = functools.partial(csv_file.write)
    else:
        csv_file = tqdm.tqdm
        write_func = functools.partial(csv_file.write, end="")

    sample_model = ma.SampleModel(model, args.batch_size)

    for smi, nll in ul.progress_bar(sample_model.run(args.num), total=args.num):
        output_row = [smi]
        if args.with_nll:
            output_row.append("{:.8f}".format(nll))
        write_func("\t".join(output_row) + "\n")

    if args.output_smiles_path:
        csv_file.close()


LOG = ul.get_logger(name="sample_from_model")
if __name__ == "__main__":
    main()