defmodule Bumblebee.Text do @moduledoc """ High-level tasks related to text processing. """ @type token_classification_input :: String.t() @type token_classification_output :: %{entities: list(token_classification_entity())} @typedoc """ A single entity label. Note that `start` and `end` indices are expressed in terms of UTF-8 bytes. """ @type token_classification_entity :: %{ start: non_neg_integer(), end: non_neg_integer(), score: float(), label: String.t(), phrase: String.t() } @doc """ Builds serving for token classification. The serving accepts `t:token_classification_input/0` and returns `t:token_classification_output/0`. A list of inputs is also supported. This function can be used for tasks such as named entity recognition (NER) or part of speech tagging (POS). The recognized entities can optionally be aggregated into groups based on the given strategy. ## Options * `:aggregation` - an optional strategy for aggregating adjacent tokens. Token classification models output probabilities for each possible token class. The aggregation strategy takes scores for each token (which possibly represents subwords) and groups tokens into phrases which are readily interpretable as entities of a certain class. Supported aggregation strategies: * `nil` (default) - corresponds to no aggregation and returns the most likely label for each input token * `:same` - groups adjacent tokens with the same label. If the labels use beginning-inside-outside (BIO) tagging, the boundaries are respected and the prefix is omitted in the output labels * `:word_first` - uses `:same` strategy except that word tokens cannot end up with different labels. With this strategy word gets the label of the first token of that word when there is ambiguity. Note that this works only on word based models * `:word_average` - uses `:same` strategy except that word tokens cannot end up with different labels. With this strategy scores are averaged across word tokens and then the maximum label is taken. Note that this works only on word based models * `:word_max` - uses `:same` strategy except that word tokens cannot end up with different labels. With this strategy word gets the label of the token with the maximum score. Note that this works only on word based models * `:ignored_labels` - the labels to ignore in the final output. The labels should be specified without BIO prefix. Defaults to `["O"]` * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:scores_function` - the function to use for converting logits to scores. Should be one of `:softmax`, `:sigmoid`, or `:none`. Defaults to `:softmax` * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` ## Examples {:ok, bert} = Bumblebee.load_model({:hf, "dslim/bert-base-NER"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "google-bert/bert-base-cased"}) serving = Bumblebee.Text.token_classification(bert, tokenizer, aggregation: :same) text = "Rachel Green works at Ralph Lauren in New York City in the sitcom Friends" Nx.Serving.run(serving, text) #=> %{ #=> entities: [ #=> %{end: 12, label: "PER", phrase: "Rachel Green", score: 0.9997024834156036, start: 0}, #=> %{end: 34, label: "ORG", phrase: "Ralph Lauren", score: 0.9968731701374054, start: 22}, #=> %{end: 51, label: "LOC", phrase: "New York City", score: 0.9995547334353129, start: 38}, #=> %{end: 73, label: "MISC", phrase: "Friends", score: 0.6997143030166626, start: 66} #=> ] #=>} """ @spec token_classification( Bumblebee.model_info(), Bumblebee.Tokenizer.t(), keyword() ) :: Nx.Serving.t() defdelegate token_classification(model_info, tokenizer, opts \\ []), to: Bumblebee.Text.TokenClassification @type generation_input :: String.t() | %{:text => String.t(), optional(:seed) => integer() | nil} @type generation_output :: %{results: list(generation_result())} @type generation_result :: %{text: String.t(), token_summary: token_summary()} @type token_summary :: %{ input: pos_integer(), outout: pos_integer(), padding: non_neg_integer() } @doc """ Builds serving for prompt-driven text generation. The serving accepts `t:generation_input/0` and returns `t:generation_output/0`. A list of inputs is also supported. ## Options * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` * `:stream` - when `true`, the serving immediately returns a stream that emits text chunks as they are generated. Note that when using streaming, only a single input can be given to the serving. To process a batch, call the serving with each input separately. Defaults to `false` * `:stream_done` - when `:stream` is enabled, this enables a final event, after all chunks have been emitted. The event has the shape `{:done, result}`, where `result` includes the same fields as `t:generation_result/0`, except for `:text`, which has been already streamed. Defaults to `false` ## Examples {:ok, model_info} = Bumblebee.load_model({:hf, "openai-community/gpt2"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai-community/gpt2"}) {:ok, generation_config} = Bumblebee.load_generation_config({:hf, "openai-community/gpt2"}) generation_config = Bumblebee.configure(generation_config, max_new_tokens: 15) serving = Bumblebee.Text.generation(model_info, tokenizer, generation_config) Nx.Serving.run(serving, "Elixir is a functional") #=> %{ #=> results: [ #=> %{ #=> text: "Elixir is a functional programming language that is designed to be used in a variety of applications. It" #=> } #=> ] #=> } We can stream the result by creating the serving with `stream: true`: {:ok, model_info} = Bumblebee.load_model({:hf, "openai-community/gpt2"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai-community/gpt2"}) {:ok, generation_config} = Bumblebee.load_generation_config({:hf, "openai-community/gpt2"}) generation_config = Bumblebee.configure(generation_config, max_new_tokens: 15) serving = Bumblebee.Text.generation(model_info, tokenizer, generation_config, stream: true) Nx.Serving.run(serving, "Elixir is a functional") |> Enum.to_list() #=> [" programming", " language", " that", " is", " designed", " to", " be", " used", " in", " a", #=> " variety", " of", " applications.", " It"] """ @spec generation( Bumblebee.model_info(), Bumblebee.Tokenizer.t(), Bumblebee.Text.GenerationConfig.t(), keyword() ) :: Nx.Serving.t() defdelegate generation(model_info, tokenizer, generation_config, opts \\ []), to: Bumblebee.Text.TextGeneration @type text_classification_input :: String.t() @type text_classification_output :: %{predictions: list(text_classification_prediction())} @type text_classification_prediction :: %{score: number(), label: String.t()} @doc """ Builds serving for text classification. The serving accepts `t:text_classification_input/0` and returns `t:text_classification_output/0`. A list of inputs is also supported. ## Options * `:top_k` - the number of top predictions to include in the output. If the configured value is higher than the number of labels, all labels are returned. Defaults to `5` * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:scores_function` - the function to use for converting logits to scores. Should be one of `:softmax`, `:sigmoid`, or `:none`. Defaults to `:softmax` * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` ## Examples {:ok, bertweet} = Bumblebee.load_model({:hf, "finiteautomata/bertweet-base-sentiment-analysis"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "vinai/bertweet-base"}) serving = Bumblebee.Text.text_classification(bertweet, tokenizer) text = "Cats are cute." Nx.Serving.run(serving, text) #=> %{ #=> predictions: [ #=> %{label: "POS", score: 0.9876555800437927}, #=> %{label: "NEU", score: 0.010068908333778381}, #=> %{label: "NEG", score: 0.002275536535307765} #=> ] #=> } """ @spec text_classification( Bumblebee.model_info(), Bumblebee.Tokenizer.t(), keyword() ) :: Nx.Serving.t() defdelegate text_classification(model_info, tokenizer, opts \\ []), to: Bumblebee.Text.TextClassification @type text_embedding_input :: String.t() @type text_embedding_output :: %{embedding: Nx.Tensor.t()} @doc """ Builds serving for text embeddings. The serving accepts `t:text_embedding_input/0` and returns `t:text_embedding_output/0`. A list of inputs is also supported. ## Options * `:output_attribute` - the attribute of the model output map to retrieve. When the output is a single tensor (rather than a map), this option is ignored. Defaults to `:pooled_state` * `:output_pool` - pooling to apply on top of the model output, in case it is not already a pooled embedding. Supported values: `:mean_pooling`. By default no pooling is applied * `:embedding_processor` - a post-processing step to apply to the embedding. Supported values: `:l2_norm`. By default the output is returned as is * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` ## Examples {:ok, model_info} = Bumblebee.load_model({:hf, "intfloat/e5-large"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "intfloat/e5-large"}) serving = Bumblebee.Text.text_embedding(model_info, tokenizer) text = "query: Cats are cute." Nx.Serving.run(serving, text) #=> %{ #=> embedding: #Nx.Tensor< #=> f32[1024] #=> EXLA.Backend #=> [-0.9789889454841614, -0.9814645051956177, -0.5015208125114441, 0.9867952466011047, 0.9917466640472412, -0.5557178258895874, -0.18618212640285492, 0.797040581703186, 0.8922086954116821, 0.7599573135375977, -0.16524426639080048, -0.8740050792694092, 0.9433475732803345, 0.7217797636985779, 0.9437620639801025, 0.4694959223270416, 0.40594056248664856, -0.20143413543701172, 0.7144518494606018, -0.8689796924591064, 0.94001305103302, 0.17163503170013428, -0.9896315932273865, 0.4455447494983673, 0.41139301657676697, 0.01911175064742565, -0.11275406181812286, -0.734498143196106, -0.6410953402519226, -0.628239095211029, -0.2570168673992157, 0.475137323141098, -0.7534396052360535, -0.9492156505584717, -0.17271563410758972, 0.9081271886825562, -0.4851466119289398, -0.9440935254096985, -0.20976334810256958, -0.684502899646759, -0.11581139266490936, 0.17509342730045319, 0.05547652021050453, 0.31042391061782837, 0.955132007598877, -0.35595986247062683, 0.016105204820632935, -0.3154579997062683, 0.9630348682403564, ...] #=> > #=> } """ @spec text_embedding( Bumblebee.model_info(), Bumblebee.Tokenizer.t(), keyword() ) :: Nx.Serving.t() defdelegate text_embedding(model_info, tokenizer, opts \\ []), to: Bumblebee.Text.TextEmbedding @type fill_mask_input :: String.t() @type fill_mask_output :: %{predictions: list(fill_mask_prediction())} @type fill_mask_prediction :: %{score: number(), token: String.t()} @doc """ Builds serving for the fill-mask task. The serving accepts `t:fill_mask_input/0` and returns `t:fill_mask_output/0`. A list of inputs is also supported. In the fill-mask task, the objective is to predict a masked word in the text. The serving expects the input to have exactly one such word, denoted as `[MASK]`. ## Options * `:top_k` - the number of top predictions to include in the output. If the configured value is higher than the number of labels, all labels are returned. Defaults to `5` * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` ## Examples {:ok, bert} = Bumblebee.load_model({:hf, "google-bert/bert-base-uncased"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "google-bert/bert-base-uncased"}) serving = Bumblebee.Text.fill_mask(bert, tokenizer) text = "The capital of [MASK] is Paris." Nx.Serving.run(serving, text) #=> %{ #=> predictions: [ #=> %{score: 0.9279842972755432, token: "france"}, #=> %{score: 0.008412551134824753, token: "brittany"}, #=> %{score: 0.007433671969920397, token: "algeria"}, #=> %{score: 0.004957548808306456, token: "department"}, #=> %{score: 0.004369721747934818, token: "reunion"} #=> ] #=> } """ @spec fill_mask(Bumblebee.model_info(), Bumblebee.Tokenizer.t(), keyword()) :: Nx.Serving.t() defdelegate fill_mask(model_info, tokenizer, opts \\ []), to: Bumblebee.Text.FillMask @type question_answering_input :: %{question: String.t(), context: String.t()} @type question_answering_output :: %{ predictions: list(question_answering_result()) } @type question_answering_result :: %{ text: String.t(), start: number(), end: number(), score: number() } @doc """ Builds serving for the question answering task. The serving accepts `t:question_answering_input/0` and returns `t:question_answering_output/0`. A list of inputs is also supported. The question answering task finds the most probable answer to a question within the given context text. ## Options * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size. Note that the batch size refers to the number of prompts to classify, while the model prediction is made for every combination of prompt and label * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` ## Examples {:ok, roberta} = Bumblebee.load_model({:hf, "deepset/roberta-base-squad2"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "FacebookAI/roberta-base"}) serving = Bumblebee.Text.question_answering(roberta, tokenizer) input = %{question: "What's my name?", context: "My name is Sarah and I live in London."} Nx.Serving.run(serving, input) #=> %{results: [%{end: 16, score: 0.81039959192276, start: 11, text: "Sarah"}]} """ @spec question_answering( Bumblebee.model_info(), Bumblebee.Tokenizer.t(), keyword() ) :: Nx.Serving.t() defdelegate question_answering(model_info, tokenizer, opts \\ []), to: Bumblebee.Text.QuestionAnswering @type zero_shot_classification_input :: String.t() @type zero_shot_classification_output :: %{ predictions: list(zero_shot_classification_prediction()) } @type zero_shot_classification_prediction :: %{score: number(), label: String.t()} @doc """ Builds serving for the zero-shot classification task. The serving accepts `t:zero_shot_classification_input/0` and returns `t:zero_shot_classification_output/0`. A list of inputs is also supported. The zero-shot task predicts zero-shot labels for a given sequence by proposing each label as a premise-hypothesis pairing. ## Options * `:top_k` - the number of top predictions to include in the output. If the configured value is higher than the number of labels, all labels are returned. Defaults to `5` * `:hypothesis_template` - an arity-1 function which accepts a label and returns a hypothesis. The default hypothesis format is: "This example is #\{label\}". * `:compile` - compiles all computations for predefined input shapes during serving initialization. Should be a keyword list with the following keys: * `:batch_size` - the maximum batch size of the input. Inputs are optionally padded to always match this batch size. Note that the batch size refers to the number of prompts to classify, while the model prediction is made for every combination of prompt and label * `:sequence_length` - the maximum input sequence length. Input sequences are always padded/truncated to match that length. A list can be given, in which case the serving compiles a separate computation for each length and then inputs are matched to the smallest bounding length It is advised to set this option in production and also configure a defn compiler using `:defn_options` to maximally reduce inference time. * `:defn_options` - the options for JIT compilation. Defaults to `[]` * `:preallocate_params` - when `true`, explicitly allocates params on the device configured by `:defn_options`. You may want to set this option when using partitioned serving, to allocate params on each of the devices. When using this option, you should first load the parameters into the host. This can be done by passing `backend: {EXLA.Backend, client: :host}` to `load_model/1` and friends. Defaults to `false` ## Examples {:ok, model} = Bumblebee.load_model({:hf, "facebook/bart-large-mnli"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "facebook/bart-large-mnli"}) labels = ["cooking", "traveling", "dancing"] zero_shot_serving = Bumblebee.Text.zero_shot_classification(model, tokenizer, labels) output = Nx.Serving.run(zero_shot_serving, "One day I will see the world") #=> %{ #=> predictions: [ #=> %{label: "cooking", score: 0.0070497458800673485}, #=> %{label: "traveling", score: 0.985000491142273}, #=> %{label: "dancing", score: 0.007949736900627613} #=> ] #=> } """ @spec zero_shot_classification( Bumblebee.model_info(), Bumblebee.Tokenizer.t(), list(String.t()), keyword() ) :: Nx.Serving.t() defdelegate zero_shot_classification(model_info, tokenizer, labels, opts \\ []), to: Bumblebee.Text.ZeroShotClassification end