#include "module.h"

#include <ctranslate2/encoder.h>

#include "replica_pool.h"

namespace ctranslate2 {

namespace python {

class EncoderWrapper : public ReplicaPoolHelper<Encoder> {

public:

using ReplicaPoolHelper::ReplicaPoolHelper;

EncoderForwardOutput

forward_batch(const std::variant<BatchTokens, BatchIds, StorageView>& inputs,

const std::optional<StorageView>& lengths,

const std::optional<BatchIds>& token_type_ids) {

std::future<EncoderForwardOutput> future;

std::shared_lock lock(_mutex);

assert_model_is_ready();

switch (inputs.index()) {

case 0:

future = _pool->forward_batch_async(

std::get<BatchTokens>(inputs),

token_type_ids.value_or(std::vector<std::vector<size_t>>()));

break;

case 1:

future = _pool->forward_batch_async(

std::get<BatchIds>(inputs),

token_type_ids.value_or(std::vector<std::vector<size_t>>()));

break;

case 2:

if (!lengths)

throw std::invalid_argument("lengths vector is required when passing a dense input");

future = _pool->forward_batch_async(

std::get<StorageView>(inputs),

lengths.value(),

token_type_ids.value_or(std::vector<std::vector<size_t>>()));

break;

}

return future.get();

}

};

void register_encoder(py::module& m) {

py::class_<EncoderForwardOutput>(m, "EncoderForwardOutput",

"Forward output of an encoder model.")

.def_readonly("last_hidden_state", &EncoderForwardOutput::last_hidden_state,

"Output of the last layer.")

.def_readonly("pooler_output", &EncoderForwardOutput::pooler_output,

"Output of the pooling layer.")

.def("__repr__", [](const EncoderForwardOutput& output) {

return "EncoderForwardOutput(last_hidden_state="

+ std::string(py::repr(py::cast(output.last_hidden_state)))

+ ", pooler_output=" + std::string(py::repr(py::cast(output.pooler_output)))

+ ")";

})

;

py::class_<EncoderWrapper>(

m, "Encoder",

R"pbdoc(

A text encoder.

Example:

>>> encoder = ctranslate2.Encoder("model/", device="cpu")

>>> encoder.forward_batch([["▁Hello", "▁world", "!"]])

)pbdoc")

.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>(),

py::arg("model_path"),

py::arg("device")="cpu",

py::kw_only(),

py::arg("device_index")=0,

py::arg("compute_type")="default",

py::arg("inter_threads")=1,

py::arg("intra_threads")=0,

py::arg("max_queued_batches")=0,

py::arg("flash_attention")=false,

py::arg("tensor_parallel")=false,

py::arg("files")=py::none(),

R"pbdoc(

Initializes the encoder.

Arguments:

model_path: Path to the CTranslate2 model directory.

device: Device to use (possible values are: cpu, cuda, auto).

device_index: Device IDs where to place this encoder on.

compute_type: Model computation type or a dictionary mapping a device name

to the computation type (possible values are: default, auto, int8, int8_float32,

int8_float16, int8_bfloat16, int16, float16, bfloat16, float32).

inter_threads: Maximum number of parallel generations.

intra_threads: Number of OpenMP threads per encoder (0 to use a default value).

max_queued_batches: Maximum numbers of batches in the queue (-1 for unlimited,

0 for an automatic value). When the queue is full, future requests will block

until a free slot is available.

flash_attention: run model with flash attention 2 for self-attention layer

tensor_parallel: run model with tensor parallel mode

files: Load model files from the memory. This argument is a dictionary mapping

file names to file contents as file-like or bytes objects. If this is set,

:obj:`model_path` acts as an identifier for this model.

)pbdoc")

.def_property_readonly("device", &EncoderWrapper::device,

"Device this encoder is running on.")

.def_property_readonly("device_index", &EncoderWrapper::device_index,

"List of device IDs where this encoder is running on.")

.def_property_readonly("compute_type", &EncoderWrapper::compute_type,

"Computation type used by the model.")

.def_property_readonly("num_encoders", &EncoderWrapper::num_replicas,

"Number of encoders backing this instance.")

.def_property_readonly("num_queued_batches", &EncoderWrapper::num_queued_batches,

"Number of batches waiting to be processed.")

.def_property_readonly("tensor_parallel", &EncoderWrapper::tensor_parallel,

"Run model with tensor parallel mode.")

.def_property_readonly("num_active_batches", &EncoderWrapper::num_active_batches,

"Number of batches waiting to be processed or currently processed.")

.def("forward_batch", &EncoderWrapper::forward_batch,

py::arg("inputs"),

py::arg("lengths")=py::none(),

py::arg("token_type_ids")=py::none(),

py::call_guard<py::gil_scoped_release>(),

R"pbdoc(

Forwards a batch of sequences in the encoder.

Arguments:

inputs: A batch of sequences either as string tokens or token IDs.

This argument can also be a dense int32 array with shape

``[batch_size, max_length]`` (e.g. created from a Numpy array or PyTorch tensor).

lengths: The length of each sequence as a int32 array with shape

``[batch_size]``. Required when :obj:`inputs` is a dense array.

token_type_ids: A batch of token type IDs of same shape as :obj:`inputs`.

``[batch_size, max_length]``.

Returns:

The encoder model output.

)pbdoc")

.def("unload_model", &EncoderWrapper::unload_model,

py::arg("to_cpu")=false,

py::call_guard<py::gil_scoped_release>(),

R"pbdoc(

Unloads the model attached to this encoder but keep enough runtime context

to quickly resume encoder on the initial device.

Arguments:

to_cpu: If ``True``, the model is moved to the CPU memory and not fully unloaded.

)pbdoc")

.def("load_model", &EncoderWrapper::load_model,

py::arg("keep_cache")=false,

py::call_guard<py::gil_scoped_release>(),

R"pbdoc(

Loads the model back to the initial device.

Arguments:

keep_cache: If ``True``, the model cache in the CPU memory is not deleted if it exists.

)pbdoc")

.def_property_readonly("model_is_loaded", &EncoderWrapper::model_is_loaded,

"Whether the model is loaded on the initial device and ready to be used.")

;

}

}

}