Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save barkerja/c84f616e5aa60b0fff35c1d49e6a1476 to your computer and use it in GitHub Desktop.
Save barkerja/c84f616e5aa60b0fff35c1d49e6a1476 to your computer and use it in GitHub Desktop.
This notebook showcases two new neural network tasks added to Livebook Kino Bumblebee v0.1.3

What's new in Kino Bumblebee v0.1.3

Mix.install(
  [
    {:kino_bumblebee, "~> 0.1.3"},
    {:exla, "~> 0.4.1"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

Token classification with Livebook and Bumblebee

{:ok, model_info} = Bumblebee.load_model({:hf, "dslim/bert-base-NER"}, log_params_diff: false)

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

serving =
  Bumblebee.Text.token_classification(model_info, tokenizer,
    aggregation: :same,
    compile: [batch_size: 1, sequence_length: 100],
    defn_options: [compiler: EXLA]
  )

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

form = Kino.Control.form([text: text_input], submit: "Run")
frame = Kino.Frame.new()

form
|> Kino.Control.stream()
|> Kino.listen(fn %{data: %{text: text}} ->
  Kino.Frame.render(frame, Kino.Markdown.new("Running..."))
  output = Nx.Serving.run(serving, text)
  Kino.Frame.render(frame, Kino.Bumblebee.HighlightedText.new(text, output.entities))
end)

Kino.Layout.grid([form, frame], boxed: true, gap: 16)

Fill mask with Livebook and Bumblebee

{:ok, model_info} = Bumblebee.load_model({:hf, "distilroberta-base"}, log_params_diff: false)

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

serving =
  Bumblebee.Text.fill_mask(model_info, tokenizer,
    compile: [batch_size: 1, sequence_length: 100],
    defn_options: [compiler: EXLA]
  )

text_input = Kino.Input.textarea("Text", default: "Elixir is a [MASK] programming language.")

form = Kino.Control.form([text: text_input], submit: "Run")
frame = Kino.Frame.new()

form
|> Kino.Control.stream()
|> Kino.listen(fn %{data: %{text: text}} ->
  one_mask? = match?([_, _], String.split(text, "[MASK]"))

  if one_mask? do
    Kino.Frame.render(frame, Kino.Markdown.new("Running..."))
    output = Nx.Serving.run(serving, text)

    output.predictions
    |> Enum.map(&{&1.token, &1.score})
    |> Kino.Bumblebee.ScoredList.new()
    |> then(&Kino.Frame.render(frame, &1))
  else
    Kino.Frame.render(
      frame,
      Kino.Markdown.new("The text must include exactly one [MASK].")
    )
  end
end)

Kino.Layout.grid([form, frame], boxed: true, gap: 16)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment