Executor in Action

Fastai

This Executor uses the ResNet18 network for object classification on images provided by fastai.

The encode function of this executor generates a feature vector for each image in each Document of the input DocumentArray. The feature vector generated is the output activations of the neural network (a vector of 1000 components).

Note

The embedding of each text is performed in a joined operation (all embeddings are created for all images in a single function call) to achieve higher performance.

As a result each Document in the input DocumentArray docs will have an embedding after encode() has completed.

  1. import torch
  2. from fastai.vision.models import resnet18
  3. from jina import Executor, requests
  4. class ResnetImageEncoder(Executor):
  5. def __init__(self, *args, **kwargs):
  6. super().__init__(*args, **kwargs)
  7. self.model = resnet18()
  8. self.model.eval()
  9. @requests
  10. def encode(self, docs, **kwargs):
  11. batch = torch.Tensor(docs.get_attributes('blob'))
  12. with torch.no_grad():
  13. batch_embeddings = self.model(batch).detach().numpy()
  14. for doc, emb in zip(docs, batch_embeddings):
  15. doc.embedding = emb

Pytorch Lightning

This code snippet uses an autoencoder pretrained from cifar10-resnet18 to build an executor that encodes Document blob( an ndarray that could for example be an image) into embedding . It demonstrates the use of prebuilt model from PyTorch Lightning Bolts to build a Jina encoder.”

  1. from pl_bolts.models.autoencoders import AE
  2. from jina import Executor, requests
  3. import torch
  4. class PLMwuAutoEncoder(Executor):
  5. def __init__(self, **kwargs):
  6. super().__init__()
  7. self.ae = AE(input_height=32).from_pretrained('cifar10-resnet18')
  8. self.ae.freeze()
  9. @requests
  10. def encode(self, docs, **kwargs):
  11. with torch.no_grad():
  12. for doc in docs:
  13. input_tensor = torch.from_numpy(doc.blob)
  14. output_tensor = self.ae(input_tensor)
  15. doc.embedding = output_tensor.detach().numpy()

Paddle

The example below uses the PaddlePaddle Ernie model as the encoder. The Executor loads the pre-trained Ernie tokenizer and model, converts Jina Documents’ doc.text into Paddle Tensors and encodes the text as embeddings. As a result, each Document in the DocumentArray will have an embedding after encode() has completed.

  1. import paddle as P # paddle==2.1.0
  2. import numpy as np
  3. from ernie.modeling_ernie import ErnieModel # paddle-ernie 0.2.0.dev1
  4. from ernie.tokenizing_ernie import ErnieTokenizer
  5. from jina import Executor, requests
  6. class PaddleErineExecutor(Executor):
  7. def __init__(self, **kwargs):
  8. super().__init__()
  9. self.tokenizer = ErnieTokenizer.from_pretrained('ernie-1.0')
  10. self.model = ErnieModel.from_pretrained('ernie-1.0')
  11. self.model.eval()
  12. @requests
  13. def encode(self, docs, **kwargs):
  14. for doc in docs:
  15. ids, _ = self.tokenizer.encode(doc.text)
  16. ids = P.to_tensor(np.expand_dims(ids, 0))
  17. pooled, encoded = self.model(ids)
  18. doc.embedding = pooled.numpy()

Tensorflow

This Executor uses the MobileNetV2 network for object classification on images.

It extracts the buffer field (which is the actual byte array) from each input Document in the DocumentArray docs , preprocesses the byte array and uses MobileNet to predict the classes (dog/car/…) found in the image. Those predictions are Numpy arrays encoding the probability for each class supported by the model (1000 in this case). The Executor stores those arrays then in the embedding for each Document.

As a result each Document in the input DocumentArray docs will have an embedding after encode() has completed.

  1. import numpy as np
  2. import tensorflow as tf
  3. from keras.applications.mobilenet_v2 import MobileNetV2, preprocess_input
  4. from tensorflow.python.framework.errors_impl import InvalidArgumentError
  5. from jina import Executor, requests
  6. class TfMobileNetEncoder(Executor):
  7. def __init__(self, **kwargs):
  8. super().__init__()
  9. self.image_dim = 224
  10. self.model = MobileNetV2(pooling='avg', input_shape=(self.image_dim, self.image_dim, 3))
  11. @requests
  12. def encode(self, docs, **kwargs):
  13. buffers, docs = docs.get_attributes_with_docs('buffer')
  14. tensors = [tf.io.decode_image(contents=b, channels=3) for b in buffers]
  15. resized_tensors = preprocess_input(np.array(self._resize_images(tensors)))
  16. embeds = self.model.predict(np.stack(resized_tensors))
  17. for d, b in zip(docs, embeds):
  18. d.embedding = b
  19. def _resize_images(self, tensors):
  20. resized_tensors = []
  21. for t in tensors:
  22. try:
  23. resized_tensors.append(tf.keras.preprocessing.image.smart_resize(t, (self.image_dim, self.image_dim)))
  24. except InvalidArgumentError:
  25. # this can happen if you include empty or other malformed images
  26. pass
  27. return resized_tensors

MindSpore

