| 8 | #define MACRO_STRINGIFY(x) STRINGIFY(x) |
| 9 | |
| 10 | PYBIND11_MODULE(_core, m) { |
| 11 | m.doc() = R"pbdoc( |
| 12 | FlashBertTokenizer Python bindings |
| 13 | --------------------------------- |
| 14 | |
| 15 | .. currentmodule:: flash_tokenizer |
| 16 | |
| 17 | .. autosummary:: |
| 18 | :toctree: _generate |
| 19 | |
| 20 | FlashBertTokenizer |
| 21 | )pbdoc"; |
| 22 | |
| 23 | py::class_<FlashBertTokenizer>(m, "FlashBertTokenizer") |
| 24 | .def(py::init<const std::string &, bool, int>(), |
| 25 | py::arg("vocab_file"), |
| 26 | py::arg("do_lower_case") = true, |
| 27 | py::arg("max_input_chars_per_word") = 256) |
| 28 | .def("tokenize", [](FlashBertTokenizer &self, py::str text, int max_length) -> std::list<std::string> { |
| 29 | std::string text_str = py::str(text).cast<std::string>(); |
| 30 | return self.tokenize(text_str, max_length); |
| 31 | }, py::arg("text"), py::arg("max_length") = -1) |
| 32 | .def("convert_tokens_to_ids", &FlashBertTokenizer::convert_tokens_to_ids) |
| 33 | .def("convert_ids_to_tokens", &FlashBertTokenizer::convert_ids_to_tokens) |
| 34 | .def("__call__", |
| 35 | [](FlashBertTokenizer &self, |
| 36 | py::str text, |
| 37 | py::str padding, |
| 38 | int max_length, |
| 39 | py::str return_tensors, |
| 40 | bool truncation) -> std::vector<int> { |
| 41 | std::string text_str = text.cast<std::string>(); |
| 42 | std::string padding_str = padding.cast<std::string>(); |
| 43 | std::string return_tensors_str = return_tensors.cast<std::string>(); |
| 44 | return self(text_str, padding_str, max_length, return_tensors_str, truncation); |
| 45 | }, |
| 46 | py::arg("text"), |
| 47 | py::arg("padding") = "longest", |
| 48 | py::arg("max_length") = -1, |
| 49 | py::arg("return_tensors") = "np", |
| 50 | py::arg("truncation") = true); |
| 51 | |
| 52 | #ifdef VERSION_INFO |
| 53 | m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO); |
| 54 | #else |
| 55 | m.attr("__version__") = "dev"; |
| 56 | #endif |
| 57 | } |