| 22 | #include "turbo_transformers/core/config.h" |
| 23 | |
| 24 | static bool test(const std::string &model_path, bool use_cuda = false) { |
| 25 | // construct a bert model using n_layers and n_heads, |
| 26 | // the hidden_size can be infered from the parameters |
| 27 | BertModel model(model_path, |
| 28 | use_cuda ? DLDeviceType::kDLGPU : DLDeviceType::kDLCPU, |
| 29 | 12, /* n_layers */ |
| 30 | 12 /* *n_heads */); |
| 31 | std::vector<std::vector<int64_t>> position_ids{{1, 0, 0, 0}, {1, 1, 1, 0}}; |
| 32 | std::vector<std::vector<int64_t>> segment_ids{{1, 1, 1, 0}, {1, 0, 0, 0}}; |
| 33 | auto vec = model({{12166, 10699, 16752, 4454}, {5342, 16471, 817, 16022}}, |
| 34 | position_ids, segment_ids, PoolType::kFirst, |
| 35 | true /* use a pooler after the encoder output */); |
| 36 | // bert-base-uncased (2020.04.23 version), you may need to change it to |
| 37 | assert(fabs(vec.data()[0] - -0.5503) < 1e-3); |
| 38 | assert(fabs(vec.data()[1] - 0.1295) < 1e-3); |
| 39 | assert(fabs(vec.data()[768] - -0.5545) < 1e-3); |
| 40 | assert(fabs(vec.data()[768 + 1] - -0.1182) < 1e-3); |
| 41 | |
| 42 | return true; |
| 43 | } |
| 44 | |
| 45 | static std::vector<float> CallBackFunction( |
| 46 | const std::shared_ptr<BertModel> model, |