feat/fp16 onnx export (#12)

Author: elcuervo

.github/workflows/tests.yml

@@ -39,7 +39,7 @@ jobs:
         uses: actions/cache@v4
         with:
           path: tmp/models
-          key: ${{ runner.os }}-models-${{ matrix.ruby }}-fp32
+          key: ${{ runner.os }}-models-${{ matrix.ruby }}-variants
           restore-keys: |
             ${{ runner.os }}-models-${{ matrix.ruby }}
 
@@ -48,5 +48,5 @@ jobs:
 
       - name: Run integration tests
         env:
-          MODEL_FILE: "model.onnx"
+          MODEL_FILE: "model_fp16.onnx"
         run: nix develop . --command bash -lc "bundle exec rake spec:integration"

.gitignore

@@ -5,4 +5,4 @@
 *.gem
 .rspec_status
 Gemfile.lock
-/onnx/gliner2-multi-v1/
+/onnx/onnx/

README.md

@@ -104,7 +104,7 @@ result = Gliner[{ "order" => ["status::[pending|processing|shipped]::str"] }]["S
 This implementation expects a directory containing:
 
 - `tokenizer.json`
-- `model.onnx` or `model_int8.onnx`
+- `model.onnx`, `model_fp16.onnx`, or `model_int8.onnx`
 - (optional) `config.json` with `max_width` and `max_seq_len`
 
 One publicly available ONNX export is `cuerbot/gliner2-multi-v1` on Hugging Face.

Rakefile

@@ -10,17 +10,20 @@ task default: :spec
 
 namespace :model do
   DEFAULT_REPO_ID = 'cuerbot/gliner2-multi-v1'
-  DEFAULT_MODEL_FILE = 'model_int8.onnx'
+  DEFAULT_MODEL_FILE = 'model_fp16.onnx'
+  DEFAULT_MODEL_SUBDIR = 'onnx'
 
-  desc 'Downloads a test model to tmp/ (REPO_ID=... MODEL_FILE=model_int8.onnx)'
+  desc 'Downloads a test model to tmp/ (REPO_ID=... MODEL_FILE=model_fp16.onnx)'
   task :pull do
     repo_id = ENV['REPO_ID'] || DEFAULT_REPO_ID
     model_file = ENV['MODEL_FILE'] || DEFAULT_MODEL_FILE
+    model_subdir = ENV['MODEL_SUBDIR'] || DEFAULT_MODEL_SUBDIR
 
     dir = File.expand_path("tmp/models/#{repo_id.tr('/', '__')}", __dir__)
     FileUtils.mkdir_p(dir)
 
     base = "https://huggingface.co/#{repo_id}/resolve/main"
+    base = "#{base}/#{model_subdir}" unless model_subdir.nil? || model_subdir.empty?
     files = ['tokenizer.json', 'config.json', model_file]
     client = HTTPX.plugin(:follow_redirects).with(max_redirects: 5)
 
@@ -44,6 +47,7 @@ namespace :spec do
   task :integration do
     repo_id = ENV['REPO_ID'] || DEFAULT_REPO_ID
     model_file = ENV['MODEL_FILE'] || DEFAULT_MODEL_FILE
+    model_subdir = ENV['MODEL_SUBDIR'] || DEFAULT_MODEL_SUBDIR
 
     Rake::Task['model:pull'].invoke unless ENV['GLINER_MODEL_DIR'] && !ENV['GLINER_MODEL_DIR'].empty?
 
@@ -53,7 +57,8 @@ namespace :spec do
     env = {
       'GLINER_INTEGRATION' => '1',
       'GLINER_MODEL_DIR' => model_dir,
-      'GLINER_MODEL_FILE' => model_file
+      'GLINER_MODEL_FILE' => model_file,
+      'GLINER_MODEL_SUBDIR' => model_subdir
     }
     sh env, 'rspec', 'spec/integration_spec.rb'
   end

bin/console

@@ -12,13 +12,15 @@ require "httpx"
 require "irb"
 
 DEFAULT_REPO_ID = "cuerbot/gliner2-multi-v1"
-DEFAULT_MODEL_FILE = "model_int8.onnx"
+DEFAULT_MODEL_FILE = "model_fp16.onnx"
+DEFAULT_MODEL_SUBDIR = "onnx"
 
-def ensure_model_dir!(repo_id:, model_file:)
+def ensure_model_dir!(repo_id:, model_file:, model_subdir:)
   dir = File.expand_path("../tmp/models/#{repo_id.tr('/', '__')}", __dir__)
   FileUtils.mkdir_p(dir)
 
   base = "https://huggingface.co/#{repo_id}/resolve/main"
+  base = "#{base}/#{model_subdir}" unless model_subdir.nil? || model_subdir.empty?
   files = ["tokenizer.json", "config.json", model_file]
 
   files.each do |file|
@@ -40,13 +42,14 @@ end
 model_dir = ARGV[0] || ENV["GLINER_MODEL_DIR"]
 repo_id = ENV["GLINER_REPO_ID"] || DEFAULT_REPO_ID
 model_file = ENV["GLINER_MODEL_FILE"] || DEFAULT_MODEL_FILE
+model_subdir = ENV["GLINER_MODEL_SUBDIR"] || DEFAULT_MODEL_SUBDIR
 
 if model_dir && !model_dir.empty?
   $gliner_model = Gliner.load(model_dir, file: model_file)
 else
   begin
     require "fileutils"
-    model_dir = ensure_model_dir!(repo_id: repo_id, model_file: model_file)
+    model_dir = ensure_model_dir!(repo_id: repo_id, model_file: model_file, model_subdir: model_subdir)
     $gliner_model = Gliner.load(model_dir, file: model_file)
   rescue => e
     warn "No model loaded (auto-download failed: #{e.class}: #{e.message})"

onnx/README.md

@@ -1,6 +1,14 @@
+---
+license: mit
+base_model:
+- fastino/gliner2-multi-v1
+---
+
 # GLiNER2 ONNX export
+## gliner2-multi-v1 ONNX
 
 - `model.onnx` (FP32 export)
+- `model_fp16.onnx` (FP16 weights converted from FP32)
 - `model_int8.onnx` (dynamic INT8 quantization via onnxruntime)
 - tokenizer files copied verbatim from the HF model
 - a small `config.json` describing runtime constraints
@@ -12,6 +20,7 @@ The export script follows the same design:
 - Inputs: `input_ids`, `attention_mask` (optionally `token_type_ids`)
 - Output: `span_logits`
 - Export with `torch.onnx.export` (opset 19) and dynamic batch/sequence axes
+- Convert FP32 weights to FP16 with `convert_float_to_float16`
 - Quantize with `onnxruntime.quantization.quantize_dynamic(QInt8)`
 
 ## Usage
@@ -47,11 +56,13 @@ validation or to load the quantized model, use:
 pipenv run python export.py --no-validate
 pipenv run python export.py --no-validate-extraction
 pipenv run python export.py --validate-quantized
+pipenv run python export.py --no-fp16
 ```
 
 The output directory will include:
 
 - `model.onnx`
+- `model_fp16.onnx`
 - `model_int8.onnx`
 - `tokenizer.json`
 - `tokenizer_config.json`

onnx/export.py

@@ -1,11 +1,13 @@
 #!/usr/bin/env python3
 import argparse
+import inspect
 import json
 import shutil
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Dict, List, Sequence, Tuple
 
+import onnx
 import torch
 from huggingface_hub import snapshot_download
 from onnxruntime.quantization import quantize_dynamic, QuantType
@@ -40,6 +42,7 @@ class ExportConfig:
     max_seq_len: int
     opset: int
     include_token_type_ids: bool
+    fp16: bool
     quantize: bool
     validate: bool
     validate_quantized: bool
@@ -151,6 +154,33 @@ def export_onnx(
     )
 
 
+def load_fp16_converter():
+    try:
+        from onnxruntime.transformers.float16 import convert_float_to_float16
+
+        return convert_float_to_float16
+    except ImportError:
+        try:
+            from onnxconverter_common.float16 import convert_float_to_float16
+
+            return convert_float_to_float16
+        except ImportError as exc:
+            raise RuntimeError(
+                "FP16 conversion requires onnxruntime.transformers.float16 or onnxconverter-common."
+            ) from exc
+
+
+def export_fp16(onnx_path: Path, fp16_path: Path) -> None:
+    converter = load_fp16_converter()
+    model = onnx.load(onnx_path.as_posix())
+    params = inspect.signature(converter).parameters
+    kwargs = {}
+    if "keep_io_types" in params:
+        kwargs["keep_io_types"] = True
+    fp16_model = converter(model, **kwargs)
+    onnx.save(fp16_model, fp16_path.as_posix())
+
+
 def export(config: ExportConfig) -> None:
     config.output_dir.mkdir(parents=True, exist_ok=True)
 
@@ -177,6 +207,9 @@ def export(config: ExportConfig) -> None:
         opset=config.opset,
     )
 
+    if config.fp16:
+        export_fp16(onnx_path, config.output_dir / "model_fp16.onnx")
+
     if config.quantize:
         quantize_dynamic(
             onnx_path.as_posix(),
@@ -243,6 +276,11 @@ def parse_args() -> ExportConfig:
         action="store_true",
         help="Include token_type_ids input in the ONNX graph.",
     )
+    parser.add_argument(
+        "--no-fp16",
+        action="store_true",
+        help="Skip FP16 conversion step.",
+    )
     parser.add_argument(
         "--no-validate",
         action="store_true",
@@ -270,6 +308,7 @@ def parse_args() -> ExportConfig:
         max_seq_len=args.max_seq_len,
         opset=args.opset,
         include_token_type_ids=args.include_token_type_ids,
+        fp16=not args.no_fp16,
         quantize=not args.no_quantize,
         validate=not args.no_validate,
         validate_quantized=args.validate_quantized,

spec/integration_spec.rb

@@ -6,12 +6,14 @@
 
 describe 'Gliner Integration', if: ENV.key?('GLINER_INTEGRATION') do
   REPO_ID = 'cuerbot/gliner2-multi-v1'
+  REPO_SUBDIR = 'onnx'
+  MODEL_FILE = 'model_fp16.onnx'
 
   describe 'real model inference' do
     context 'with entities extraction' do
       it 'extracts entities correctly' do
         model_dir = ensure_model_dir!
-        Gliner.load(model_dir)
+        Gliner.load(model_dir, file: MODEL_FILE)
 
         text = 'Apple CEO Tim Cook announced iPhone 15 in Cupertino yesterday.'
         labels = {
@@ -35,7 +37,7 @@
     context 'with text classification' do
       it 'classifies sentiment correctly' do
         model_dir = ensure_model_dir!
-        Gliner.load(model_dir)
+        Gliner.load(model_dir, file: MODEL_FILE)
 
         text = 'This laptop has amazing performance but terrible battery life!'
         schema = { 'sentiment' => %w[positive negative neutral] }
@@ -49,7 +51,7 @@
     context 'with structured extraction' do
       it 'extracts JSON structure correctly' do
         model_dir = ensure_model_dir!
-        Gliner.load(model_dir)
+        Gliner.load(model_dir, file: MODEL_FILE)
 
         text = 'iPhone 15 Pro Max with 256GB storage, A17 Pro chip, priced at $1199.'
         complex_schema = {
@@ -72,7 +74,7 @@
 
       it 'supports choices and multiple instances' do
         model_dir = ensure_model_dir!
-        Gliner.load(model_dir)
+        Gliner.load(model_dir, file: MODEL_FILE)
 
         text = <<~TEXT
           Transaction 1
@@ -130,13 +132,13 @@ def ensure_model_dir!
 
     download(hf_resolve_url('tokenizer.json').to_s, File.join(dir, 'tokenizer.json'))
     download(hf_resolve_url('config.json').to_s, File.join(dir, 'config.json'))
-    download(hf_resolve_url('model_int8.onnx').to_s, File.join(dir, 'model_int8.onnx'))
+    download(hf_resolve_url(MODEL_FILE).to_s, File.join(dir, MODEL_FILE))
 
     dir
   end
 
   def hf_resolve_url(filename)
-    "https://huggingface.co/#{REPO_ID}/resolve/main/#{filename}"
+    "https://huggingface.co/#{REPO_ID}/resolve/main/#{REPO_SUBDIR}/#{filename}"
   end
 
   def download(url, dest)