SOLD/test_encoder_decoder.py at main · SAIS-LifeScience/SOLD · 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
import warnings
warnings.filterwarnings("ignore")

from utils import *
from dataset import *
from torch.utils.data import DataLoader
from models import *
import torch.nn as nn

def test_encoder_decoder(config):
    print("Testing Encoder-Decoder model...")
    data_config = config['data_config']
    train_config = config['train_config']
    model_config = config['model_config']
    seed = config['seed']
    seeding(seed)

    llm_config = model_config['llm_config']
    llm_config['device'] = config['device']
    data_config['llm_config'] = llm_config

    test_dataset = GraphDataset(
        data_path=data_config['test_data_path'],
        data_config=data_config
    )

    loss_fn = nn.CrossEntropyLoss()

    test_dataloader = DataLoader(test_dataset,
                                 batch_size=train_config['batch_size'],
                                 shuffle=False,
                                 collate_fn=test_dataset.collate_fn)

    if model_config['decoder']['name'] == 'mlp':
        decoder = MLPDecoder(**model_config['decoder']['config'])
        if model_config.get("encoder"):
            if model_config["encoder"]["name"] == "mlp":
                encoder = MLPEncoder(**model_config["encoder"]["config"])
            else:
                raise NotImplementedError(model_config["encoder"]["name"])
            encoder_decoder = EncoderMlpDecoder(decoder=decoder, encoder=encoder).to(config['device'])
        else:
            encoder_decoder = EncoderMlpDecoder(decoder=decoder).to(config['device'])
    else:
        raise NotImplementedError(model_config['decoder']['name'])

    state_dict = torch.load(train_config['ckpt_path'], map_location='cpu')['encoder_decoder']

    encoder_decoder.load_state_dict(state_dict, strict=True)
    encoder_decoder.eval()

    test_loss, test_recovery_seq, test_recovery_aa = validate_encoder_decoder(encoder_decoder,
                                                                           config['device'],
                                                                           test_dataloader,
                                                                           loss_fn,
                                                                           'testing',
                                                                           0)
    print("test loss: {}, seq recovery {}, aa recovery {}".format(test_loss, test_recovery_seq, test_recovery_aa))
    print("testing complete ......")