| 20 | |
| 21 | |
| 22 | void register_wav2vec2bert(py::module& m) { |
| 23 | py::class_<Wav2Vec2BertWrapper>( |
| 24 | m, "Wav2Vec2Bert", |
| 25 | R"pbdoc( |
| 26 | Implements the Wav2Vec2Bert speech recognition model published by Facebook. |
| 27 | |
| 28 | See Also: |
| 29 | https://github.com/facebookresearch/fairseq/tree/main/examples/wav2vec |
| 30 | )pbdoc") |
| 31 | |
| 32 | .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>(), |
| 33 | py::arg("model_path"), |
| 34 | py::arg("device")="cpu", |
| 35 | py::kw_only(), |
| 36 | py::arg("device_index")=0, |
| 37 | py::arg("compute_type")="default", |
| 38 | py::arg("inter_threads")=1, |
| 39 | py::arg("intra_threads")=0, |
| 40 | py::arg("max_queued_batches")=0, |
| 41 | py::arg("flash_attention")=false, |
| 42 | py::arg("tensor_parallel")=false, |
| 43 | py::arg("files")=py::none(), |
| 44 | R"pbdoc( |
| 45 | Initializes a Wav2Vec2Bert model from a converted model. |
| 46 | |
| 47 | Arguments: |
| 48 | model_path: Path to the CTranslate2 model directory. |
| 49 | device: Device to use (possible values are: cpu, cuda, auto). |
| 50 | device_index: Device IDs where to place this model on. |
| 51 | compute_type: Model computation type or a dictionary mapping a device name |
| 52 | to the computation type (possible values are: default, auto, int8, int8_float32, |
| 53 | int8_float16, int8_bfloat16, int16, float16, bfloat16, float32). |
| 54 | inter_threads: Number of workers to allow executing multiple batches in parallel. |
| 55 | intra_threads: Number of OpenMP threads per worker (0 to use a default value). |
| 56 | max_queued_batches: Maximum numbers of batches in the worker queue (-1 for unlimited, |
| 57 | 0 for an automatic value). When the queue is full, future requests will block |
| 58 | until a free slot is available. |
| 59 | flash_attention: run model with flash attention 2 for self-attention layer |
| 60 | tensor_parallel: run model with tensor parallel mode |
| 61 | files: Load model files from the memory. This argument is a dictionary mapping |
| 62 | file names to file contents as file-like or bytes objects. If this is set, |
| 63 | :obj:`model_path` acts as an identifier for this model. |
| 64 | )pbdoc") |
| 65 | |
| 66 | .def_property_readonly("device", &Wav2Vec2BertWrapper::device, |
| 67 | "Device this model is running on.") |
| 68 | .def_property_readonly("device_index", &Wav2Vec2BertWrapper::device_index, |
| 69 | "List of device IDs where this model is running on.") |
| 70 | .def_property_readonly("compute_type", &Wav2Vec2BertWrapper::compute_type, |
| 71 | "Computation type used by the model.") |
| 72 | .def_property_readonly("num_workers", &Wav2Vec2BertWrapper::num_replicas, |
| 73 | "Number of model workers backing this instance.") |
| 74 | .def_property_readonly("num_queued_batches", &Wav2Vec2BertWrapper::num_queued_batches, |
| 75 | "Number of batches waiting to be processed.") |
| 76 | .def_property_readonly("tensor_parallel", &Wav2Vec2BertWrapper::tensor_parallel, |
| 77 | "Run model with tensor parallel mode.") |
| 78 | .def_property_readonly("num_active_batches", &Wav2Vec2BertWrapper::num_active_batches, |
| 79 | "Number of batches waiting to be processed or currently processed.") |