Transmuter.Serving (transmuter v2.0.0)

Copy Markdown View Source

Builds the Nx.Serving of a task from a model, a tokenizer or a featurizer, and the functions of the task.

text/7 and image/6 build the servings that run the model once. The client side turns the inputs into tensors and gives each batch a key: the sequence length for text, and the image size for images. The serving compiles the model once for each key in its init function. Init runs in the serving process, so the cache of compiled graph segments lives as long as the serving. Each batch is padded to the batch size, so the model only runs the compiled shapes.

each/2 builds the servings that run each input as a batch of its own, such as text-to-speech, whose outputs have a shape that depends on the input.

generation/4 builds the servings that generate tokens with Transmuter.Generation. The rows of a batch generate together in one loop. With stream: true, the loop sends each new token to the client through the :token hook of Nx.Serving.streaming/2, and the client decodes the text.

Summary

Functions

The options that all audio tasks with a Transmuter.Audio.Featurizer take. See Transmuter.Audio.

Raises unless the :batch_size of the :compile option is 1, for a task of each/2. reason says why each input is a batch of its own.

Returns the {height, width} of the images that a serving of featurizer compiles when it starts: the size of the featurizer, or else the size that it gives for a 640x480 image.

Returns an Nx.Serving that runs each input as a batch of one row, for outputs whose shapes depend on the input, such as the audio of text-to-speech. The serving streams the outputs of each batch to the client, so the outputs of the inputs of a list can have different shapes.

Returns an Nx.Serving that generates tokens with a context of Transmuter.Generation. The rows of a batch generate together. tokenizer decodes the text of a stream.

Validates the options of a generation task that Keyword.validate!/2 does not check, before the task loads its models, and returns them. :token_logprobs is true or false. A stream gives the log-probabilities in its done event, so token_logprobs: true with stream: true needs stream_done: true.

Returns the image inputs of model for pixels of shape {batch, channels, height, width}: "pixel_values" and, when the model has a "pixel_mask" input, a mask of ones. The mask has the shape that the model declares, or else the size of the images.

The options that all image tasks take. See Transmuter.Vision.

Returns the {height, width} of the images of featurizer. Raises when the size depends on the image, because images of different sizes cannot go in one batch.

Resizes the last two axes of tensor to {height, width} as torch.nn.functional.interpolate with align_corners=False and without antialias, which the post-processing of Python transformers uses. mode is :bilinear or :bicubic, whose cubic coefficient is -0.75 as in PyTorch.

Divides each vector on the last axis by its L2 norm.

Returns the label of class index from the id2label map of config.json, or "LABEL_<index>".

Loads the repository of source, its model and its tokenizer. See load_model!/4.

Loads the repository of source, its encoder (encoder_model) and its merged decoder (decoder_model_merged). The default variant is fp32. The encoder gives only last_hidden_state, and the decoder gives only the outputs that decoder_outputs selects, by default those of Transmuter.Generation.outputs/1.

Loads the repository of source and its model. The model is the first of components that the repository has, in the default variant variant of the task. The :outputs in opts go to Transmuter.load!/2, so the model gives only the outputs that the task reads.

Returns an Nx.Serving that runs model on the tensors of a task. text/7, image/6 and audio/6 build their servings with this function.

The options that all text tasks take. See Transmuter.Text.

Returns %{predictions: [%{label: label, score: score}]} from the scores of one input, from the highest score. label takes the index of a class and returns its label. top_k is the number of predictions, or nil for all.

Applies a scores function (:softmax, :sigmoid or :none) to logits.

Returns the default scores function of a classification model, as the classification pipelines of Python transformers: :none for regression, :sigmoid for one label or multi-label classification, and else :softmax.

The default sequence lengths of a text task: powers of 4 from 32, up to the smallest of max_length and the limits in tokenizer_config.json and config.json.

The softmax over the last axis.

Accepts a string as one sequence.

Returns the :postprocess of a task whose run generates text: %{results: [%{text: text, token_summary: summary}]}, with :token_logprobs when logprobs? is true, the :token_logprobs? of the context. The token summary of the run gets the number of new tokens.

