elixir-nx/bumblebee Examples - NER

Mix.install([
  {:bumblebee, "~> 0.3.1"},
  {:nx, "~> 0.6.0"},
  {:exla, "~> 0.6.0"},
  {:axon, "~> 0.6.0"},
  {:kino, "~> 0.10.0"}
])

Nx.global_default_backend(EXLA.Backend)

Introduction

In this notebook we go through a number of examples to get a quick overview of what Bumblebee brings to the table.

Image classification

Let’s start with image classification. First, we load a pre-trained ResNet-50 model from a HuggingFace repository. We also load the corresponding featurizer for preprocessing input images.

{:ok, resnet} = Bumblebee.load_model({:hf, "microsoft/resnet-50"})
{:ok, featurizer} = Bumblebee.load_featurizer({:hf, "microsoft/resnet-50"})

:ok

Next, we use the high-level API to build an end-to-end task definition on top of the model we just loaded. We will also need an image to work with, so let’s show an image input.

serving = Bumblebee.Vision.image_classification(resnet, featurizer)

image_input = Kino.Input.image("Image", size: {224, 224})

Bumblebee implements end-to-end tasks using Nx.Serving. With serving we can choose to either do a one-off run, or to start a supervised process that automatically batches multiple inference requests. Thanks to this abstraction we can do quick experimentation and then plug the task into a production app with minimal effort.

In this case we will do the one-off run for the selected image:

image = Kino.Input.read(image_input)

# Build a tensor from the raw pixel data
image =
  image.data
  |> Nx.from_binary(:u8)
  |> Nx.reshape({image.height, image.width, 3})

Nx.Serving.run(serving, image)

Manual inference

Note that we are dealing with regular Axon models and the high-level API is just a convenience. If you need full control over the inference flow, you can do it manually. In this case, we would pass the image through the featurizer to get normalized model inputs, then we would run the model and finally extract the most probable label.

inputs = Bumblebee.apply_featurizer(featurizer, image)
outputs = Axon.predict(resnet.model, resnet.params, inputs)

id = outputs.logits |> Nx.argmax() |> Nx.to_number()
resnet.spec.id_to_label[id]

You can try a number of other models, just replace the repository id with one of these:

• facebook/convnext-tiny-224 (ConvNeXT)

• google/vit-base-patch16-224 (ViT)

• facebook/deit-base-distilled-patch16-224 (DeiT)

Fill-mask

Now time for some text processing. Specifically, we want to fill in the missing word in a sentence. This time we load the BERT model together with a matching tokenizer. We will use the tokenizer to preprocess our text input.

{:ok, bert} = Bumblebee.load_model({:hf, "bert-base-uncased"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "bert-base-uncased"})

serving = Bumblebee.Text.fill_mask(bert, tokenizer)

text_input = Kino.Input.text("Sentence with mask", default: "The capital of [MASK] is Paris.")

text = Kino.Input.read(text_input)

Nx.Serving.run(serving, text)

Again, you can try other models, such as albert-base-v2 or roberta-base.

Text classification

In this example we will analyze text sentiment.

We will use the BERTweet model, trained to classify text into one of three categories: positive (POS), negative (NEG) or neutral (NEU).

{: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_input = Kino.Input.text("Text", default: "This cat is so cute.")

Note: this time we need to load a matching tokenizer from a different repository.

text = Kino.Input.read(text_input)
Nx.Serving.run(serving, text)

Named-entity recognition

In this section we look at token classification, more specifically named-entity recognition (NER), where the objective is to identify and categorize entities in text. We will once again use a fine-tuned BERT model.

{:ok, bert} = Bumblebee.load_model({:hf, "dslim/bert-base-NER"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "bert-base-cased"})

serving = Bumblebee.Text.token_classification(bert, tokenizer, aggregation: :same)

text_input =
  Kino.Input.text("Text",
    default: "Rachel Green works at Ralph Lauren in New York City in the sitcom Friends"
  )

text = Kino.Input.read(text_input)
Nx.Serving.run(serving, text)

Text generation

Generation is where things get even more exciting. In this section w will use the extremely popular GPT-2 model to generate text continuation.

Generation generally is an iterative process, where the model predicts the sentence token by token, adhering to some constraints. Again, we will make use of a higher-level API based on Nx.Serving.

{:ok, gpt2} = Bumblebee.load_model({:hf, "gpt2"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "gpt2"})
{:ok, generation_config} = Bumblebee.load_generation_config({:hf, "gpt2"})

serving = Bumblebee.Text.generation(gpt2, tokenizer, generation_config)

text_input = Kino.Input.text("Text", default: "Yesterday, I was reading a book and")

text = Kino.Input.read(text_input)
Nx.Serving.run(serving, text)

There is also gpt2-medium and gpt2-large - heavier versions of the model with much more parameters.

Question answering

Another text-related task is question answering, where the objective is to retrieve the answer to a question based on a given text. We will work with a RoBERTa model trained to do just that.

{:ok, roberta} = Bumblebee.load_model({:hf, "deepset/roberta-base-squad2"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "roberta-base"})

serving = Bumblebee.Text.question_answering(roberta, tokenizer)

question_input =
  Kino.Input.text("Question",
    default: "Which name is also used to describe the Amazon rainforest in English?"
  )

context_input =
  Kino.Input.textarea("Context",
    default:
      ~s/The Amazon rainforest (Portuguese: Floresta Amazônica or Amazônia; Spanish: Selva Amazónica, Amazonía or usually Amazonia; French: Forêt amazonienne; Dutch: Amazoneregenwoud), also known in English as Amazonia or the Amazon Jungle, is a moist broadleaf forest that covers most of the Amazon basin of South America. This basin encompasses 7,000,000 square kilometres (2,700,000 sq mi), of which 5,500,000 square kilometres (2,100,000 sq mi) are covered by the rainforest. This region includes territory belonging to nine nations. The majority of the forest is contained within Brazil, with 60% of the rainforest, followed by Peru with 13%, Colombia with 10%, and with minor amounts in Venezuela, Ecuador, Bolivia, Guyana, Suriname and French Guiana. States or departments in four nations contain "Amazonas" in their names. The Amazon represents over half of the planet's remaining rainforests, and comprises the largest and most biodiverse tract of tropical rainforest in the world, with an estimated 390 billion individual trees divided into 16,000 species./
  )

Kino.Layout.grid([question_input, context_input])

question = Kino.Input.read(question_input)
context = Kino.Input.read(context_input)
Nx.Serving.run(serving, %{question: question, context: context})

Final notes

The examples we covered should give you a good idea of what Bumblebee is about. We are excited about enabling easy access to the pre-trained, powerful deep learning models in Elixir. We are actively working on adding more models and high-level APIs, so stay tuned 🚀