docarray.array.mixins.embed module#

class docarray.array.mixins.embed.EmbedMixin[source]#

Bases: object

Helper functions for embedding with a model

embed(embed_model, device='cpu', batch_size=256, to_numpy=False, collate_fn=None)[source]#

Fill embedding of Documents inplace by using embed_model

Parameters
  • embed_model (AnyDNN) – the embedding model written in Keras/Pytorch/Paddle

  • device (str) – the computational device for embed_model, can be either cpu or cuda.

  • batch_size (int) – number of Documents in a batch for embedding

  • to_numpy (bool) – if to store embeddings back to Document in numpy.ndarray or original framework format.

  • collate_fn (Optional[ForwardRef]) – create a mini-batch of Input(s) from the given DocumentArray. Default built-in collate_fn is to use the tensors of the documents.

Return type

T

Returns

itself after modified.

docarray.array.mixins.embed.get_framework(dnn_model)[source]#

Return the framework that powers a DNN model.

Note

This is not a solid implementation. It is based on __module__ name, the key idea is to tell dnn_model without actually importing the framework.

Parameters

dnn_model – a DNN model

Return type

str

Returns

keras, torch, paddle or ValueError