Bert.h 2.18 KB
Newer Older
liucong's avatar
liucong committed
1
2
#ifndef __BERT_H__
#define __BERT_H__
3

liucong's avatar
liucong committed
4
5
#include <tokenization.h>

6
7
#include <cstdint>
#include <migraphx/program.hpp>
liucong's avatar
liucong committed
8
#include <string>
9

liucong's avatar
liucong committed
10
11
namespace migraphxSamples {
typedef enum _ErrorCode
12
{
liucong's avatar
liucong committed
13
14
15
16
17
18
    SUCCESS = 0,
    MODEL_NOT_EXIST,
    CONFIG_FILE_NOT_EXIST,
    FAIL_TO_LOAD_MODEL,
    FAIL_TO_OPEN_CONFIG_FILE,
} ErrorCode;
liucong's avatar
liucong committed
19

liucong's avatar
liucong committed
20
21
22
23
24
typedef struct _Sort_st
{
    int index;
    float value;
} Sort_st;
25

liucong's avatar
liucong committed
26
27
28
29
30
31
32
typedef struct _ResultOfPredictions
{
    int start_index;
    int end_index;
    float start_predictionvalue;
    float end_predictionvalue;
} ResultOfPredictions;
33
34
35

class Bert
{
liucong's avatar
liucong committed
36
    public:
37
    Bert();
liucong's avatar
liucong committed
38

39
40
    ~Bert();

liucong's avatar
liucong committed
41
    ErrorCode Initialize();
42

liucong's avatar
liucong committed
43
44
45
46
47
    ErrorCode Inference(const std::vector<std::vector<long unsigned int>>& input_ids,
                        const std::vector<std::vector<long unsigned int>>& input_masks,
                        const std::vector<std::vector<long unsigned int>>& segment_ids,
                        std::vector<float>& start_position,
                        std::vector<float>& end_position);
48
49

    ErrorCode Preprocessing(cuBERT::FullTokenizer tokenizer,
liucong's avatar
liucong committed
50
51
52
53
54
55
56
                            int batch_size,
                            int max_seq_length,
                            const char* text,
                            char* question,
                            std::vector<std::vector<long unsigned int>>& input_ids,
                            std::vector<std::vector<long unsigned int>>& input_masks,
                            std::vector<std::vector<long unsigned int>>& segment_ids);
57

liucong's avatar
liucong committed
58
59
60
61
62
    ErrorCode Postprocessing(int n_best_size,
                             int max_answer_length,
                             const std::vector<float>& start_position,
                             const std::vector<float>& end_position,
                             std::string& answer);
63

liucong's avatar
liucong committed
64
    private:
65
66
67
68
69
70
71
72
73
74
75
76
77
78
    std::vector<std::string> tokens_text;
    std::vector<std::string> tokens_question;

    migraphx::program net;
    std::string inputName1;
    std::string inputName2;
    std::string inputName3;
    std::string inputName4;
    migraphx::shape inputShape1;
    migraphx::shape inputShape2;
    migraphx::shape inputShape3;
    migraphx::shape inputShape4;
};

liucong's avatar
liucong committed
79
} // namespace migraphxSamples
80
81

#endif