chem-encdec/t5_infer.py at main · cambridgeltl/chem-encdec · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
import os

from datasets import load_dataset
from rdkit import Chem

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import argparse
import random

from timeit import default_timer as timer
import datetime

import numpy as np
import torch
from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer,
                          GenerationConfig, set_seed)

from utils_other import generate_batch_custom

from tokenization.tokenization_utils import smi_tokenizer_spaces, simple_spaces
from generation.generation_utils import set_generation_config
from evaluation.evaluation_utils import evaluate_batch

device = "cuda" if torch.cuda.is_available() else "cpu"

def load_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, default=None)
    parser.add_argument("--test_path", type=str, required=True, default=None)
    parser.add_argument("--seed",type=int, required=False, default=42)
    parser.add_argument("--smiles_check",type=str, required=False, default="no", choices=['yes', 'no'])

    parser.add_argument("--num_beams", type=int, default=1)
    parser.add_argument("--top_k", type=int, default=210)
    parser.add_argument("--top_p", type=float, default=1.0)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--num_return_sequences",type=int, default=1)
    parser.add_argument("--max_new_tokens", type=int, default=200, required=False)
    parser.add_argument("--max_length", type=int, default=400, required=True)
    parser.add_argument("--penalty_alpha", type=float, default=0.0, required=False)
    parser.add_argument("--diversity_penalty", type=float, default=1.0, required=False)
    parser.add_argument("--length_penalty", type=float, default=1.0, required=False)
    parser.add_argument("--num_beam_groups", type=int, default=1, required=False)

    parser.add_argument("--infer_batch", type=int, default=1)
    parser.add_argument("--final_n", type=int, default=1)
    parser.add_argument("--load_in_8bit", type=bool, default=False, action=argparse.BooleanOptionalAction)
    parser.add_argument("--infer_mode", type=str, default="greedy", required=True, choices=['greedy', 'search_beam', 'diverse_beam', 'contrastive', 'sampling_beam', 'nucleus'])

    parser.add_argument("--tokenization", type=str, required=False, default="none", choices=['none','map','shrink', 'map_shrink', 'spaces', 'shrink_spaces', 'simple_spaces'])

    args = parser.parse_args()

    return args


def main():

    # Set arguments
    args = load_arguments()
    print(args)

    # Set the seed
    np.random.seed(args.seed)
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    set_seed(args.seed)

    # Set the generation configuration
    generation_config = set_generation_config(args)

    # Load the test dataset now
    dataset = load_dataset("csv", data_files={"test": [args.test_path]})["test"]

    if args.load_in_8bit:
        model = AutoModelForSeq2SeqLM.from_pretrained(args.model_path, load_in_8bit=True)
    else:
        model = AutoModelForSeq2SeqLM.from_pretrained(args.model_path).to(device)

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)

    start_time = timer()
    in_gold_mol_pairs = []

    model.eval()

    spaces=False

    for i in range(len(dataset)):

        orig_input, orig_output = dataset[i]["Input"], dataset[i]["Output"]
        if args.tokenization == "none":
            tokenized_input, tokenized_output = orig_input, orig_output
        elif args.tokenization == "spaces":
            tokenized_input = smi_tokenizer_spaces(orig_input)
            tokenized_output = orig_output
            spaces = True
        elif args.tokenization == "simple_spaces":
            tokenized_input = simple_spaces(orig_input)
            tokenized_output = orig_output
            spaces = True
        else:
            print("ERROR: Unsupported preprocessing tokenization!")
            sys.exit(-1)

        in_gold_mol_pair = (tokenized_input, tokenized_output)
        in_gold_mol_pairs.append(in_gold_mol_pair)

    in_gold_mol_batches = generate_batch_custom(in_gold_mol_pairs, args.infer_batch)

    all_predictions = []
    all_golds = []

    for in_gold_mol_batch in in_gold_mol_batches:
        batch_inputs = [x[0] for x in in_gold_mol_batch]
        batch_golds = [x[1] for x in in_gold_mol_batch]

        inputs = tokenizer(batch_inputs, return_tensors="pt", max_length=args.max_length, padding=True, truncation=True).to(device)
        with torch.no_grad():
            outputs = model.generate(**inputs, generation_config=generation_config)
            output_ids = outputs.sequences
            batch_pred_mols = tokenizer.batch_decode(output_ids, skip_special_tokens = True)

        ex_batch_pred_mols = [batch_pred_mols[j:j+args.num_return_sequences] for j in range(0, len(batch_pred_mols), args.num_return_sequences)]

        all_predictions.append(ex_batch_pred_mols)
        all_golds.append(batch_golds)

    # Finally, run evaluation over different K-s

    for K in [1,2,3,5]:
        total = 0
        correct = 0
        for batch_golds, ex_batch_pred_mols in zip(all_golds, all_predictions):
            correct_batch, total_batch = evaluate_batch(batch_golds, ex_batch_pred_mols, K, spaces)
            correct += correct_batch
            total += total_batch

        acc = float(correct)/total
        print("K="+str(K), "|||", "Correct / Total:", correct, "/", total, "["+str(round(acc*100,2))+"%]")

    end_time = timer()
    print("Processing time:" + str(datetime.timedelta(seconds=(end_time-start_time))))


if __name__ == "__main__":
    main()