| 123 | |
| 124 | |
| 125 | void register_generator(py::module& m) { |
| 126 | py::class_<GeneratorWrapper>( |
| 127 | m, "Generator", |
| 128 | R"pbdoc( |
| 129 | A text generator. |
| 130 | |
| 131 | Example: |
| 132 | |
| 133 | >>> generator = ctranslate2.Generator("model/", device="cpu") |
| 134 | >>> generator.generate_batch([["<s>"]], max_length=50, sampling_topk=20) |
| 135 | )pbdoc") |
| 136 | |
| 137 | .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>(), |
| 138 | py::arg("model_path"), |
| 139 | py::arg("device")="cpu", |
| 140 | py::kw_only(), |
| 141 | py::arg("device_index")=0, |
| 142 | py::arg("compute_type")="default", |
| 143 | py::arg("inter_threads")=1, |
| 144 | py::arg("intra_threads")=0, |
| 145 | py::arg("max_queued_batches")=0, |
| 146 | py::arg("flash_attention")=false, |
| 147 | py::arg("tensor_parallel")=false, |
| 148 | py::arg("files")=py::none(), |
| 149 | R"pbdoc( |
| 150 | Initializes the generator. |
| 151 | |
| 152 | Arguments: |
| 153 | model_path: Path to the CTranslate2 model directory. |
| 154 | device: Device to use (possible values are: cpu, cuda, auto). |
| 155 | device_index: Device IDs where to place this generator on. |
| 156 | compute_type: Model computation type or a dictionary mapping a device name |
| 157 | to the computation type (possible values are: default, auto, int8, int8_float32, |
| 158 | int8_float16, int8_bfloat16, int16, float16, bfloat16, float32). |
| 159 | inter_threads: Maximum number of parallel generations. |
| 160 | intra_threads: Number of OpenMP threads per generator (0 to use a default value). |
| 161 | max_queued_batches: Maximum numbers of batches in the queue (-1 for unlimited, |
| 162 | 0 for an automatic value). When the queue is full, future requests will block |
| 163 | until a free slot is available. |
| 164 | flash_attention: run model with flash attention 2 for self-attention layer |
| 165 | tensor_parallel: run model with tensor parallel mode. |
| 166 | files: Load model files from the memory. This argument is a dictionary mapping |
| 167 | file names to file contents as file-like or bytes objects. If this is set, |
| 168 | :obj:`model_path` acts as an identifier for this model. |
| 169 | )pbdoc") |
| 170 | |
| 171 | .def_property_readonly("device", &GeneratorWrapper::device, |
| 172 | "Device this generator is running on.") |
| 173 | .def_property_readonly("device_index", &GeneratorWrapper::device_index, |
| 174 | "List of device IDs where this generator is running on.") |
| 175 | .def_property_readonly("compute_type", &GeneratorWrapper::compute_type, |
| 176 | "Computation type used by the model.") |
| 177 | .def_property_readonly("num_generators", &GeneratorWrapper::num_replicas, |
| 178 | "Number of generators backing this instance.") |
| 179 | .def_property_readonly("num_queued_batches", &GeneratorWrapper::num_queued_batches, |
| 180 | "Number of batches waiting to be processed.") |
| 181 | .def_property_readonly("tensor_parallel", &GeneratorWrapper::tensor_parallel, |
| 182 | "Run model with tensor parallel mode.") |