Puts logprobs, the log-probability of each token of result, in result under :token_logprobs when logprobs? is true, the :token_logprobs? of the context.

Returns the tokenizer of a task: the :tokenizer of opts, or else the tokenizer of repo. The tokenizer pads on the right. The padding then comes after the text, so the text keeps the positions that it has without padding, and a decoder does not see the padding through its causal mask (see Transmuter.Generation.prompt_inputs/3).

Returns the token ids and their log-probabilities of the integers of a row whose values are the tokens of Transmuter.Generation.generate/5. Without log-probabilities, the integers are the ids, and the list of log-probabilities is empty.

Functions

audio(model, featurizer, opts, fun, postprocess, inputs \\ %{})

Returns an Nx.Serving that runs model on audio.

The featurizer (Transmuter.Audio.Featurizer) turns the audio into the input of Transmuter.Audio.Featurizer.input_name/1, such as "input_values", on the :backend of opts. When the features of all audio have one shape, such as the 1024 frames of AST, the serving compiles that shape, and the audio goes into batches of the :batch_size of :compile. When the shape depends on the length of the audio, such as the samples of wav2vec2, the batch size must be 1. Each audio is then a batch, and the compiled model compiles each new length on its first call. It keeps up to 64 lengths.

The input is an audio of Transmuter.Audio.Featurizer.audio/2. The serving joins the chunks of a stream into one tensor.

fun and postprocess are as in text/7. An audio is one row, so postprocess takes a list of one row and a list of one audio tensor. inputs are other inputs of the model that are the same for all audio, as in image/6.

audio_options()

The options that all audio tasks with a Transmuter.Audio.Featurizer take. See Transmuter.Audio.

batch_size_one!(compile, reason)

Raises unless the :batch_size of the :compile option is 1, for a task of each/2. reason says why each input is a batch of its own.

compile_size(featurizer)

Returns the {height, width} of the images that a serving of featurizer compiles when it starts: the size of the featurizer, or else the size that it gives for a 640x480 image.

each(task, opts)

Returns an Nx.Serving that runs each input as a batch of one row, for outputs whose shapes depend on the input, such as the audio of text-to-speech. The serving streams the outputs of each batch to the client, so the outputs of the inputs of a list can have different shapes.

task is a map:

  • :validate takes one input and returns {:ok, input} or {:error, message}.

  • :preprocess takes a valid input and returns a map of Nx.BinaryBackend tensors without a batch axis.

  • :init takes the defn options and returns a function. The serving calls :init in its process, so it can compile models there. The function takes the tensors of one input, with a batch axis of size 1, and returns a map of tensors with a batch axis of size 1.

  • :postprocess takes the map of tensors of one input, without the batch axis, and returns the result.

The serving reads :defn_options of opts.

generation(ctx, tokenizer, task, opts)

Returns an Nx.Serving that generates tokens with a context of Transmuter.Generation. The rows of a batch generate together. tokenizer decodes the text of a stream.

task is a map:

  • :validate takes one input and returns {:ok, input} or {:error, message}.

  • :preprocess takes a valid input and returns an enumerable of runs, such as the windows of long audio. The client makes the rows and the batches of a list of runs, so an error in them raises in the caller. A stream of runs, such as the windows of an audio stream, goes into batches when Nx.Serving reads them, once, in a process of its own. An error in such a stream ends the call with an exit, not an exception that the caller can rescue. A run is a map with the token ids of the prompt of the decoder in :prompt, the seed of sampling in :seed, a map of Nx.BinaryBackend tensors without a batch axis for :init in :tensors, a batch key in :key, and %{input: count, padding: count} in :summary for the token summary of a stream. Runs with the same key have tensors of the same shapes. Each run is one row of a batch.

  • :init takes the defn options and returns a function. The serving calls :init in its process, so it can compile models there. The function takes the tensors of a batch, the prompts of its rows and a function generate. It calls generate with the other inputs of the decoder, such as the outputs of an encoder, and the prompts, one or more times. An empty prompt is a padding row. generate returns the tokens of each row, see Transmuter.Generation.generate/5. The function returns a list for each row, such as the result of generate. Its values are integers, f32 values such as log-probabilities, and lists of them. The serving sends the row as a flat list of integers, with each f32 value as its bits (see Transmuter.Generation.from_bits/1).

  • :postprocess takes the list of the integers of each run of an input, and the runs from :preprocess for this input, and returns the result.

  • :keys lists the keys of the runs. The default is [nil].

