| 58 | } |
| 59 | |
| 60 | bool CheckCppBertWithPooler(bool use_cuda, bool only_input) { |
| 61 | BertModel model("models/bert.npz", |
| 62 | use_cuda ? DLDeviceType::kDLGPU : DLDeviceType::kDLCPU, 12, |
| 63 | 12); |
| 64 | std::vector<std::vector<int64_t>> position_ids{{1, 0, 0, 0}, {1, 1, 1, 0}}; |
| 65 | std::vector<std::vector<int64_t>> segment_ids{{1, 1, 1, 0}, {1, 0, 0, 0}}; |
| 66 | if (only_input) { |
| 67 | position_ids.clear(); |
| 68 | segment_ids.clear(); |
| 69 | } |
| 70 | auto vec = model({{12166, 10699, 16752, 4454}, {5342, 16471, 817, 16022}}, |
| 71 | position_ids, segment_ids, PoolType::kFirst, |
| 72 | /*use_pooler*/ true); |
| 73 | REQUIRE(vec.size() == 768 * 2); |
| 74 | // Write a better UT |
| 75 | for (size_t i = 0; i < vec.size(); ++i) { |
| 76 | REQUIRE(!std::isnan(vec.data()[i])); |
| 77 | REQUIRE(!std::isinf(vec.data()[i])); |
| 78 | } |
| 79 | if (only_input) { |
| 80 | std::cerr << vec.data()[0] << std::endl; |
| 81 | REQUIRE(fabs(vec.data()[0] - 0.9671) < 1e-3); |
| 82 | REQUIRE(fabs(vec.data()[1] - 0.9860) < 1e-3); |
| 83 | REQUIRE(fabs(vec.data()[768] - 0.9757) < 1e-3); |
| 84 | REQUIRE(fabs(vec.data()[768 + 1] - 0.9794) < 1e-3); |
| 85 | } else { |
| 86 | REQUIRE(fabs(vec.data()[0] - 0.9151) < 1e-3); |
| 87 | REQUIRE(fabs(vec.data()[1] - 0.5919) < 1e-3); |
| 88 | REQUIRE(fabs(vec.data()[768] - 0.9802) < 1e-3); |
| 89 | REQUIRE(fabs(vec.data()[768 + 1] - 0.9321) < 1e-3); |
| 90 | } |
| 91 | return true; |
| 92 | } |
| 93 | |
| 94 | TEST_CASE("Bert", "Cpp interface") { |
| 95 | CheckCppBert(false /*use_cuda*/, true /* only_input*/); |
no test coverage detected