Integrations

Integration classes for Sentence Transformers and other frameworks.

Sentence Transformers

RKSentenceTransformer

class rktransformers.integrations.sentence_transformers.RKSentenceTransformer[source]

Bases: SentenceTransformer

RKNN-compatible SentenceTransformer implementation.

This class extends SentenceTransformers’ SentenceTransformer to support the RKNN backend. It overrides the encode method to handle RKNN-specific batch size requirements.

__init__(model_name_or_path=None, modules=None, device=None, prompts=None, default_prompt_name=None, similarity_fn_name=None, cache_folder=None, trust_remote_code=False, revision=None, local_files_only=False, token=None, use_auth_token=None, truncate_dim=None, model_kwargs=None, tokenizer_kwargs=None, config_kwargs=None, model_card_data=None, backend='rknn')[source]

Initialize the RKSentenceTransformer.

Parameters:

  • model_name_or_path (str, optional) – Model name or path to load.

  • backend (str) – The backend to use for model loading (default: “rknn”).

  • SentenceTransformer.__init__. (All other arguments are passed to)

  • device (str | None)

  • prompts (dict[str, str] | None)

  • default_prompt_name (str | None)

  • cache_folder (str | None)

  • trust_remote_code (bool)

  • revision (str | None)

  • local_files_only (bool)

  • truncate_dim (int | None)

  • model_kwargs (dict[str, Any] | None)

  • tokenizer_kwargs (dict[str, Any] | None)

  • config_kwargs (dict[str, Any] | None)

Return type:

None

encode(sentences, **kwargs)[source]

Encode sentences with RKNN backend batch size requirements.

Overrides SentenceTransformer.encode to enforce the batch size specified in the RKNN model configuration when encoding sentences.

Parameters:

  • sentences (str or list of str) – The sentence or list of sentences to encode.

  • **kwargs – Additional keyword arguments passed to the SentenceTransformer.encode method. The ‘batch_size’ argument will be overridden if the RKNN backend is active.

Returns:

The encoded sentence embeddings as a numpy array.

Return type:

np.ndarray

RKCrossEncoder

class rktransformers.integrations.sentence_transformers.RKCrossEncoder[source]

Bases: CrossEncoder

RKNN-compatible CrossEncoder implementation.

This class extends SentenceTransformers’ CrossEncoder to support the RKNN backend. It overrides model loading and prediction methods to handle RKNN-specific requirements such as fixed input shapes and quantization.

__init__(model_name_or_path, num_labels=None, max_length=None, activation_fn=None, device=None, cache_folder=None, trust_remote_code=False, revision=None, local_files_only=False, token=None, model_kwargs=None, tokenizer_kwargs=None, config_kwargs=None, model_card_data=None, backend='rknn')[source]

Initialize the RKCrossEncoder.

Parameters:

  • model_name_or_path (str) – The model name on Hugging Face or the path to a local model directory.

  • num_labels (int, optional) – Number of labels for the classifier.

  • max_length (int, optional) – Maximum sequence length.

  • activation_fn (callable, optional) – Activation function.

  • device (str, optional) – Device to use for computation.

  • cache_folder (str, optional) – Directory to cache the model files.

  • trust_remote_code (bool) – Whether to trust remote code.

  • revision (str, optional) – Model revision to use.

  • local_files_only (bool) – Whether to only use local files.

  • token (bool or str, optional) – Hugging Face authentication token.

  • model_kwargs (dict, optional) – Additional model loading arguments.

  • tokenizer_kwargs (dict, optional) – Additional tokenizer loading arguments.

  • config_kwargs (dict, optional) – Additional config loading arguments.

  • model_card_data (optional) – Model card data.

  • backend (str) – The backend to use for model loading (default: “rknn”).

Return type:

None

predict(sentences, batch_size=32, show_progress_bar=None, activation_fn=None, apply_softmax=False, convert_to_numpy=True, convert_to_tensor=False)[source]

Performs predictions with the CrossEncoder on the given sentence pairs.

Overrides CrossEncoder.predict to enforce max_length padding for RKNN models.

Parameters:

  • sentences (Union[List[Tuple[str, str]], Tuple[str, str]]) – A list of sentence pairs [(Sent1, Sent2), (Sent3, Sent4)] or one sentence pair (Sent1, Sent2).

  • batch_size (int, optional) – Batch size for encoding. Overridden by RKNN model’s configured batch size if set.

  • show_progress_bar (bool, optional) – Output progress bar. Defaults to None.

  • activation_fn (callable, optional) – Activation function applied on the logits output of the CrossEncoder. If None, the model.activation_fn will be used, which defaults to torch.nn.Sigmoid if num_labels=1, else torch.nn.Identity. Defaults to None.

  • convert_to_numpy (bool, optional) – Convert the output to a numpy matrix. Defaults to True.

  • apply_softmax (bool, optional) – If set to True and model.num_labels > 1, applies softmax on the logits output such that for each sample, the scores of each class sum to 1. Defaults to False.

  • convert_to_numpy – Whether the output should be a list of numpy vectors. If False, output a list of PyTorch tensors. Defaults to True.

  • convert_to_tensor (bool, optional) – Whether the output should be one large tensor. Overwrites convert_to_numpy. Defaults to False.

Return type:

list[Tensor] | ndarray | Tensor