| 46 | |
| 47 | |
| 48 | void register_encoder(py::module& m) { |
| 49 | py::class_<EncoderForwardOutput>(m, "EncoderForwardOutput", |
| 50 | "Forward output of an encoder model.") |
| 51 | |
| 52 | .def_readonly("last_hidden_state", &EncoderForwardOutput::last_hidden_state, |
| 53 | "Output of the last layer.") |
| 54 | .def_readonly("pooler_output", &EncoderForwardOutput::pooler_output, |
| 55 | "Output of the pooling layer.") |
| 56 | |
| 57 | .def("__repr__", [](const EncoderForwardOutput& output) { |
| 58 | return "EncoderForwardOutput(last_hidden_state=" |
| 59 | + std::string(py::repr(py::cast(output.last_hidden_state))) |
| 60 | + ", pooler_output=" + std::string(py::repr(py::cast(output.pooler_output))) |
| 61 | + ")"; |
| 62 | }) |
| 63 | ; |
| 64 | |
| 65 | py::class_<EncoderWrapper>( |
| 66 | m, "Encoder", |
| 67 | R"pbdoc( |
| 68 | A text encoder. |
| 69 | |
| 70 | Example: |
| 71 | |
| 72 | >>> encoder = ctranslate2.Encoder("model/", device="cpu") |
| 73 | >>> encoder.forward_batch([["▁Hello", "▁world", "!"]]) |
| 74 | )pbdoc") |
| 75 | |
| 76 | .def(py::init<const std::string&, const std::string&, const std::variant<int, std::vector<int>>&, const StringOrMap&, size_t, size_t, long, bool, bool, py::object>(), |
| 77 | py::arg("model_path"), |
| 78 | py::arg("device")="cpu", |
| 79 | py::kw_only(), |
| 80 | py::arg("device_index")=0, |
| 81 | py::arg("compute_type")="default", |
| 82 | py::arg("inter_threads")=1, |
| 83 | py::arg("intra_threads")=0, |
| 84 | py::arg("max_queued_batches")=0, |
| 85 | py::arg("flash_attention")=false, |
| 86 | py::arg("tensor_parallel")=false, |
| 87 | py::arg("files")=py::none(), |
| 88 | R"pbdoc( |
| 89 | Initializes the encoder. |
| 90 | |
| 91 | Arguments: |
| 92 | model_path: Path to the CTranslate2 model directory. |
| 93 | device: Device to use (possible values are: cpu, cuda, auto). |
| 94 | device_index: Device IDs where to place this encoder on. |
| 95 | compute_type: Model computation type or a dictionary mapping a device name |
| 96 | to the computation type (possible values are: default, auto, int8, int8_float32, |
| 97 | int8_float16, int8_bfloat16, int16, float16, bfloat16, float32). |
| 98 | inter_threads: Maximum number of parallel generations. |
| 99 | intra_threads: Number of OpenMP threads per encoder (0 to use a default value). |
| 100 | max_queued_batches: Maximum numbers of batches in the queue (-1 for unlimited, |
| 101 | 0 for an automatic value). When the queue is full, future requests will block |
| 102 | until a free slot is available. |
| 103 | flash_attention: run model with flash attention 2 for self-attention layer |
| 104 | tensor_parallel: run model with tensor parallel mode |
| 105 | files: Load model files from the memory. This argument is a dictionary mapping |
no test coverage detected