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)[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.

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