KU-Hackathon-DL-Server/test.py at main · ryuni-dev/KU-Hackathon-DL-Server · 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
import torch
import numpy as np
from kobert.utils import get_tokenizer
from kobert.pytorch_kobert import get_pytorch_kobert_model
import gluonnlp as nlp
from dataset.dataset import BERTDataset


def predict(predict_sentence, model):
    device = 'cpu'

    batch_size = 64
    max_len = 64
    _, vocab = get_pytorch_kobert_model()
    tokenizer = get_tokenizer()
    tok=nlp.data.BERTSPTokenizer(tokenizer, vocab, lower = False)

    data = [predict_sentence, '0']
    dataset_another = [data]

    another_test = BERTDataset(dataset_another, 0, 1, tok, max_len, True, False)
    test_dataloader = torch.utils.data.DataLoader(another_test, batch_size=batch_size, num_workers=2)

    model.eval()

    for _, (token_ids, valid_length, segment_ids, label) in enumerate(test_dataloader):
        token_ids = token_ids.long().to(device)
        segment_ids = segment_ids.long().to(device)

        valid_length= valid_length
        label = label.long().to(device)

        out = model(token_ids, valid_length, segment_ids)


        test_eval=[]
        for i in out:
            logits=i
            logits = logits.detach().cpu().numpy()

            #HAPPINESS, ANXIOUS, PANIC, SADNESS, ANGER, HURT
            if np.argmax(logits) == 0:
                test_eval.append("HAPPINESS")
            elif np.argmax(logits) == 1:
                test_eval.append("ANXIOUS")
            elif np.argmax(logits) == 2:
                test_eval.append("PANIC")
            elif np.argmax(logits) == 3:
                test_eval.append("SADNESS")
            elif np.argmax(logits) == 4:
                test_eval.append("ANGER")
            elif np.argmax(logits) == 5:
                test_eval.append("HURT")

    return test_eval[0]