BERT Question and Answer

Mix.install([
  {:tflite_elixir, "~> 0.3.3"},
  {:req, "~> 0.3.0"},
  {:kino, "~> 0.9.0"}
])

How it works

The model can be used to build a system that can answer users’ questions in natural language. It was created using a pre-trained BERT model fine-tuned on SQuAD 1.1 dataset.

BERT, or Bidirectional Encoder Representations from Transformers, is a method of pre-training language representations which obtains state-of-the-art results on a wide array of Natural Language Processing tasks.

This app uses a compressed version of BERT, MobileBERT, that runs 4x faster and has 4x smaller model size.

SQuAD, or Stanford Question Answering Dataset, is a reading comprehension dataset consisting of articles from Wikipedia and a set of question-answer pairs for each article.

The model takes a passage and a question as input, then returns a segment of the passage that most likely answers the question. It requires semi-complex pre-processing including tokenization and post-processing steps that are described in the BERT paper and implemented in the sample app.

alias TFLiteElixir.MobileBert

Download model

Download the pre-trained TensorFlow Lite MobileBert model.

# /data is the writable portion of a Nerves system
downloads_dir =
  if Code.ensure_loaded?(Nerves.Runtime), do: "/data/livebook", else: System.tmp_dir!()

download = fn url, save_as ->
  save_as = Path.join(downloads_dir, save_as)
  unless File.exists?(save_as), do: Req.get!(url, output: save_as)
  save_as
end

data_files =
  [
    mobiler_bert: {
        "https://tfhub.dev/tensorflow/lite-model/mobilebert/1/metadata/1?lite-format=tflite",
        "mobilebert.tflite"
    },
  ]
  |> Enum.map(fn {key, {url, save_as}} -> {key, download.(url, save_as)} end)
  |> Map.new()

data_files
|> Enum.map(fn {k, v} -> [name: k, location: v] end)
|> Kino.DataTable.new(name: "Data files")

Load MobileBert

alias TFLiteElixir.MobileBert
{:ok, bert} = MobileBert.init(data_files.mobile_bert)

Example

Passage (Input)

content = """
Google LLC is an American multinational technology company 
that specializes in Internet-related services and products, 
which include online advertising technologies, search engine, 
cloud computing, software, and hardware. It is considered one 
of the Big Four technology companies, alongside Amazon, Apple, 
and Facebook.

Google was founded in September 1998 by Larry Page and Sergey 
Brin while they were Ph.D. students at Stanford University in
California. Together they own about 14 percent of its shares
and control 56 percent of the stockholder voting power through
supervoting stock. They incorporated Google as a California
privately held company on September 4, 1998, in California.
Google was then reincorporated in Delaware on October 22, 2002.
An initial public offering (IPO) took place on August 19, 2004,
and Google moved to its headquarters in Mountain View, California,
nicknamed the Googleplex. In August 2015, Google announced plans
to reorganize its various interests as a conglomerate called
Alphabet Inc. Google is Alphabet's leading subsidiary and will
continue to be the umbrella company for Alphabet's Internet
interests. Sundar Pichai was appointed CEO of Google, replacing
Larry Page who became the CEO of Alphabet.
"""

:ok

Question (Input)

query = "Who is the CEO of Google?"

Answer (Output)

MobileBert.run(bert, query, content)
|> Enum.map(fn {score, answer} -> [score: Float.round(score, 6), answer: answer] end)
|> Kino.DataTable.new(name: "Answer")