TensorVox/bert.cpp at master · ZDisket/TensorVox · 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
#include "bert.h"
#include <windows.h>

BERT::BERT()
{

}

BERT::BERT(const std::string &Path, const std::string &DictPath)
{
    Initialize(Path, DictPath);
}

void BERT::Initialize(const std::string &Path, const std::string &DictPath)
{
    Model = torch::jit::load(Path);
    Tokenizer = std::make_unique<FullTokenizer>(DictPath,true);


}

std::pair<TFTensor<float>, TFTensor<float> > BERT::Infer(const std::string &InText)
{
    torch::NoGradGuard no_grad;

    auto Tokens = Tokenizer->tokenize(InText + "\n");
    auto Ids = Tokenizer->convertTokensToIds(Tokens);

    std::vector<int32_t> InTokens(Ids.begin(),Ids.end());


    auto InIDS = torch::tensor(InTokens).unsqueeze(0); // (1, tokens)
    std::pair<TFTensor<float>,TFTensor<float>> BERTOutputs;

    try{
        auto Output = Model({InIDS}).toTuple(); // (hidden states, pooled)
        BERTOutputs.first = VoxUtil::CopyTensor<float>(Output.get()->elements()[0].toTensor());
        BERTOutputs.second = VoxUtil::CopyTensor<float>(Output.get()->elements()[1].toTensor());


    }

    catch (const std::exception& e) {
        int msgboxID = MessageBox(
                    NULL,
                    (LPCWSTR)QString::fromStdString(e.what()).toStdWString().c_str(),
                    (LPCWSTR)L"Error1!!",
                    MB_ICONWARNING | MB_CANCELTRYCONTINUE | MB_DEFBUTTON2
                    );


        return BERTOutputs;

    }


    return BERTOutputs;


}