opts are the options of the task, after generation_options!/1. The serving reads :batch_size of :compile, :stream, :stream_done and :defn_options. The done event of a stream has the log-probabilities when the context has token_logprobs: true.

generation_options!(opts)

Validates the options of a generation task that Keyword.validate!/2 does not check, before the task loads its models, and returns them. :token_logprobs is true or false. A stream gives the log-probabilities in its done event, so token_logprobs: true with stream: true needs stream_done: true.

image(model, featurizer, opts, fun, postprocess, inputs \\ %{})

Returns an Nx.Serving that runs model on images.

The featurizer turns the images into "pixel_values" on the :backend of opts. When the model has a "pixel_mask" input, each image also gets a mask of ones. The mask has the shape that the model declares, such as 64x64 in the DETR exports of Xenova, or else the size of the image. The images of a batch have one size, so no pixel is padding. inputs are other inputs of the model that are the same for all images, such as the text of the labels in zero-shot classification. Their first axis is not the batch.

When the featurizer gives one image size, the serving compiles that shape, and the images go into batches of the :batch_size of :compile. When the ONNX file declares a fixed size of the batch axis of pixel_values, the batch size must be that size. When the size depends on the image, such as with the shortest_edge and longest_edge of DETR, the batch size must be 1. Each image is then a batch, and the compiled model compiles each new image size on its first call. It keeps up to 64 sizes.

fun and postprocess are as in text/7. The tensors of fun must have the same shape for all image sizes. An image is one row, so postprocess takes a list of one row and a list of one image tensor.

image_inputs(model, pixels)

Returns the image inputs of model for pixels of shape {batch, channels, height, width}: "pixel_values" and, when the model has a "pixel_mask" input, a mask of ones. The mask has the shape that the model declares, or else the size of the images.

image_options()

The options that all image tasks take. See Transmuter.Vision.

image_size!(featurizer)

Returns the {height, width} of the images of featurizer. Raises when the size depends on the image, because images of different sizes cannot go in one batch.

interpolate(tensor, arg, mode)

Resizes the last two axes of tensor to {height, width} as torch.nn.functional.interpolate with align_corners=False and without antialias, which the post-processing of Python transformers uses. mode is :bilinear or :bicubic, whose cubic coefficient is -0.75 as in PyTorch.

The resize is the ONNX Resize operator with the pytorch_half_pixel coordinates. Transmuter.Featurizer.resize/3 gives the pixels of Pillow, whose filters widen when the image gets smaller. This filter does not widen.

l2_norm(tensor)

Divides each vector on the last axis by its L2 norm.

label(repo, index)

Returns the label of class index from the id2label map of config.json, or "LABEL_<index>".

load!(source, opts, variant, components \\ ["model"])

Loads the repository of source, its model and its tokenizer. See load_model!/4.

load_encoder_decoder!(source, opts, decoder_outputs \\ &Generation.outputs/1)

Loads the repository of source, its encoder (encoder_model) and its merged decoder (decoder_model_merged). The default variant is fp32. The encoder gives only last_hidden_state, and the decoder gives only the outputs that decoder_outputs selects, by default those of Transmuter.Generation.outputs/1.

load_model!(source, opts, variant, components \\ ["model"])

Loads the repository of source and its model. The model is the first of components that the repository has, in the default variant variant of the task. The :outputs in opts go to Transmuter.load!/2, so the model gives only the outputs that the task reads.

new(model, opts, shapes, preprocess, fun, postprocess, validate)

