Bumblebee - New Text Generation Model

TwistingTwists
Tue, 13 Dec 2022 12:37:14 +0000
Nx Forum

I plan to use the following model

{:ok, bertqca} = Bumblebee.load_model({:hf, "voidful/bart-eqg-question-generator"})

with tokenizer

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

But on running

serving = Bumblebee.Text.text_classification(bertqca, tokenizer)

I get

** (ArgumentError) expected a model with architecture :for_sequence_classification, got :for_conditional_generation
    (bumblebee 0.1.0) lib/bumblebee/shared.ex:209: Bumblebee.Shared.validate_architecture!/2
    (bumblebee 0.1.0) lib/bumblebee/text/text_classification.ex:9: Bumblebee.Text.TextClassification.text_classification/3
    /Documents/qca.livemd#cell:tzn72neqi6ygxd42wtiunchytys3yjz2:1: (file)

So, following module does not support :for_conditional_generation .

github.com

elixir-nx/bumblebee/blob/8331b0f48766538be865e72384126e1aeb1c663b/lib/bumblebee/text/text_classification.ex#L9

1. defmodule Bumblebee.Text.TextClassification do


2.   @moduledoc false






4.   alias Bumblebee.Utils


5.   alias Bumblebee.Shared






7.   def text_classification(model_info, tokenizer, opts \\ []) do


8.     %{model: model, params: params, spec: spec} = model_info


9.     Shared.validate_architecture!(spec, :for_sequence_classification)


10.     opts = Keyword.validate!(opts, [:compile, top_k: 5, defn_options: []])






12.     top_k = opts[:top_k]


13.     compile = opts[:compile]


14.     defn_options = opts[:defn_options]






16.     batch_size = compile[:batch_size]


17.     sequence_length = compile[:sequence_length]






19.     if compile != nil and (batch_size == nil or sequence_length == nil) do

I am new to machine learning. I would like to extend bumblebee to support above model.
Please show the way. @josevalim

Willing to raise a PR.

Thank you.

4 posts - 2 participants