The code snippet below takes docs as input and perform matrix multiplication with self.encoding_matrix. It leverages Mindspore Tensor conversion and operation. Finally, the embedding of each document will be set as the multiplied result as np.ndarray.

  1. import numpy as np
  2. from mindspore import Tensor # mindspore 1.2.0
  3. import mindspore.ops as ops
  4. import mindspore.context as context
  5. from jina import Executor, requests
  6. class MindsporeMwuExecutor(Executor):
  7. def __init__(self, **kwargs):
  8. super().__init__()
  9. context.set_context(mode=context.PYNATIVE_MODE, device_target='CPU')
  10. self.encoding_mat = Tensor(np.random.rand(5, 5))
  11. @requests
  12. def encode(self, docs, **kwargs):
  13. matmul = ops.MatMul()
  14. for doc in docs:
  15. input_tensor = Tensor(doc.blob) # convert the ``ndarray`` of the doc to ``Tensor``
  16. output_tensor = matmul(self.encoding_mat, input_tensor) # multiply the input with the encoding matrix.
  17. doc.embedding = output_tensor.asnumpy() # assign the encoding results to ``embedding``

Scikit-learn

This Executor uses a TF-IDF feature vector to generate sparse embeddings for text search.

The class TFIDFTextEncoder extracts stores a tfidf_vectorizer object that it is fitted with a dataset already present in sklearn. The executor provides an encode method that receives a DocumentArray and updates each document in the DocumentArray with an embedding attribute that is the tf-idf representation of the text found in the document.

Note

The embedding of each text is perfomed in a joined operation (all embeddings are creted for all texts in a single function call) to achieve higher performance.

As a result, each Document in the DocumentArray will have an embedding after encode() has completed.

  1. from sklearn.datasets import fetch_20newsgroups
  2. from sklearn.feature_extraction.text import TfidfVectorizer
  3. from jina import Executor, requests, DocumentArray
  4. class TFIDFTextEncoder(Executor):
  5. def __init__(self, *args, **kwargs):
  6. super().__init__(*args, **kwargs)
  7. from sklearn import datasets
  8. dataset = fetch_20newsgroups()
  9. tfidf_vectorizer = TfidfVectorizer()
  10. tfidf_vectorizer.fit(dataset.data)
  11. self.ttfidf_vectorizer = tfidf_vectorizer
  12. @requests
  13. def encode(self, docs: DocumentArray, *args, **kwargs):
  14. iterable_of_texts = docs.get_attributes('text')
  15. embedding_matrix = self.tfidf_vectorizer.transform(iterable_of_texts)
  16. for i, doc in enumerate(docs):
  17. doc.embedding = embedding_matrix[i]

PyTorch

The code snippet below takes docs as input and perform feature extraction with modelnet v2. It leverages Pytorch’s Tensor conversion and operation. Finally, the embedding of each document will be set as the extracted features.

  1. import torch # 1.8.1
  2. import torchvision.models as models # 0.9.1
  3. from jina import Executor, requests
  4. class PytorchMobilNetExecutor(Executor):
  5. def __init__(self, **kwargs):
  6. super().__init__()
  7. self.model = models.quantization.mobilenet_v2(pretrained=True, quantize=True)
  8. self.model.eval()
  9. @requests
  10. def encode(self, docs, **kwargs):
  11. blobs = torch.Tensor(docs.get_attributes('blob'))
  12. with torch.no_grad():
  13. embeds = self.model(blobs).detach().numpy()
  14. for doc, embed in zip(docs, embeds):
  15. doc.embedding = embed

ONNX

The code snippet bellow converts a Pytorch model to the ONNX and leverage onnxruntime to run inference tasks on models from hugging-face transformers.

  1. from pathlib import Path
  2. import numpy as np
  3. import onnxruntime
  4. from jina import Executor, requests
  5. from transformers import BertTokenizerFast, convert_graph_to_onnx
  6. class ONNXBertExecutor(Executor):
  7. def __init__(self, **kwargs):
  8. super().__init__()
  9. # export your huggingface model to onnx
  10. convert_graph_to_onnx.convert(
  11. framework="pt",
  12. model="bert-base-cased",
  13. output=Path("onnx/bert-base-cased.onnx"),
  14. opset=11,
  15. )
  16. # create the tokenizer
  17. self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
  18. # create the inference session
  19. options = onnxruntime.SessionOptions()
  20. options.intra_op_num_threads = 1 # have an impact on performances
  21. options.graph_optimization_level = (
  22. onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
  23. )
  24. # Load the model as a graph and prepare the CPU backend
  25. self.session = onnxruntime.InferenceSession(
  26. "onnx/bert-base-cased.onnx", options
  27. )
  28. self.session.disable_fallback()
  29. @requests
  30. def encode(self, docs, **kwargs):
  31. for doc in docs:
  32. tokens = self.tokenizer.encode_plus(doc.text)
  33. inputs = {name: np.atleast_2d(value) for name, value in tokens.items()}
  34. output, pooled = self.session.run(None, inputs)
  35. # assign the encoding results to ``embedding``
  36. doc.embedding = pooled[0]