Returns an Nx.Serving that runs model on the tensors of a task. text/7, image/6 and audio/6 build their servings with this function.

shapes is a map:

  • :batch_size - the number of rows of a batch.

  • :templates - the input tensors of each batch key. The serving compiles the model for them.

  • :inputs - inputs of the model that are the same for all batches. Their first axis is not the batch.

  • :cache_size - optional, the :cache_size option of Transmuter.compile/3.

  • :steps - optional, the inputs and the static inputs of a step of a decoder for each batch key (see Transmuter.Generation.prompt_inputs/3).

  • :prepare - optional, a function that takes the inputs of a batch, padded to the batch size, and the number of rows that are not padding. It returns the inputs of the batch for the model, such as inputs whose shapes depend on the values of the batch. The fixed inputs and the inputs of the steps join them after this function. The serving compiles the model for the templates through this function too, as a full batch. The rows of the batch stay first. fun gets the outputs of the prepared batch, which can have one more row than the inputs. The serving gives the outputs of the rows that are not padding to the client.

  • :streaming - optional, true when the outputs of a batch can have other shapes than those of another batch, such as with :prepare. The serving then streams the outputs of each batch to the client, so that Nx.Serving does not join them.

validate gives the sequences of each input. preprocess takes the sequences of all inputs and gives a list of {tensors, batch key}, one for each batch, and for each sequence the list of its rows. The rows of the batches are those of the sequences, in order. fun and postprocess are as in text/7.

options()

The options that all text tasks take. See Transmuter.Text.

predictions(scores, label, top_k)

Returns %{predictions: [%{label: label, score: score}]} from the scores of one input, from the highest score. label takes the index of a class and returns its label. top_k is the number of predictions, or nil for all.

scores(logits, atom)

Applies a scores function (:softmax, :sigmoid or :none) to logits.

scores_function(repo)

Returns the default scores function of a classification model, as the classification pipelines of Python transformers: :none for regression, :sigmoid for one label or multi-label classification, and else :softmax.

sequence_lengths(repo, max_length)

The default sequence lengths of a text task: powers of 4 from 32, up to the smallest of max_length and the limits in tokenizer_config.json and config.json.

softmax(logits)

The softmax over the last axis.

string(input)

Accepts a string as one sequence.

text(repo, model, tokenizer, opts, fun, postprocess, validate \\ &string/1)

Returns an Nx.Serving that runs model on text.

validate takes one input and returns {:ok, sequences} or {:error, message}. The sequences are the strings or {text, text_pair} tuples to encode for the input. Each sequence gives one row of the batch, or one row for each window when the tokenizer has a :stride (see Transmuter.Tokenizer.load!/2).

fun takes the outputs of the model and the inputs of a batch, and returns a map of tensors with the rows on the first axis. postprocess takes the rows of one input and returns the result for that input. Its arguments are the list of rows of the map from fun, and a list of {sequence, encoding} for these rows.

The :max_length option of a task limits the default sequence lengths, as do the limits in tokenizer_config.json and config.json. The :encode_options of a task go to Transmuter.Tokenizer.encode/3.

text_postprocess(tokenizer, logprobs?)

Returns the :postprocess of a task whose run generates text: %{results: [%{text: text, token_summary: summary}]}, with :token_logprobs when logprobs? is true, the :token_logprobs? of the context. The token summary of the run gets the number of new tokens.

token_logprobs(result, logprobs, logprobs?)

Puts logprobs, the log-probability of each token of result, in result under :token_logprobs when logprobs? is true, the :token_logprobs? of the context.

tokenizer(repo, opts)

Returns the tokenizer of a task: the :tokenizer of opts, or else the tokenizer of repo. The tokenizer pads on the right. The padding then comes after the text, so the text keeps the positions that it has without padding, and a decoder does not see the padding through its causal mask (see Transmuter.Generation.prompt_inputs/3).

tokens(integers, bool)

Returns the token ids and their log-probabilities of the integers of a row whose values are the tokens of Transmuter.Generation.generate/5. Without log-probabilities, the integers are the ids, and the list of log-probabilities is empty.