| 272 | |
| 273 | |
| 274 | void register_translator(py::module& m) { |
| 275 | py::class_<TranslatorWrapper>( |
| 276 | m, "Translator", |
| 277 | R"pbdoc( |
| 278 | A text translator. |
| 279 | |
| 280 | Example: |
| 281 | |
| 282 | >>> translator = ctranslate2.Translator("model/", device="cpu") |
| 283 | >>> translator.translate_batch([["▁Hello", "▁world", "!"]]) |
| 284 | )pbdoc") |
| 285 | |
| 286 | .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>(), |
| 287 | py::arg("model_path"), |
| 288 | py::arg("device")="cpu", |
| 289 | py::kw_only(), |
| 290 | py::arg("device_index")=0, |
| 291 | py::arg("compute_type")="default", |
| 292 | py::arg("inter_threads")=1, |
| 293 | py::arg("intra_threads")=0, |
| 294 | py::arg("max_queued_batches")=0, |
| 295 | py::arg("flash_attention")=false, |
| 296 | py::arg("tensor_parallel")=false, |
| 297 | py::arg("files")=py::none(), |
| 298 | R"pbdoc( |
| 299 | Initializes the translator. |
| 300 | |
| 301 | Arguments: |
| 302 | model_path: Path to the CTranslate2 model directory. |
| 303 | device: Device to use (possible values are: cpu, cuda, auto). |
| 304 | device_index: Device IDs where to place this generator on. |
| 305 | compute_type: Model computation type or a dictionary mapping a device name |
| 306 | to the computation type (possible values are: default, auto, int8, int8_float32, |
| 307 | int8_float16, int8_bfloat16, int16, float16, bfloat16, float32). |
| 308 | inter_threads: Maximum number of parallel translations. |
| 309 | intra_threads: Number of OpenMP threads per translator (0 to use a default value). |
| 310 | max_queued_batches: Maximum numbers of batches in the queue (-1 for unlimited, |
| 311 | 0 for an automatic value). When the queue is full, future requests will block |
| 312 | until a free slot is available. |
| 313 | flash_attention: run model with flash attention 2 for self-attention layer |
| 314 | tensor_parallel: run model with tensor parallel mode |
| 315 | files: Load model files from the memory. This argument is a dictionary mapping |
| 316 | file names to file contents as file-like or bytes objects. If this is set, |
| 317 | :obj:`model_path` acts as an identifier for this model. |
| 318 | )pbdoc") |
| 319 | |
| 320 | .def_property_readonly("device", &TranslatorWrapper::device, |
| 321 | "Device this translator is running on.") |
| 322 | .def_property_readonly("device_index", &TranslatorWrapper::device_index, |
| 323 | "List of device IDs where this translator is running on.") |
| 324 | .def_property_readonly("compute_type", &TranslatorWrapper::compute_type, |
| 325 | "Computation type used by the model.") |
| 326 | .def_property_readonly("num_translators", &TranslatorWrapper::num_replicas, |
| 327 | "Number of translators backing this instance.") |
| 328 | .def_property_readonly("num_queued_batches", &TranslatorWrapper::num_queued_batches, |
| 329 | "Number of batches waiting to be processed.") |
| 330 | .def_property_readonly("tensor_parallel", &TranslatorWrapper::tensor_parallel, |
| 331 | "Run model with tensor parallel mode.") |