Merge branch 'elixir-nx:main' into bool_tensors

Author: gregszumel

examples/distilbert/README.md

@@ -0,0 +1,9 @@
+# DistilBert exported to ONNX with HuggingFace transformers
+
+### Running
+
+Run `python export.py` to create the ONNX model for distilbert/distilbert-base-uncased-finetuned-sst-2-english, then `mix run` the `distilbert_classification.exs` script.
+
+### Labels
+
+When exporting the model from huggingface transformers to ONNX, a `config.json` file is added to the chosen directory. This file has the id to label mappings and you can extract them directly to give a label to the input, as shwon in `distilbert_classification.exs`. 

examples/distilbert/distilbert_classification.exs

@@ -0,0 +1,34 @@
+defmodule Inference do
+  def id_to_label(id) do
+    {:ok, config_json} = File.read("./models/distilbert-onnx/config.json")
+    {:ok, %{"id2label" => id2label}} = Jason.decode(config_json)
+    Map.get(id2label, to_string(id))
+  end
+
+  def run() do
+    model = Ortex.load("./models/distilbert-onnx/model.onnx")
+
+    text =
+      "the movie had a lot of nuance and interesting artistic choices, would like to see more support in the industry for these types of productions"
+
+    {:ok, tokenizer} = Tokenizers.Tokenizer.from_file("./models/distilbert-onnx/tokenizer.json")
+    {:ok, encoding} = Tokenizers.Tokenizer.encode(tokenizer, text)
+
+    input = Nx.tensor([Tokenizers.Encoding.get_ids(encoding)])
+    mask = Nx.tensor([Tokenizers.Encoding.get_attention_mask(encoding)])
+
+    {output} = Ortex.run(model, {input, mask})
+
+    IO.inspect(output)
+
+    IO.inspect(
+      output
+      |> Nx.backend_transfer()
+      |> Nx.argmax()
+      |> Nx.to_number()
+      |> id_to_label()
+    )
+  end
+end
+
+Inference.run()

examples/distilbert/export.py

@@ -0,0 +1,20 @@
+"""
+### Install dependencies:
+
+   $ pip install transformers
+   $ pip install optimum
+   $ pip install "transformers[onnx]"
+
+"""
+
+from transformers import DistilBertTokenizer
+from optimum.onnxruntime import ORTModelForSequenceClassification
+
+save_directory = "./models/distilbert-onnx/"
+
+tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+model = ORTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english", export=True)
+print(model)
+
+model.save_pretrained(save_directory)
+tokenizer.save_pretrained(save_directory)

lib/ortex/native.ex

@@ -1,16 +1,27 @@
 defmodule Ortex.Native do
   @moduledoc false
 
+  @rustler_version Application.spec(:rustler, :vsn) |> to_string() |> Version.parse!()
+
   # We have to compile the crate before `use Rustler` compiles the crate since
   # cargo downloads the onnxruntime shared libraries and they are not available
   # to load or copy into Elixir's during the on_load or Elixir compile steps.
   # In the future, this may be configurable in Rustler.
-  Rustler.Compiler.compile_crate(__MODULE__, otp_app: :ortex, crate: :ortex)
+  if Version.compare(@rustler_version, "0.30.0") in [:gt, :eq] do
+    Rustler.Compiler.compile_crate(:ortex, Application.compile_env(:ortex, __MODULE__, []),
+      otp_app: :ortex,
+      crate: :ortex
+    )
+  else
+    Rustler.Compiler.compile_crate(__MODULE__, otp_app: :ortex, crate: :ortex)
+  end
+
   Ortex.Util.copy_ort_libs()
 
   use Rustler,
     otp_app: :ortex,
-    crate: :ortex
+    crate: :ortex,
+    skip_compilation?: true
 
   # When loading a NIF module, dummy clauses for all NIF function are required.
   # NIF dummies usually just error out when called when the NIF is not loaded, as that should never normally happen.

mix.exs

@@ -31,7 +31,7 @@ defmodule Ortex.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:rustler, "~> 0.29.0"},
+      {:rustler, "~> 0.27"},
       {:nx, "~> 0.6"},
       {:tokenizers, "~> 0.4", only: :dev},
       {:ex_doc, "0.29.4", only: :dev, runtime: false},

python/export_resnet.py

@@ -15,4 +15,5 @@
     output_names=["output"],
     dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
     export_params=True,
+    opset_version=19,
 )

python/multi_input.py

@@ -37,5 +37,5 @@ def forward(self, x, y):
         "x": {0: "batch_size"},
         "y": {0: "batch_size"},
     },
-    opset_version=19
+    opset_version=19,
 )