diff --git a/lib/bumblebee.ex b/lib/bumblebee.ex index 28314c44..1d362c69 100644 --- a/lib/bumblebee.ex +++ b/lib/bumblebee.ex @@ -100,6 +100,7 @@ defmodule Bumblebee do "BertForSequenceClassification" => {Bumblebee.Text.Bert, :for_sequence_classification}, "BertForTokenClassification" => {Bumblebee.Text.Bert, :for_token_classification}, "BertLMHeadModel" => {Bumblebee.Text.Bert, :for_causal_language_modeling}, + "BoundaryExtractor" => {Bumblebee.Text.Gliner, :boundary}, "BertModel" => {Bumblebee.Text.Bert, :base}, "BlenderbotForConditionalGeneration" => {Bumblebee.Text.Blenderbot, :for_conditional_generation}, diff --git a/lib/bumblebee/text.ex b/lib/bumblebee/text.ex index c92371d2..1984ffc6 100644 --- a/lib/bumblebee/text.ex +++ b/lib/bumblebee/text.ex @@ -3,6 +3,65 @@ defmodule Bumblebee.Text do High-level tasks related to text processing. """ + @doc """ + Loads the GLiNER 2.5 encoder, boundary head, span scorer, and tokenizer. + + The repository follows `t:Bumblebee.repository/0`, including local directories + and pinned Hugging Face revisions. It must include `encoder_config/config.json`. + Only the boundary architecture with first-token pooling and a shared candidate + pool is supported. Returns `{:ok, extractor}` or a loading error. + + ## Options + + * `:defn_options` - numerical compiler options, such as `[compiler: EXLA]` + * `:backend` - backend on which to load parameters + * `:log_params_diff` - log checkpoint parameter differences + * `:sequence_length` - optional fixed padded encoder length + * `:max_sequence_length` - reject longer packed inputs; defaults to 4096. + Inputs include both the label schema and document. No truncation is applied. + + """ + defdelegate load_entity_extraction(repository, opts \\ []), + to: Bumblebee.Text.GlinerEntityExtraction, + as: :load + + @doc """ + Builds an entity extraction serving with labels supplied for each document. + + Accepts `%{text: "Apple hired Alice.", labels: ["company", "person"]}`. + Returns `%{"entities" => %{"company" => ["Apple"], "person" => ["Alice"]}}`. + Repeated labels are deduplicated, and an empty label list returns an empty + map (`%{}`) without model inference. + + Processes one document per batch. Keep `batch_size: 1` when starting the + serving process; document batching and accelerator partitioning are not + currently supported. Distinct packed lengths may trigger separate compilations. + + ## Options + + * `:threshold` - inclusive entity confidence cutoff, defaults to 0.5 + * `:include_spans` - include `"text"`, `"start"`, and `"end"` fields. + Offsets count Unicode codepoints in the original document, matching GLiNER. + Defaults to `false` + * `:include_confidence` - include a `"confidence"` field, defaults to `false` + + Scores are numerically approximate across backends. Decisions extremely close + to a cutoff can differ from another execution of the reference model. + + ## Example + + {:ok, extractor} = Bumblebee.Text.load_entity_extraction( + {:hf, "fastino/gliner2.5-base-v1", revision: "78cea040597df251eedefa9d7ee2a756af39fe64"}, + defn_options: [compiler: EXLA] + ) + serving = Bumblebee.Text.entity_extraction(extractor, include_spans: true) + Nx.Serving.run(serving, %{text: "Apple hired Alice.", labels: ["company", "person"]}) + + """ + defdelegate entity_extraction(extractor, opts \\ []), + to: Bumblebee.Text.GlinerEntityExtraction, + as: :new + @type token_classification_input :: String.t() @type token_classification_output :: %{entities: list(token_classification_entity())} diff --git a/lib/bumblebee/text/gliner.ex b/lib/bumblebee/text/gliner.ex new file mode 100644 index 00000000..e51436a4 --- /dev/null +++ b/lib/bumblebee/text/gliner.ex @@ -0,0 +1,234 @@ +defmodule Bumblebee.Text.Gliner do + alias Bumblebee.Shared + alias Bumblebee.Text.{DebertaV2, Gliner} + + @moduledoc """ + GLiNER 2.5 boundary-based named entity extraction model stages. + + The `:encoder` stage wraps `Bumblebee.Text.DebertaV2`; `:boundary` computes + boundary states and marginals from routed text and query states; and + `:span_scoring` scores a host-selected shared candidate pool. These stages + support inference with one unpadded document and an arbitrary number of labels. + + Load the nested encoder configuration from `encoder_config/config.json` + and set `:encoder_spec` when loading `:encoder`. Set `:hidden_size` on the + other stages to match that encoder. The default hidden size is 768. + + This implementation supports first-token word pooling, a shared candidate + pool, span content, inside evidence, abstention and flat overlap resolution. + Unsupported architectural variants raise during configuration. Classification, + records and relations are outside this model's extraction API. + + Boundary configuration options use the names in the checkpoint's + `boundary_head` object. Model outputs use the same string keys as their + downstream stage inputs. `:span_scoring` returns a tensor shaped + `{1, number_of_labels, number_of_candidates}`. + """ + + defstruct architecture: :boundary, + encoder_spec: nil, + hidden_size: 768, + abstention_threshold: 0.5, + adaptive_threshold: false, + boundary_attention_heads: 4, + boundary_attention_layers: 2, + boundary_attention_window: 128, + boundary_dim: 128, + boundary_ffn_multiplier: 2.0, + boundary_refinement_layers: 1, + candidate_attention_layers: 0, + candidate_pool: "shared", + content_dim: 64, + content_soft_max_pool: false, + enable_abstention: true, + enable_span_content: true, + min_pool_per_query: 8, + overlap_policy: "flat", + pair_dim: 128, + pair_temperature: 1.0, + pool_boundary_top_k: 32, + pool_size: 192, + query_attention_layers: 0, + use_inside_evidence: true + + @behaviour Bumblebee.ModelSpec + @behaviour Bumblebee.Configurable + + @impl true + def architectures, do: [:encoder, :boundary, :span_scoring] + + @impl true + def config(spec, opts) do + spec = Shared.put_config_attrs(spec, opts) + + validate_architecture!(spec) + validate_dimensions!(spec) + validate_scoring!(spec) + + spec + end + + defp validate_architecture!(spec) do + required = [ + candidate_pool: "shared", + candidate_attention_layers: 0, + query_attention_layers: 0, + content_soft_max_pool: false, + enable_span_content: true, + enable_abstention: true, + use_inside_evidence: true, + adaptive_threshold: false, + overlap_policy: "flat" + ] + + for {key, value} <- required, Map.fetch!(spec, key) != value do + raise ArgumentError, "unsupported GLiNER option #{key}: #{inspect(Map.fetch!(spec, key))}" + end + end + + defp validate_dimensions!(spec) do + for key <- [ + :hidden_size, + :boundary_dim, + :boundary_attention_heads, + :pair_dim, + :content_dim, + :pool_size, + :pool_boundary_top_k + ] do + value = Map.fetch!(spec, key) + + unless is_integer(value) and value > 0, + do: raise(ArgumentError, "#{key} must be a positive integer") + end + + unless rem(spec.boundary_dim, spec.boundary_attention_heads) == 0, + do: raise(ArgumentError, "boundary_dim must be divisible by boundary_attention_heads") + + for key <- [ + :boundary_attention_layers, + :boundary_refinement_layers, + :boundary_attention_window, + :min_pool_per_query + ] do + value = Map.fetch!(spec, key) + + unless is_integer(value) and value >= 0, + do: raise(ArgumentError, "#{key} must be a nonnegative integer") + end + end + + defp validate_scoring!(spec) do + for key <- [:boundary_ffn_multiplier, :pair_temperature] do + value = Map.fetch!(spec, key) + unless is_number(value) and value > 0, do: raise(ArgumentError, "#{key} must be positive") + end + + unless is_number(spec.abstention_threshold) and spec.abstention_threshold >= 0 and + spec.abstention_threshold <= 1 do + raise ArgumentError, "abstention_threshold must be between zero and one" + end + end + + @doc false + def boundary_config(spec) do + spec + |> Map.from_struct() + |> Map.drop([:architecture, :encoder_spec, :hidden_size]) + |> Map.new(fn {key, value} -> {Atom.to_string(key), value} end) + end + + @impl true + def input_template(%{architecture: :encoder, encoder_spec: %DebertaV2{} = encoder}), + do: DebertaV2.input_template(encoder) + + def input_template(%{architecture: :boundary} = spec) do + %{ + "text_states" => Nx.template({1, 2, spec.hidden_size}, :f32), + "query_states" => Nx.template({1, 1, spec.hidden_size}, :f32) + } + end + + def input_template(%{architecture: :span_scoring} = spec) do + Map.merge(input_template(%{spec | architecture: :boundary}), %{ + "boundary_states" => Nx.template({1, 3, spec.boundary_dim}, :f32), + "starts" => Nx.template({1, spec.pool_size}, :s64), + "ends" => Nx.template({1, spec.pool_size}, :s64), + "compat" => Nx.template({1, spec.pool_size}, :f32), + "valid" => Nx.template({1, spec.pool_size}, :u8), + "start_logits" => Nx.template({1, 1, 3}, :f32), + "end_logits" => Nx.template({1, 1, 3}, :f32), + "inside_logits" => Nx.template({1, 1, 2}, :f32) + }) + end + + @impl true + def model(%{architecture: :encoder, encoder_spec: %DebertaV2{} = encoder}), + do: DebertaV2.model(encoder) + + def model(%{architecture: :encoder}), + do: raise(ArgumentError, "encoder_spec must be a configured DebertaV2 encoder") + + def model(%{architecture: :boundary} = spec), + do: Gliner.BoundaryHead.model(boundary_config(spec), spec.hidden_size) + + def model(%{architecture: :span_scoring} = spec), + do: Gliner.PairScorer.model(boundary_config(spec), spec.hidden_size) + + defimpl Bumblebee.HuggingFace.Transformers.Config do + def load(spec, data) do + unless data["architecture"] == "boundary" and data["token_pooling"] == "first" do + raise ArgumentError, "expected a GLiNER boundary checkpoint with first-token pooling" + end + + boundary = Map.fetch!(data, "boundary_head") + + opts = + for {key, _} <- @for.boundary_config(spec), + Map.has_key?(boundary, key), + do: {String.to_existing_atom(key), boundary[key]} + + @for.config(spec, opts) + end + end + + defimpl Bumblebee.HuggingFace.Transformers.Model do + alias Bumblebee.HuggingFace.Transformers + alias Bumblebee.Utils + + def params_mapping(%{architecture: :encoder, encoder_spec: encoder}) do + Map.new(Transformers.Model.params_mapping(encoder), fn {target, source} -> + {target, + Transformers.Utils.map_params_source_layer_names( + source, + &String.replace_prefix(&1, "deberta.", "encoder.") + )} + end) + end + + def params_mapping(spec) do + for {node, name} <- @for.model(spec) |> Utils.Axon.nodes_with_names(), + node.parameters != [], + into: %{} do + cond do + String.ends_with?(name, ".bos") -> + {name, + %{ + "bos_state" => + {[{String.trim_trailing(name, ".bos"), "bos_state"}], fn [value] -> value end} + }} + + String.ends_with?(name, ".eos") -> + {name, + %{ + "eos_state" => + {[{String.trim_trailing(name, ".eos"), "eos_state"}], fn [value] -> value end} + }} + + true -> + {name, name} + end + end + end + end +end diff --git a/lib/bumblebee/text/gliner/boundary_head.ex b/lib/bumblebee/text/gliner/boundary_head.ex new file mode 100644 index 00000000..42fa23cc --- /dev/null +++ b/lib/bumblebee/text/gliner/boundary_head.ex @@ -0,0 +1,135 @@ +defmodule Bumblebee.Text.Gliner.BoundaryHead do + @moduledoc false + import Nx.Defn + + def model(config, hidden_size) do + text = Axon.input("text_states", shape: {1, nil, hidden_size}) + query = Axon.input("query_states", shape: {1, nil, hidden_size}) + dim = config["boundary_dim"] + prefix = "boundary_head.boundary_encoder" + + left = + Axon.layer(&left_states/3, [text, Axon.param("bos_state", {hidden_size})], + name: prefix <> ".bos" + ) + + right = + Axon.layer(&right_states/3, [text, Axon.param("eos_state", {hidden_size})], + name: prefix <> ".eos" + ) + + left = Axon.dense(left, dim, name: prefix <> ".left_projection") + right = Axon.dense(right, dim, name: prefix <> ".right_projection") + + boundary = + Axon.concatenate([left, right]) + |> Axon.dense(dim, name: prefix <> ".output_projection") + |> Axon.layer_norm(name: prefix <> ".layer_norm") + + boundary = + Enum.reduce(0..(config["boundary_attention_layers"] - 1)//1, boundary, fn index, states -> + name = prefix <> ".attention_blocks.#{index}" + + projected = + states + |> Axon.layer_norm(name: name <> ".norm") + |> Axon.dense(3 * dim, name: name <> ".qkv_projection") + + update = + Axon.layer(&self_attention/2, [projected], + heads: config["boundary_attention_heads"], + window: config["boundary_attention_window"] + ) + |> Axon.dense(dim, name: name <> ".output_projection") + + Axon.add(states, update) + end) + + boundary = + Enum.reduce(0..(config["boundary_refinement_layers"] - 1)//1, boundary, fn index, states -> + name = prefix <> ".refinement_blocks.#{index}" + inner = max(1, trunc(dim * config["boundary_ffn_multiplier"])) + + update = + states + |> Axon.layer_norm(name: name <> ".norm") + |> Axon.dense(2 * inner, name: name <> ".input_projection") + |> Axon.nx(&swiglu/1) + |> Axon.dense(dim, name: name <> ".output_projection") + + Axon.add(states, update) + end) + + start = marginal(boundary, query, dim, "start_boundary_projection", "start_query_projection") + ending = marginal(boundary, query, dim, "end_boundary_projection", "end_query_projection") + inside = marginal(text, query, dim, "inside_text_projection", "inside_query_projection") + null = Axon.dense(query, 1, name: "boundary_head.null_projection") + + Axon.container(%{ + "boundary_states" => boundary, + "start_logits" => start, + "end_logits" => ending, + "inside_logits" => inside, + "null_logits" => null, + "pool_start" => + Axon.dense(boundary, dim, name: "boundary_head.shared_pool_builder.start_projection"), + "pool_end" => + Axon.dense(boundary, dim, name: "boundary_head.shared_pool_builder.end_projection") + }) + end + + defp marginal(states, query, dim, state_name, query_name) do + prefix = "boundary_head.boundary_query_head." + projected_states = Axon.dense(states, dim, name: prefix <> state_name) + projected_query = Axon.dense(query, dim, name: prefix <> query_name) + Axon.layer(&query_scores/3, [projected_states, projected_query]) + end + + defnp(left_states(text, bos, _opts), + do: Nx.concatenate([Nx.reshape(bos, {1, 1, Nx.axis_size(bos, 0)}), text], axis: 1) + ) + + defnp(right_states(text, eos, _opts), + do: Nx.concatenate([text, Nx.reshape(eos, {1, 1, Nx.axis_size(eos, 0)})], axis: 1) + ) + + defn query_scores(states, query, _opts) do + Nx.dot(query, [2], [0], states, [2], [0]) / Nx.sqrt(Nx.axis_size(states, 2)) + end + + defnp swiglu(states) do + {value, gate} = Nx.split(states, 0.5, axis: -1) + value * Axon.Activations.silu(gate) + end + + defn self_attention(projected, opts) do + opts = keyword!(opts, heads: 4, window: 128, mode: :inference) + {batch, length, triple_dim} = Nx.shape(projected) + dim = div(triple_dim, 3) + heads = opts[:heads] + width = div(dim, heads) + + qkv = + projected + |> Nx.reshape({batch, length, 3, heads, width}) + |> Nx.transpose(axes: [2, 0, 3, 1, 4]) + + query = qkv[0] + key = qkv[1] + value = qkv[2] + scores = Nx.dot(query, [3], [0, 1], key, [3], [0, 1]) / Nx.sqrt(width) + ids = Nx.iota({length}) + + allowed = + if opts[:window] > 0, + do: Nx.abs(Nx.new_axis(ids, 1) - Nx.new_axis(ids, 0)) <= opts[:window], + else: Nx.broadcast(1, {length, length}) + + scores = Nx.select(Nx.broadcast(allowed, Nx.shape(scores)), scores, -1.0e4) + weights = Axon.Activations.softmax(scores, axis: -1) + + Nx.dot(weights, [3], [0, 1], value, [2], [0, 1]) + |> Nx.transpose(axes: [0, 2, 1, 3]) + |> Nx.reshape({batch, length, dim}) + end +end diff --git a/lib/bumblebee/text/gliner/candidate_pool.ex b/lib/bumblebee/text/gliner/candidate_pool.ex new file mode 100644 index 00000000..e2322e75 --- /dev/null +++ b/lib/bumblebee/text/gliner/candidate_pool.ex @@ -0,0 +1,84 @@ +defmodule Bumblebee.Text.Gliner.CandidatePool do + @moduledoc false + + def select(output, config) do + starts = rows(output["start_logits"]) + endings = rows(output["end_logits"]) + union_start = starts |> Enum.zip_with(&Enum.max/1) |> List.to_tuple() + union_end = endings |> Enum.zip_with(&Enum.max/1) |> List.to_tuple() + top_start = top_indices(union_start, config["pool_boundary_top_k"]) + top_end = top_indices(union_end, config["pool_boundary_top_k"]) + compatibilities = compatibility(output, top_start, top_end) + + pairs = + for {start, i} <- Enum.with_index(top_start), + {ending, j} <- Enum.with_index(top_end), + ending > start do + %{start: start, end: ending, compat: elem(compatibilities, i * length(top_end) + j)} + end + + global = + Enum.map( + pairs, + &Map.put(&1, :priority, &1.compat + elem(union_start, &1.start) + elem(union_end, &1.end)) + ) + + quota = config["min_pool_per_query"] + + reserved = + Enum.zip_with(starts, endings, fn start_logits, end_logits -> + start_logits = List.to_tuple(start_logits) + end_logits = List.to_tuple(end_logits) + + pairs + |> Enum.sort_by(&(-(&1.compat + elem(start_logits, &1.start) + elem(end_logits, &1.end)))) + |> Enum.take(quota) + |> Enum.with_index(fn pair, rank -> Map.put(pair, :priority, 5000.0 + quota - rank) end) + end) + |> List.flatten() + + selected = + (reserved ++ global) + |> Enum.sort_by(&(-&1.priority)) + |> Enum.uniq_by(&{&1.start, &1.end}) + |> Enum.sort_by(&{-&1.priority, &1.start, &1.end}) + |> Enum.take(config["pool_size"]) + + padding = config["pool_size"] - length(selected) + + %{ + "starts" => + Nx.tensor([Enum.map(selected, & &1.start) ++ List.duplicate(0, padding)], type: :s64), + "ends" => + Nx.tensor([Enum.map(selected, & &1.end) ++ List.duplicate(0, padding)], type: :s64), + "compat" => + Nx.tensor([Enum.map(selected, & &1.compat) ++ List.duplicate(0.0, padding)], type: :f32), + "valid" => + Nx.tensor([List.duplicate(1, length(selected)) ++ List.duplicate(0, padding)], type: :u8) + } + end + + defp rows(tensor) do + tensor |> Nx.to_flat_list() |> Enum.chunk_every(Nx.axis_size(tensor, -1)) + end + + defp top_indices(logits, count) do + logits + |> Tuple.to_list() + |> Enum.with_index() + |> Enum.sort_by(fn {score, index} -> {-score, index} end) + |> Enum.take(count) + |> Enum.map(&elem(&1, 1)) + end + + defp compatibility(output, starts, endings) do + left = Nx.take(output["pool_start"], Nx.tensor(starts), axis: 1) |> Nx.squeeze(axes: [0]) + right = Nx.take(output["pool_end"], Nx.tensor(endings), axis: 1) |> Nx.squeeze(axes: [0]) + + Nx.multiply(Nx.new_axis(left, 1), Nx.new_axis(right, 0)) + |> Nx.sum(axes: [2]) + |> Nx.divide(:math.sqrt(Nx.axis_size(left, 1))) + |> Nx.to_flat_list() + |> List.to_tuple() + end +end diff --git a/lib/bumblebee/text/gliner/decode.ex b/lib/bumblebee/text/gliner/decode.ex new file mode 100644 index 00000000..649b8b20 --- /dev/null +++ b/lib/bumblebee/text/gliner/decode.ex @@ -0,0 +1,118 @@ +defmodule Bumblebee.Text.Gliner.Decode do + @moduledoc false + + def entities(packed, outputs, config, opts \\ []) do + threshold = Keyword.get(opts, :threshold, 0.5) + null = outputs["null_logits"] |> sigmoid() |> Nx.to_flat_list() + + probabilities = + outputs["pair_logits"] + |> Nx.divide(config["pair_temperature"]) + |> sigmoid() + |> Nx.to_flat_list() + + starts = Nx.to_flat_list(outputs["starts"]) + endings = Nx.to_flat_list(outputs["ends"]) + valid = Nx.to_flat_list(outputs["valid"]) + scores = Enum.chunk_every(probabilities, length(starts)) + + entities = + Enum.zip([packed.labels, null, scores]) + |> Map.new(fn {label, abstention, scores} -> + candidates = Enum.zip([scores, starts, endings, valid]) + spans = decode_query(candidates, abstention, config, threshold, packed, opts) + + {label, spans} + end) + + %{"entities" => entities} + end + + defp sigmoid(x) do + Axon.Activations.sigmoid(x) + end + + defp decode_query(candidates, abstention, config, threshold, packed, opts) do + if abstention > config["abstention_threshold"] do + [] + else + candidates + |> Enum.filter(fn {score, _, _, valid} -> valid == 1 and score >= threshold end) + |> Enum.flat_map(&original_span(&1, packed)) + |> resolve_flat() + |> Enum.map(&format_span(&1, packed, opts)) + end + end + + def resolve_flat(spans) do + ordered = + spans + |> Enum.sort_by(&rank/1) + |> Enum.uniq_by(fn {_, start, ending} -> {start, ending} end) + |> Enum.sort_by(fn {score, start, ending} -> {ending, start, -score} end) + + best = + Enum.reduce(ordered, [{0.0, []}], fn {score, start, _} = span, best -> + predecessor = + Enum.reduce_while(ordered, 0, fn + {_, _, ending}, count when ending <= start -> {:cont, count + 1} + _, count -> {:halt, count} + end) + + {previous_score, previous_spans} = Enum.at(best, predecessor) + with_span = {previous_score + score, [span | previous_spans]} + without_span = List.last(best) + best ++ [choose(with_span, without_span)] + end) + + best |> List.last() |> elem(1) |> Enum.sort_by(&rank/1) + end + + defp choose({a, left}, {b, right}) do + cond do + a > b -> {a, left} + a < b -> {b, right} + length(left) > length(right) -> {a, left} + length(left) < length(right) -> {b, right} + Enum.sort(Enum.map(left, &rank/1)) < Enum.sort(Enum.map(right, &rank/1)) -> {a, left} + true -> {b, right} + end + end + + defp rank({score, start, ending}), do: {-score, start, ending} + + # Project before overlap resolution: a synthetic-only span must not compete + # with real entities, and spans differing only by the suffix are duplicates. + defp original_span({score, start, ending, _valid}, packed) do + first = Enum.at(packed.words, start).start + last = min(Enum.at(packed.words, ending - 1).end, packed.original_length) + raw = slice(packed.original_text, first, max(last - first, 0)) + left_trimmed = String.trim_leading(raw) + text = String.trim_trailing(left_trimmed) + + if text == "" do + [] + else + first = first + length(String.codepoints(raw)) - length(String.codepoints(left_trimmed)) + [{score, first, first + length(String.codepoints(text))}] + end + end + + defp format_span({score, char_start, char_end}, packed, opts) do + surface = slice(packed.original_text, char_start, char_end - char_start) + + fields = %{"text" => surface} + + fields = + if opts[:include_spans], + do: Map.merge(fields, %{"start" => char_start, "end" => char_end}), + else: fields + + fields = if opts[:include_confidence], do: Map.put(fields, "confidence", score), else: fields + if opts[:include_spans] || opts[:include_confidence], do: fields, else: surface + end + + defp slice(text, start, length) do + text |> String.codepoints() |> Enum.slice(start, length) |> Enum.join() + end +end diff --git a/lib/bumblebee/text/gliner/pair_scorer.ex b/lib/bumblebee/text/gliner/pair_scorer.ex new file mode 100644 index 00000000..d9e02ec4 --- /dev/null +++ b/lib/bumblebee/text/gliner/pair_scorer.ex @@ -0,0 +1,131 @@ +defmodule Bumblebee.Text.Gliner.PairScorer do + @moduledoc false + import Nx.Defn + @prefix "boundary_head.shared_pool_scorer." + + def model(config, hidden_size) do + dim = config["pair_dim"] + states = Axon.input("boundary_states", shape: {1, nil, config["boundary_dim"]}) + query = Axon.input("query_states", shape: {1, nil, hidden_size}) + text = Axon.input("text_states", shape: {1, nil, hidden_size}) + starts = Axon.input("starts", shape: {1, nil}) + endings = Axon.input("ends", shape: {1, nil}) + compat = Axon.input("compat", shape: {1, nil}) + valid = Axon.input("valid", shape: {1, nil}) + start_logits = Axon.input("start_logits", shape: {1, nil, nil}) + end_logits = Axon.input("end_logits", shape: {1, nil, nil}) + inside_logits = Axon.input("inside_logits", shape: {1, nil, nil}) + start_rep = gather(dense(states, dim, "start_projection"), starts) + end_rep = gather(dense(states, dim, "end_projection"), endings) + + lengths = + Axon.layer(&length_features/4, [starts, endings, text]) |> dense(dim, "length_projection") + + prior = compat |> Axon.nx(&Nx.new_axis(&1, -1)) |> dense(dim, "prior_projection") + values = dense(text, config["content_dim"], "content_pooler.value_projection") + + content = + Axon.layer(&content_mean/4, [values, starts, endings]) + |> Axon.layer_norm(name: @prefix <> "content_pooler.layer_norm") + |> dense(dim, "content_projection") + + candidate = + Axon.add([start_rep, end_rep, lengths, prior, content]) + |> Axon.layer_norm(name: @prefix <> "candidate_norm") + + candidate = Axon.layer(&mask_states/3, [candidate, valid]) + projected_query = dense(query, dim, "query_projection") + film = dense(projected_query, 2 * dim, "film") + + conditioned = + Axon.layer(&condition/3, [candidate, film]) + |> dense(64, "film_output.0") + |> Axon.gelu() + |> dense(1, "film_output.3") + + score = + Axon.layer(&score/7, [ + candidate, + projected_query, + conditioned, + start_logits, + end_logits, + Axon.container({starts, endings}) + ]) + + score = Axon.layer(&inside_evidence/5, [score, inside_logits, starts, endings]) + Axon.layer(&mask_scores/3, [score, valid]) + end + + defp dense(input, units, name), do: Axon.dense(input, units, name: @prefix <> name) + defp gather(input, indices), do: Axon.layer(&gather_rows/3, [input, indices]) + + defnp(gather_rows(input, indices, _opts), + do: Nx.take(input, Nx.squeeze(indices, axes: [0]), axis: 1) + ) + + defnp(mask_states(states, valid, _opts), do: states * Nx.new_axis(valid, -1)) + + defnp length_features(starts, endings, text, _opts) do + length = Nx.as_type(Nx.max(endings - starts, 1), :f32) + Nx.stack([Nx.log1p(length), length / Nx.axis_size(text, 1), Nx.rsqrt(length)], axis: -1) + end + + defnp content_mean(values, starts, endings, _opts) do + prefix = + Nx.concatenate( + [ + Nx.broadcast(0.0, {1, 1, Nx.axis_size(values, 2)}), + Nx.cumulative_sum(Nx.as_type(values, :f64), axis: 1) |> Nx.as_type(:f32) + ], + axis: 1 + ) + + (gather_rows(prefix, endings, []) - gather_rows(prefix, starts, [])) / + Nx.new_axis(Nx.max(endings - starts, 1), -1) + end + + defnp condition(candidate, film, _opts) do + {gamma, beta} = Nx.split(film, 0.5, axis: -1) + Nx.new_axis(candidate, 2) * (1.0 + Nx.new_axis(gamma, 1)) + Nx.new_axis(beta, 1) + end + + defnp score(candidate, query, conditioned, start_logits, end_logits, indices, _opts) do + {starts, endings} = indices + base = Nx.dot(candidate, [2], [0], query, [2], [0]) / Nx.sqrt(Nx.axis_size(candidate, 2)) + + start_scores = + Nx.take(start_logits, Nx.squeeze(starts, axes: [0]), axis: 2) + |> Nx.transpose(axes: [0, 2, 1]) + + end_scores = + Nx.take(end_logits, Nx.squeeze(endings, axes: [0]), axis: 2) + |> Nx.transpose(axes: [0, 2, 1]) + + base + Nx.squeeze(conditioned, axes: [3]) + start_scores + end_scores + end + + defnp inside_evidence(scores, inside, starts, endings, _opts) do + mean = Nx.mean(inside, axes: [2], keep_axes: true) + + prefix = + Nx.concatenate( + [ + Nx.broadcast(0.0, {1, Nx.axis_size(inside, 1), 1}), + Nx.cumulative_sum(Nx.as_type(inside - mean, :f64), axis: 2) |> Nx.as_type(:f32) + ], + axis: 2 + ) + + start_prefix = Nx.take(prefix, Nx.squeeze(starts, axes: [0]), axis: 2) + end_prefix = Nx.take(prefix, Nx.squeeze(endings, axes: [0]), axis: 2) + length = Nx.new_axis(endings - starts, 1) + evidence = (end_prefix - start_prefix + mean * length) / Nx.sqrt(Nx.max(length, 1)) + scores + Nx.transpose(evidence, axes: [0, 2, 1]) + end + + defnp mask_scores(scores, valid, _opts) do + Nx.select(Nx.broadcast(Nx.new_axis(valid, -1), Nx.shape(scores)), scores, -1.0e4) + |> Nx.transpose(axes: [0, 2, 1]) + end +end diff --git a/lib/bumblebee/text/gliner/preprocessor.ex b/lib/bumblebee/text/gliner/preprocessor.ex new file mode 100644 index 00000000..1ab74220 --- /dev/null +++ b/lib/bumblebee/text/gliner/preprocessor.ex @@ -0,0 +1,61 @@ +defmodule Bumblebee.Text.Gliner.Preprocessor do + @moduledoc false + alias Tokenizers.{Encoding, Tokenizer} + + # Python's Unicode IGNORECASE includes these four non-ASCII letters in [a-z]. + @word_pattern ~r"(?:https?://[^\s]+|www\.[^\s]+)|[a-zİıſK0-9._%+-]+@[a-zİıſK0-9.-]+\.[a-zİıſK]{2,}|@[a-zİıſK0-9_]+|[\p{L}\p{N}_]+(?:[-_][\p{L}\p{N}_]+)*|\S"iu + + def words(text) do + Regex.scan(@word_pattern, text, return: :index) + |> Enum.map(fn [{offset, size}] -> + word = binary_part(text, offset, size) + start = text |> binary_part(0, offset) |> String.codepoints() |> length() + + %{ + text: String.downcase(word, :greek), + start: start, + end: start + length(String.codepoints(word)) + } + end) + end + + def pack(text, labels, tokenizer) when is_binary(text) and is_list(labels) do + original_text = text + text = if String.ends_with?(text, [".", "!", "?"]), do: text, else: text <> "." + words = words(text) + schema = ["(", "[P]", "entities", "("] ++ Enum.flat_map(labels, &["[E]", &1]) ++ [")", ")"] + tokens = schema ++ ["[SEP_TEXT]"] ++ Enum.map(words, & &1.text) + + encodings = + Enum.map(tokens, fn token -> + {:ok, encoding} = Tokenizer.encode(tokenizer, token, add_special_tokens: false) + Encoding.get_ids(encoding) + end) + + {positions, _} = + Enum.map_reduce(encodings, 0, fn ids, position -> {position, position + length(ids)} end) + + ids = List.flatten(encodings) + + query_indices = + labels + |> Enum.with_index() + |> Enum.map(fn {_, index} -> Enum.at(positions, 4 + index * 2) end) + + text_indices = Enum.drop(positions, length(schema) + 1) + + %{ + inputs: %{ + "input_ids" => Nx.tensor([ids], type: :s64), + "attention_mask" => Nx.tensor([List.duplicate(1, length(ids))], type: :s64) + }, + text_indices: Nx.tensor([text_indices], type: :s64), + query_indices: Nx.tensor([query_indices], type: :s64), + words: words, + labels: labels, + text: text, + original_text: original_text, + original_length: length(String.codepoints(original_text)) + } + end +end diff --git a/lib/bumblebee/text/gliner_entity_extraction.ex b/lib/bumblebee/text/gliner_entity_extraction.ex new file mode 100644 index 00000000..3db11a3c --- /dev/null +++ b/lib/bumblebee/text/gliner_entity_extraction.ex @@ -0,0 +1,243 @@ +defmodule Bumblebee.Text.GlinerEntityExtraction do + @moduledoc """ + GLiNER 2.5 entity extraction with labels supplied for each document. + + See `Bumblebee.Text.load_entity_extraction/2` and + `Bumblebee.Text.entity_extraction/2` for loading and serving. + """ + @behaviour Nx.Serving + alias Bumblebee.Text.{DebertaV2, Gliner} + alias Gliner.{CandidatePool, Decode, Preprocessor} + + defstruct [ + :encoder, + :boundary, + :scorer, + :tokenizer, + :config, + :defn_options, + :sequence_length, + :max_sequence_length + ] + + @doc false + def load(repository, opts \\ []) do + opts = + Keyword.validate!(opts, [ + :backend, + :log_params_diff, + defn_options: [], + sequence_length: nil, + max_sequence_length: 4096 + ]) + + sequence_length = opts[:sequence_length] + maximum = opts[:max_sequence_length] + + unless is_integer(maximum) and maximum > 0, + do: raise(ArgumentError, "max_sequence_length must be a positive integer") + + unless is_nil(sequence_length) or + (is_integer(sequence_length) and sequence_length > 0 and sequence_length <= maximum), + do: + raise( + ArgumentError, + "sequence_length must be positive and at most max_sequence_length" + ) + + with {:ok, encoder_spec} <- + Bumblebee.load_spec(encoder_repository(repository), + module: DebertaV2, + architecture: :base + ), + {:ok, spec} <- Bumblebee.load_spec(repository, module: Gliner, architecture: :boundary), + {:ok, tokenizer} <- Bumblebee.load_tokenizer(repository, type: :deberta_v2) do + validate_tokenizer!(tokenizer.native_tokenizer) + + spec = + Bumblebee.configure(spec, + hidden_size: encoder_spec.hidden_size, + encoder_spec: encoder_spec + ) + + load_opts = Keyword.take(opts, [:backend, :log_params_diff]) + + with {:ok, encoder} <- + load_stage(repository, spec, :encoder, load_opts, opts[:defn_options]), + {:ok, boundary} <- + load_stage(repository, spec, :boundary, load_opts, opts[:defn_options]), + {:ok, scorer} <- + load_stage(repository, spec, :span_scoring, load_opts, opts[:defn_options]) do + {:ok, + %__MODULE__{ + encoder: encoder, + boundary: boundary, + scorer: scorer, + tokenizer: tokenizer.native_tokenizer, + config: Gliner.boundary_config(spec), + defn_options: opts[:defn_options], + sequence_length: sequence_length, + max_sequence_length: maximum + }} + end + end + end + + defp validate_tokenizer!(tokenizer) do + for token <- ["[P]", "[E]", "[SEP_TEXT]"] do + if Tokenizers.Tokenizer.token_to_id(tokenizer, token) == nil do + raise ArgumentError, "GLiNER tokenizer is missing #{token}" + end + end + end + + defp load_stage(repository, spec, architecture, load_opts, defn_options) do + spec = Bumblebee.configure(spec, architecture: architecture) + + with {:ok, info} <- Bumblebee.load_model(repository, [spec: spec] ++ load_opts) do + {_, predict} = Axon.build(info.model, defn_options) + {:ok, {predict, info.params}} + end + end + + defp encoder_repository({:local, path}), do: {:local, Path.join(path, "encoder_config")} + defp encoder_repository({:hf, id}), do: encoder_repository({:hf, id, []}) + + defp encoder_repository({:hf, id, opts}) do + subdir = if opts[:subdir], do: opts[:subdir] <> "/encoder_config", else: "encoder_config" + {:hf, id, Keyword.put(opts, :subdir, subdir)} + end + + @doc """ + Extracts entities from one document without starting a serving process. + + Options match `Bumblebee.Text.entity_extraction/2`. Offsets count Unicode + codepoints in the original text. Labels are deduplicated in input order. + """ + def extract_entities(model, text, labels, opts \\ []) do + labels = validate_input!(%{text: text, labels: labels}) + opts = validate_options!(opts) + + if labels == [] do + %{} + else + packed = pack(model, text, labels) + Decode.entities(packed, infer(model, packed), model.config, opts) + end + end + + @doc false + def new(model, opts \\ []) do + opts = validate_options!(opts) + + Nx.Serving.new(__MODULE__, model) + |> Nx.Serving.batch_size(1) + |> Nx.Serving.process_options(batch_keys: [:default, :empty]) + |> Nx.Serving.client_preprocessing(fn input -> + labels = validate_input!(input) + + if labels == [] do + {Nx.Batch.concatenate([Nx.tensor([0])]) |> Nx.Batch.key(:empty), :empty} + else + packed = pack(model, input.text, labels) + data = {packed.inputs, packed.text_indices, packed.query_indices} + {Nx.Batch.concatenate([data]), packed} + end + end) + |> Nx.Serving.client_postprocessing(fn + {_outputs, _metadata}, :empty -> %{} + {outputs, _metadata}, packed -> Decode.entities(packed, outputs, model.config, opts) + end) + end + + @impl true + def init(_type, model, _partitions), do: {:ok, model} + + @impl true + def handle_batch(%Nx.Batch{key: :empty}, _partition, model), + do: {:execute, fn -> {Nx.tensor([0]), :inference} end, model} + + def handle_batch(%Nx.Batch{size: 1} = batch, _partition, model) do + execute = fn -> + {inputs, text_indices, query_indices} = + Nx.Defn.jit_apply(&Function.identity/1, [batch], model.defn_options) + + outputs = + infer(model, %{inputs: inputs, text_indices: text_indices, query_indices: query_indices}) + + {Map.take(outputs, ~w(pair_logits null_logits starts ends valid)), :inference} + end + + {:execute, execute, model} + end + + def handle_batch(_batch, _partition, _model), + do: raise(ArgumentError, "GLiNER serving requires batch_size: 1") + + @doc false + def infer(model, packed) do + original_length = Nx.axis_size(packed.inputs["input_ids"], 1) + hidden = predict(model.encoder, pad_inputs(packed.inputs, model.sequence_length)).hidden_state + hidden = Nx.slice_along_axis(hidden, 0, original_length, axis: 1) + + inputs = %{ + "text_states" => route(hidden, packed.text_indices), + "query_states" => route(hidden, packed.query_indices) + } + + output = predict(model.boundary, inputs) + candidates = CandidatePool.select(output, model.config) + inputs = inputs |> Map.merge(output) |> Map.merge(candidates) + + inputs + |> Map.put("pair_logits", predict(model.scorer, inputs)) + |> Map.put("hidden_states", hidden) + end + + defp predict({fun, params}, inputs), do: fun.(params, inputs) + defp route(hidden, indices), do: Nx.take(hidden, Nx.squeeze(indices, axes: [0]), axis: 1) + + defp pack(model, text, labels) do + packed = Preprocessor.pack(text, labels, model.tokenizer) + limit = model.sequence_length || model.max_sequence_length + + if Nx.axis_size(packed.inputs["input_ids"], 1) > limit do + raise ArgumentError, "packed input exceeds the configured limit of #{limit} tokens" + end + + packed + end + + defp pad_inputs(inputs, nil), do: inputs + + defp pad_inputs(inputs, length) do + padding = length - Nx.axis_size(inputs["input_ids"], 1) + if padding < 0, do: raise(ArgumentError, "input exceeds sequence_length") + + Map.new(inputs, fn {key, tensor} -> {key, Nx.pad(tensor, 0, [{0, 0, 0}, {0, padding, 0}])} end) + end + + defp validate_input!(%{text: text, labels: labels}) when is_binary(text) and is_list(labels) do + unless Enum.all?(labels, &(is_binary(&1) and String.trim(&1) != "")), + do: raise(ArgumentError, "labels must be nonempty strings") + + Enum.uniq(labels) + end + + defp validate_input!(_), + do: raise(ArgumentError, "expected %{text: string, labels: list_of_strings}") + + defp validate_options!(opts) do + opts = + Keyword.validate!(opts, threshold: 0.5, include_spans: false, include_confidence: false) + + unless is_number(opts[:threshold]) and opts[:threshold] >= 0 and opts[:threshold] <= 1, + do: raise(ArgumentError, "threshold must be between zero and one") + + for key <- [:include_spans, :include_confidence], + not is_boolean(opts[key]), + do: raise(ArgumentError, "#{key} must be a boolean") + + opts + end +end diff --git a/mix.exs b/mix.exs index 5d84035b..a24ad04b 100644 --- a/mix.exs +++ b/mix.exs @@ -92,6 +92,7 @@ defmodule Bumblebee.MixProject do Bumblebee.Text.Distilbert, Bumblebee.Text.Gemma, Bumblebee.Text.Gemma3Text, + Bumblebee.Text.Gliner, Bumblebee.Text.Gpt2, Bumblebee.Text.GptBigCode, Bumblebee.Text.GptNeoX, diff --git a/test/bumblebee/text/gliner/decode_test.exs b/test/bumblebee/text/gliner/decode_test.exs new file mode 100644 index 00000000..8f53a441 --- /dev/null +++ b/test/bumblebee/text/gliner/decode_test.exs @@ -0,0 +1,85 @@ +defmodule Bumblebee.Text.Gliner.DecodeTest do + use ExUnit.Case, async: true + alias Bumblebee.Text.Gliner.Decode + alias Bumblebee.Text.Gliner.Preprocessor + + test "synthetic punctuation is removed before overlap selection and deduplication" do + result = decode("D", [{1.0, 0, 1}, {2.0, 0, 2}, {3.0, 1, 2}]) + assert [%{"text" => "D", "start" => 0, "end" => 1, "confidence" => score}] = result + assert_in_delta score, 0.880797, 1.0e-6 + end + + test "empty and whitespace-only text never produce synthetic entities" do + assert decode("", [{2.0, 0, 1}]) == [] + assert decode(" \t", [{2.0, 0, 1}]) == [] + end + + test "trailing whitespace is excluded from both surface and offsets" do + assert [%{"text" => "A", "start" => 0, "end" => 1}] = + decode("A \t", [{2.0, 0, 2}]) + end + + test "real terminal punctuation is preserved" do + assert [%{"text" => "D.", "start" => 0, "end" => 2}] = + decode("D.", [{2.0, 0, 2}]) + end + + test "projection retains Unicode codepoint offsets" do + text = "👩🏽‍💻 Café" + count = length(Preprocessor.words(text <> ".")) + + assert [%{"text" => ^text, "start" => 0, "end" => ending}] = + decode(text, [{2.0, 0, count}]) + + assert ending == length(String.codepoints(text)) + end + + test "flat decoding maximizes total score rather than choosing the largest single span" do + assert Decode.resolve_flat([{0.9, 0, 4}, {0.6, 0, 2}, {0.6, 2, 4}]) == [ + {0.6, 0, 2}, + {0.6, 2, 4} + ] + end + + test "ties prefer the larger compatible set then confidence/start/end rank" do + assert Decode.resolve_flat([{1.0, 0, 4}, {0.5, 0, 2}, {0.5, 2, 4}]) == [ + {0.5, 0, 2}, + {0.5, 2, 4} + ] + + assert Decode.resolve_flat([{0.5, 1, 3}, {0.5, 0, 2}, {0.4, 0, 2}]) == [{0.5, 0, 2}] + assert Decode.resolve_flat([]) == [] + end + + test "threshold comparison remains inclusive without an epsilon" do + assert [%{"confidence" => 0.5}] = decode("A", [{0.0, 0, 1}], 0.5) + assert decode("A", [{0.0, 0, 1}], 0.50000001) == [] + assert [%{"confidence" => 0.5}] = decode("A", [{0.0, 0, 1}], 0.49999999) + end + + defp decode(text, spans, threshold \\ 0.5) do + augmented = if String.ends_with?(text, [".", "!", "?"]), do: text, else: text <> "." + + packed = %{ + text: augmented, + original_text: text, + original_length: length(String.codepoints(text)), + words: Preprocessor.words(augmented), + labels: ["person"] + } + + output = %{ + "pair_logits" => Nx.tensor([Enum.map(spans, &elem(&1, 0))]), + "null_logits" => Nx.tensor([[-10.0]]), + "starts" => Nx.tensor([Enum.map(spans, &elem(&1, 1))]), + "ends" => Nx.tensor([Enum.map(spans, &elem(&1, 2))]), + "valid" => Nx.tensor([Enum.map(spans, fn _ -> 1 end)]) + } + + Decode.entities(packed, output, %{"pair_temperature" => 1.0, "abstention_threshold" => 0.5}, + threshold: threshold, + include_spans: true, + include_confidence: true + )["entities"]["person"] + end +end diff --git a/test/bumblebee/text/gliner/preprocessor_test.exs b/test/bumblebee/text/gliner/preprocessor_test.exs new file mode 100644 index 00000000..a26641a5 --- /dev/null +++ b/test/bumblebee/text/gliner/preprocessor_test.exs @@ -0,0 +1,22 @@ +defmodule Bumblebee.Text.Gliner.PreprocessorTest do + use ExUnit.Case, async: true + alias Bumblebee.Text.Gliner.Preprocessor + + test "Python IGNORECASE keeps dotted-I emails and handles intact" do + assert Preprocessor.words("İNFO@example.org @ıpek") == [ + %{text: "i̇nfo@example.org", start: 0, end: 16}, + %{text: "@ıpek", start: 17, end: 22} + ] + end + + test "Python lower uses contextual final sigma and preserves original offsets" do + assert Preprocessor.words("ΟΣ ΟΣΑ ΣΟΣ ΑΣ; İ") == [ + %{text: "ος", start: 0, end: 2}, + %{text: "οσα", start: 3, end: 6}, + %{text: "σος", start: 7, end: 10}, + %{text: "ας", start: 11, end: 13}, + %{text: ";", start: 13, end: 14}, + %{text: "i̇", start: 15, end: 16} + ] + end +end diff --git a/test/bumblebee/text/gliner_entity_extraction_test.exs b/test/bumblebee/text/gliner_entity_extraction_test.exs new file mode 100644 index 00000000..c74a16ab --- /dev/null +++ b/test/bumblebee/text/gliner_entity_extraction_test.exs @@ -0,0 +1,103 @@ +defmodule Bumblebee.Text.GlinerEntityExtractionTest do + use ExUnit.Case, async: false + alias Bumblebee.Text.GlinerEntityExtraction, as: Extraction + @moduletag Bumblebee.TestHelpers.serving_test_tags() + + setup_all do + reference = + Path.expand("../../fixtures/gliner/ner_reference.json", __DIR__) + |> File.read!() + |> Jason.decode!() + + repository = + case System.get_env("GLINER_MODEL_DIR") do + nil -> {:hf, reference["model"], revision: reference["revision"]} + path -> {:local, path} + end + + {:ok, model} = + Bumblebee.Text.load_entity_extraction(repository, defn_options: [compiler: EXLA]) + + %{model: model, examples: reference["examples"]} + end + + test "entity decisions and confidence match Python on representative inputs", %{ + model: model, + examples: examples + } do + for example <- examples do + actual = + Extraction.extract_entities(model, example["text"], example["labels"], + include_spans: true, + include_confidence: true + ) + + assert Map.keys(actual["entities"]) |> Enum.sort() == + Map.keys(example["prediction"]["entities"]) |> Enum.sort() + + for {label, expected} <- example["prediction"]["entities"] do + assert length(actual["entities"][label]) == length(expected) + + for {actual, expected} <- Enum.zip(actual["entities"][label], expected) do + assert Map.delete(actual, "confidence") == Map.delete(expected, "confidence") + assert_in_delta actual["confidence"], expected["confidence"], 1.0e-5 + end + end + end + end + + test "serves concurrent documents, deduplicates labels and handles empty labels", %{ + model: model + } do + serving = Bumblebee.Text.entity_extraction(model, include_spans: true) + start_supervised!({Nx.Serving, serving: serving, name: __MODULE__}) + + input = %{ + text: "Apple CEO Tim Cook announced the iPhone 15 in Cupertino.", + labels: ["company", "person", "company"] + } + + expected = %{ + "entities" => %{ + "company" => [%{"text" => "Apple", "start" => 0, "end" => 5}], + "person" => [%{"text" => "Tim Cook", "start" => 10, "end" => 18}] + } + } + + for result <- + Task.async_stream(1..4, fn _ -> Nx.Serving.batched_run(__MODULE__, input) end, + timeout: 120_000 + ) do + assert {:ok, ^expected} = result + end + + assert Nx.Serving.batched_run(__MODULE__, %{text: "Ignored", labels: []}) == %{} + + assert Extraction.extract_entities(model, "Ignored", []) == %{} + end + + test "validates inputs and limits before model execution", %{model: model} do + assert_raise ArgumentError, ~r/nonempty strings/, fn -> + Extraction.extract_entities(model, "Hello", [""]) + end + + assert_raise ArgumentError, ~r/threshold/, fn -> + Extraction.extract_entities(model, "Hello", ["person"], threshold: 1.1) + end + + assert_raise ArgumentError, ~r/exceeds/, fn -> + Extraction.extract_entities(%{model | max_sequence_length: 1}, "Hello", ["person"]) + end + end + + test "fixed encoder padding preserves decisions", %{model: model} do + model = %{model | sequence_length: 128} + + assert Extraction.extract_entities( + model, + "Apple CEO Tim Cook announced the iPhone 15 in Cupertino.", + ["company", "person"] + ) == + %{"entities" => %{"company" => ["Apple"], "person" => ["Tim Cook"]}} + end +end diff --git a/test/bumblebee/text/gliner_test.exs b/test/bumblebee/text/gliner_test.exs new file mode 100644 index 00000000..921cb35b --- /dev/null +++ b/test/bumblebee/text/gliner_test.exs @@ -0,0 +1,96 @@ +defmodule Bumblebee.Text.GlinerTest do + use ExUnit.Case, async: true + import Nx.Testing + import ExUnit.CaptureLog + + for variant <- ["local", "global", "no_refinement"] do + @variant variant + test "boundary and span scores match Python for #{@variant}" do + directory = Path.expand("../../fixtures/gliner/#{@variant}", __DIR__) + reference = directory |> Path.join("expected.json") |> File.read!() |> Jason.decode!() + + log = + capture_log(fn -> + {:ok, info} = + Bumblebee.load_model({:local, directory}, + spec_overrides: [hidden_size: 12], + log_params_diff: true + ) + + assert info.spec.architecture == :boundary + {_, predict} = Axon.build(info.model, compiler: EXLA) + inputs = Map.new(reference["inputs"], fn {key, value} -> {key, Nx.tensor(value)} end) + output = predict.(info.params, inputs) + + for {key, expected} <- reference["outputs"] do + assert_all_close(output[key], Nx.tensor(expected), atol: 2.0e-5, rtol: 2.0e-5) + end + + {:ok, info} = + Bumblebee.load_model({:local, directory}, + architecture: :span_scoring, + spec_overrides: [hidden_size: 12], + log_params_diff: true + ) + + {_, predict} = Axon.build(info.model, compiler: EXLA) + + inputs = + Map.new(reference["scorer_inputs"], fn {key, value} -> {key, Nx.tensor(value)} end) + + assert_all_close(predict.(info.params, inputs), Nx.tensor(reference["scores"]), + atol: 2.0e-5, + rtol: 2.0e-5 + ) + end) + + refute log =~ "were missing" + refute log =~ "non-matching shape" + end + end + + test "loads a nested DeBERTa encoder from GLiNER parameter names" do + directory = Path.expand("../../fixtures/gliner/local", __DIR__) + {:ok, encoder_spec} = Bumblebee.load_spec({:local, Path.join(directory, "encoder_config")}) + + {:ok, info} = + Bumblebee.load_model({:local, directory}, + architecture: :encoder, + spec_overrides: [encoder_spec: encoder_spec] + ) + + reference = directory |> Path.join("encoder_expected.json") |> File.read!() |> Jason.decode!() + inputs = Map.new(reference["inputs"], fn {key, value} -> {key, Nx.tensor(value)} end) + {_, predict} = Axon.build(info.model, compiler: EXLA) + + assert_all_close( + predict.(info.params, inputs).hidden_state, + Nx.tensor(reference["hidden_state"]), + atol: 2.0e-5, + rtol: 2.0e-5 + ) + end + + test "rejects invalid scoring and refinement options" do + for options <- [ + [pair_temperature: 0], + [boundary_ffn_multiplier: -1], + [abstention_threshold: 2] + ] do + assert_raise ArgumentError, fn -> Bumblebee.configure(Bumblebee.Text.Gliner, options) end + end + end + + test "rejects unsupported variants before loading weights" do + for options <- [ + [candidate_pool: "per_query"], + [query_attention_layers: 1], + [adaptive_threshold: true], + [content_soft_max_pool: true] + ] do + assert_raise ArgumentError, ~r/unsupported GLiNER option/, fn -> + Bumblebee.configure(Bumblebee.Text.Gliner, options) + end + end + end +end diff --git a/test/fixtures/gliner/README.md b/test/fixtures/gliner/README.md new file mode 100644 index 00000000..9157f019 --- /dev/null +++ b/test/fixtures/gliner/README.md @@ -0,0 +1,43 @@ +# GLiNER reference fixtures + +`generate.py` creates three small, seeded boundary-head checkpoints and their +Python outputs. The cases cover local and global boundary attention, disabled +attention/refinement blocks, shared-pool span scoring, and an invalid padded +candidate. The `local` checkpoint also embeds the small DeBERTa shared-attention +encoder to test loading parameters with the GLiNER prefix. Its configuration +and Python output are stored alongside the boundary-head fixture. + +The generator creates both the encoder and boundary-head fixtures. It was +run in a Python 3.12 virtual environment with torch 2.14.0, transformers 4.57.6, +gliner2 2.0.0, safetensors 0.8.0, and tokenizers 0.22.2. The generated files are +checked in, so ordinary tests do not require Python or network access. + +```sh +python test/fixtures/gliner/generate.py +mix test test/bumblebee/text/gliner_test.exs +``` + +`ner_reference.json` contains CPU/FP32 predictions from the official +`fastino/gliner2.5-base-v1` checkpoint at revision +`78cea040597df251eedefa9d7ee2a756af39fe64`, using gliner2 2.0.0. The examples +cover ordinary entities, no entities, Unicode, overlapping mentions, +punctuation, and a long document. They were recorded with: + +```python +from gliner2 import AutoExtractor +model = AutoExtractor.from_pretrained(checkpoint_directory).float().eval() +prediction = model.extract_entities(text, labels, include_spans=True, include_confidence=True) +``` + +The serving tests load that pinned checkpoint. To use an existing local copy: + +```sh +GLINER_MODEL_DIR=/path/to/checkpoint mix test \ + test/bumblebee/text/gliner_entity_extraction_test.exs --include slow +``` + +The extraction contract retains inclusive confidence cutoffs and Unicode +codepoint offsets. Spans are clipped to the original document before overlap +resolution, so the period appended during preprocessing cannot become an +entity or extend one past the supplied text. Exact floating-point equality +across platforms and decisions at arbitrarily close cutoffs are not promised. diff --git a/test/fixtures/gliner/generate.py b/test/fixtures/gliner/generate.py new file mode 100644 index 00000000..f881a9ba --- /dev/null +++ b/test/fixtures/gliner/generate.py @@ -0,0 +1,72 @@ +"""Generate offline head references with gliner2 2.0.0 / torch 2.14.0.""" +import dataclasses +import json +from pathlib import Path +import torch +from safetensors.torch import save_file +from transformers import DebertaV2Config, DebertaV2Model +from gliner2.configuration import BoundaryHeadSettings +from gliner2.models.boundary.model import BoundaryHead +from gliner2.models.boundary.pool import PooledCandidates + +root = Path(__file__).resolve().parent +torch.set_num_threads(1) +torch.manual_seed(42) +encoder = DebertaV2Model(DebertaV2Config( + vocab_size=32, hidden_size=12, num_hidden_layers=2, num_attention_heads=3, + intermediate_size=17, max_position_embeddings=32, relative_attention=True, + pos_att_type=["c2p", "p2c"], share_att_key=True, position_buckets=8, + norm_rel_ebd="layer_norm", position_biased_input=False, +)).eval() +encoder.config.architectures = ["DebertaV2Model"] +encoder_inputs = { + "input_ids": torch.tensor([[2,3,4,5,6,7,8,9,10,11,12,13], [9,8,7,6,5,4,3,2,0,0,0,0]]), + "attention_mask": torch.tensor([[1]*12, [1]*8+[0]*4]), +} +with torch.inference_mode(): + encoder_output = encoder(**encoder_inputs).last_hidden_state +for name, layers, window in [("local", 1, 2), ("global", 1, 0), ("no_refinement", 0, 0)]: + torch.manual_seed(42) + settings = BoundaryHeadSettings(boundary_dim=8, pair_dim=8, content_dim=4, + boundary_attention_heads=2, boundary_attention_layers=layers, + boundary_refinement_layers=layers, boundary_attention_window=window, + candidate_pool="shared", candidate_attention_layers=0, query_attention_layers=0, + content_soft_max_pool=False, enable_span_content=True, enable_abstention=True, + use_inside_evidence=True, adaptive_threshold=False, overlap_policy="flat", + pool_size=8, pool_boundary_top_k=4, min_pool_per_query=2) + head = BoundaryHead(12, settings).eval() + text = torch.randn(1, 7, 12) + query = torch.randn(1, 3, 12) + text_mask = torch.ones(1,7,dtype=torch.bool) + query_mask = torch.ones(1,3,dtype=torch.bool) + with torch.inference_mode(): + boundary = head.boundary_encoder(text, text_mask) + margins = head.boundary_query_head(boundary.states, boundary.mask, text, text_mask, query, query_mask) + pooled = head.shared_pool_builder(boundary.states, boundary.mask, query_mask, margins.start_logits, margins.end_logits) + # Retain a padded candidate to check score masking. + pooled.mask[:, -1] = False + scores, _ = head.shared_pool_scorer(boundary.states, query, query_mask, pooled, + margins.start_logits, margins.end_logits, margins.inside_prefix, + torch.tensor([7]), text, text_mask, margins.inside_prefix_mean) + output = {"boundary_states": boundary.states, "start_logits": margins.start_logits, + "end_logits": margins.end_logits, "inside_logits": margins.inside_logits, + "null_logits": head.null_projection(query), + "pool_start": head.shared_pool_builder.start_projection(boundary.states), + "pool_end": head.shared_pool_builder.end_projection(boundary.states)} + inputs = {"text_states": text, "query_states": query} + scorer_inputs = {**inputs, **{k:output[k] for k in ["boundary_states","start_logits","end_logits","inside_logits"]}, + "starts": pooled.indices[...,0], "ends": pooled.indices[...,1], + "valid": pooled.mask.to(torch.uint8), "compat": pooled.compat_logits} + directory = root / name + directory.mkdir(exist_ok=True) + state = {"boundary_head."+key: value for key,value in head.state_dict().items()} + if name == "local": + state.update({"encoder."+key:value for key,value in encoder.state_dict().items()}) + encoder.config.save_pretrained(directory / "encoder_config") + (directory / "encoder_expected.json").write_text(json.dumps({ + "inputs": {key: value.tolist() for key, value in encoder_inputs.items()}, + "hidden_state": encoder_output.tolist(), + }, separators=(",", ":")) + "\n") + save_file(state, directory / "model.safetensors") + (directory / "config.json").write_text(json.dumps({"architectures":["BoundaryExtractor"], "architecture":"boundary", "token_pooling":"first", "boundary_head":dataclasses.asdict(settings)}, indent=2)+"\n") + (directory / "expected.json").write_text(json.dumps({"inputs":{k:v.tolist() for k,v in inputs.items()},"outputs":{k:v.tolist() for k,v in output.items()},"scorer_inputs":{k:v.tolist() for k,v in scorer_inputs.items()},"scores":scores.transpose(1,2).tolist()},separators=(",",":"))+"\n") diff --git a/test/fixtures/gliner/global/config.json b/test/fixtures/gliner/global/config.json new file mode 100644 index 00000000..48a39270 --- /dev/null +++ b/test/fixtures/gliner/global/config.json @@ -0,0 +1,92 @@ +{ + "architectures": [ + "BoundaryExtractor" + ], + "architecture": "boundary", + "token_pooling": "first", + "boundary_head": { + "boundary_dim": 8, + "pair_dim": 8, + "boundary_refinement_layers": 1, + "boundary_ffn_multiplier": 2.0, + "start_top_k": 16, + "end_top_k": 16, + "ends_per_start": 8, + "starts_per_end": 8, + "candidate_budget": 128, + "training_candidate_budget": 160, + "max_gold_per_query": 32, + "end_block_size": 256, + "bidirectional_proposals": true, + "use_inside_evidence": true, + "dropout": 0.1, + "export_mode": "auto", + "vectorized_pair_elements": 16777216, + "boundary_negative_weight": 1.0, + "boundary_marginal_loss": "bce", + "loss_reduction": "global", + "boundary_focal_gamma_positive": 0.0, + "boundary_focal_gamma_negative": 2.0, + "boundary_focal_clip": 0.05, + "hard_negatives_per_positive": 5, + "minimum_hard_negatives": 8, + "hard_negative_keep_all_when_absent": false, + "enable_span_content": true, + "content_dim": 4, + "content_soft_max_pool": false, + "enable_rotary_endpoints": false, + "rotary_base": 10000.0, + "boundary_attention_layers": 1, + "boundary_attention_heads": 2, + "boundary_attention_window": 0, + "query_conditioned_inside_weight": false, + "endpoint_difference_features": false, + "reranker_endpoint_compat": true, + "multihead_pair_compat_heads": 8, + "boundary_top_k_alpha": 0.0, + "boundary_top_k_max": 128, + "boundary_top_k_bucket": 8, + "candidate_pool": "shared", + "pool_boundary_top_k": 4, + "pool_size": 8, + "min_pool_per_query": 2, + "candidate_attention_layers": 0, + "candidate_attention_heads": 4, + "query_attention_layers": 0, + "enable_abstention": true, + "abstention_threshold": 0.5, + "proposal_loss_weight": 0.3, + "consistency_loss_weight": 0.1, + "rerank_listwise_weight": 0.3, + "soft_iou_aux_weight": 0.2, + "soft_iou_anneal_steps": 20000, + "abstention_loss_weight": 0.2, + "consistency_warmup_steps": 2000, + "enable_count_head": true, + "count_loss_weight": 0.2, + "adaptive_threshold": false, + "overlap_policy": "flat", + "pair_temperature": 1.0, + "relation_temperature": 1.0, + "record_temperature": 1.0, + "classification_temperature": 1.0, + "negative_query_ratio": 0.5, + "max_negative_queries_per_batch": 64, + "classification_loss_weight": 1.0, + "enable_records": true, + "record_dim": 128, + "record_instance_queries": 32, + "record_anchor_proposal_threshold": 0.2, + "record_anchor_threshold": 0.5, + "record_field_threshold": 0.5, + "record_loss_weight": 1.0, + "enable_relations": true, + "relation_heads_per_type": 32, + "relation_tails_per_type": 32, + "relation_pair_cap": 128, + "relation_loss_weight": 1.0, + "relation_argument_proposal_threshold": 0.0, + "directional_relation_states": false, + "relation_biaffine_content": false + } +} diff --git a/test/fixtures/gliner/global/expected.json b/test/fixtures/gliner/global/expected.json new file mode 100644 index 00000000..676610e4 --- /dev/null +++ b/test/fixtures/gliner/global/expected.json @@ -0,0 +1 @@ +{"inputs":{"text_states":[[[0.3396981954574585,1.273138165473938,-1.5371414422988892,1.0402002334594727,-1.1433438062667847,-0.6347954869270325,1.180337905883789,0.37527406215667725,0.8895691633224487,2.015190601348877,0.12288182973861694,0.283156156539917],[0.7347974181175232,1.4138669967651367,-0.06916153430938721,-0.2501976788043976,1.7436981201171875,-0.38771310448646545,0.7723278403282166,0.22671592235565186,-0.38377487659454346,-0.7681311964988708,-1.003414511680603,-0.9655314683914185],[0.2913444936275482,-1.1624996662139893,-0.4783642590045929,1.3870203495025635,-0.4341132640838623,0.2929425835609436,0.11547855287790298,-0.17881472408771515,-0.7945829629898071,-0.5694751739501953,0.0549420528113842,-0.7646364569664001],[-1.049290657043457,1.0828652381896973,-0.9056679606437683,-1.0400710105895996,2.1871161460876465,1.1617707014083862,0.6346917748451233,-0.4325243830680847,-1.275227665901184,-0.9807081818580627,0.7780191898345947,-0.23794014751911163],[1.774613618850708,1.5161902904510498,-1.0465848445892334,1.3759161233901978,0.05762922391295433,0.35209375619888306,1.0967378616333008,-0.9882838129997253,-0.5617773532867432,-0.4788392186164856,-0.7523331642150879,-0.4768691956996918],[1.2204824686050415,-0.15345095098018646,-0.5119454264640808,0.11924171447753906,0.023615580052137375,-0.7165514230728149,0.9313072562217712,0.11275327950716019,0.5183308720588684,0.10639318078756332,-0.0708099976181984,0.1868356168270111],[-1.135615348815918,-0.6121562123298645,0.47324639558792114,0.22348645329475403,0.4591689705848694,0.3639809787273407,0.24808230996131897,-0.5714194178581238,-1.8829892873764038,1.6257703304290771,-0.7037963271141052,0.558134913444519]]],"query_states":[[[0.42072793841362,1.2927827835083008,-0.062042392790317535,-0.07447009533643723,1.8144526481628418,0.4126001298427582,-0.5384427905082703,-0.6699069142341614,0.13211850821971893,-0.45546427369117737,0.7100200653076172,-1.4318643808364868],[1.3607994318008423,0.6645160913467407,-0.19414693117141724,-1.3718643188476562,1.7688463926315308,1.2299517393112183,0.29568976163864136,0.8400743007659912,0.3205632269382477,-0.5852316617965698,0.4060876965522766,0.2572997212409973],[0.8778835535049438,0.7318398952484131,0.27381399273872375,1.7188366651535034,-1.4703203439712524,-0.3415771424770355,1.2259759902954102,1.8914308547973633,-0.6524525284767151,-0.8879274725914001,-1.0803742408752441,1.4245182275772095]]]},"outputs":{"boundary_states":[[[0.435674786567688,-0.3552878201007843,1.0446969270706177,-0.23349082469940186,1.0340025424957275,1.100759744644165,-2.3040928840637207,-0.6984513401985168],[-0.2731614410877228,0.32742464542388916,1.1776659488677979,1.5370218753814697,-0.5072778463363647,0.9318945407867432,-2.1648144721984863,-0.8701298236846924],[0.18581756949424744,2.110225200653076,-0.3917454183101654,-1.081823468208313,-0.19402827322483063,-0.059499531984329224,0.1265341341495514,-1.1011775732040405],[-1.378901481628418,-1.346583604812622,0.8190063834190369,0.9230098724365234,1.1955806016921997,1.1885589361190796,-0.4205845296382904,-0.736480176448822],[-0.015728741884231567,-0.1459733247756958,-0.5577592849731445,-1.0362745523452759,2.0171844959259033,-0.07644002139568329,-0.32235294580459595,-0.5144805312156677],[-0.42267632484436035,0.21576334536075592,0.9981375932693481,0.6242271065711975,0.5597060322761536,1.5171633958816528,-1.327799677848816,-2.2183334827423096],[-0.9369467496871948,-0.9813311100006104,0.9250056147575378,0.8637326955795288,1.349330186843872,1.1422157287597656,-1.6476943492889404,-0.7070372104644775],[-0.0252586230635643,1.9137977361679077,0.32240030169487,-1.8818881511688232,-0.098744235932827,-0.11144017428159714,-0.28226974606513977,0.4701942205429077]]],"start_logits":[[[-0.4956240952014923,0.6737099289894104,0.5640376210212708,-0.2739996016025543,-0.7868633270263672,0.4668186604976654,-0.34135890007019043,-0.14971965551376343],[0.5836982727050781,1.8898866176605225,0.4864198565483093,0.24729780852794647,-0.9303515553474426,1.4244972467422485,0.4904747009277344,-0.14380036294460297],[-0.20199060440063477,-0.33043450117111206,0.17490944266319275,-0.382266640663147,0.04883630573749542,-0.25426435470581055,-0.3888903558254242,0.07494132220745087]]],"end_logits":[[[-0.24732623994350433,-0.1219557374715805,-0.037788767367601395,-0.22911766171455383,-0.08409157395362854,-0.2427012324333191,-0.27881643176078796,-0.2082749903202057],[-0.23388980329036713,-0.23882760107517242,-0.21627427637577057,0.03876682370901108,0.13328732550144196,-0.26707398891448975,-0.025014901533722878,-0.25488826632499695],[0.1483789086341858,0.2214311957359314,-0.2682836353778839,0.3091442883014679,-0.07662579417228699,0.25217729806900024,0.3251894414424896,-0.34030088782310486]]],"inside_logits":[[[0.14138878881931305,0.5262569785118103,-0.23322194814682007,0.47248613834381104,0.21367579698562622,0.09782277047634125,0.2908124625682831],[0.19397881627082825,0.32151395082473755,0.07189678400754929,0.37444400787353516,0.4646332561969757,0.04348849877715111,0.4586145579814911],[0.13514390587806702,-0.6980278491973877,0.32489508390426636,-0.5097112655639648,-0.2958136796951294,0.03978963941335678,0.1785593181848526]]],"null_logits":[[[0.0],[0.0],[0.0]]],"pool_start":[[[0.1424010992050171,-1.2664109468460083,-0.3496744632720947,0.2284083515405655,-0.09108977019786835,0.4088525176048279,-0.3723243474960327,-0.6073645353317261],[0.2809818983078003,-1.7302943468093872,0.06907674670219421,0.10990826785564423,-0.8067101240158081,0.28423887491226196,-0.19281499087810516,0.2266453504562378],[0.6404389142990112,-0.5601590275764465,0.07159161567687988,0.2967291474342346,0.0737350583076477,-0.2878093421459198,0.18593795597553253,0.3432932198047638],[-0.034147948026657104,-0.3844667673110962,-0.7951714992523193,0.3923138380050659,-0.7382920980453491,0.7803592681884766,-0.8190670609474182,-0.5816023945808411],[-0.0793866515159607,-0.11469712853431702,-0.545879602432251,0.21922560036182404,0.2860366106033325,0.32230955362319946,-0.2154349833726883,-0.9906845688819885],[0.8201143741607666,-1.6151561737060547,-0.3219591975212097,0.30914050340652466,-0.38404205441474915,0.4628337323665619,-0.3937923312187195,0.1008448526263237],[0.0484461784362793,-0.9714303016662598,-0.6605094075202942,0.37103384733200073,-0.5901164412498474,0.7520049810409546,-0.6656572818756104,-0.7515038251876831],[0.28223752975463867,0.0958474650979042,-0.24649132788181305,0.752622663974762,-0.1236373782157898,-0.34583672881126404,-0.05648399889469147,-0.16845795512199402]]],"pool_end":[[[-0.4914288818836212,0.026747405529022217,-0.3121238052845001,0.8096144795417786,0.14824475347995758,0.6502490043640137,0.3671307861804962,0.6362134218215942],[-0.7871421575546265,0.08015817403793335,-0.42331045866012573,0.6645656228065491,0.4108622074127197,0.44630831480026245,0.008258026093244553,0.5134128332138062],[0.5806735754013062,-0.2012903392314911,0.6403052806854248,0.03131312131881714,-0.5458632707595825,-0.7429167628288269,-0.12218023836612701,0.9721102714538574],[0.4121811091899872,-0.6742359399795532,-1.6212353706359863,0.015002161264419556,-0.36485618352890015,1.5678343772888184,-0.47280293703079224,0.4175695478916168],[0.5988327264785767,-0.17072325944900513,-0.2586260139942169,-0.22670775651931763,-0.7181596755981445,0.43657124042510986,-0.3557000756263733,0.32711806893348694],[0.278384804725647,-0.2086164951324463,-1.068066954612732,0.2583248019218445,-0.21872283518314362,0.9032034873962402,-0.40400850772857666,1.1619079113006592],[-0.12734481692314148,-0.328058660030365,-1.2689820528030396,0.24951604008674622,-0.18137241899967194,1.4049160480499268,-0.24997873604297638,0.3420554995536804],[-0.011179909110069275,-0.6165403127670288,1.0077025890350342,0.5894898176193237,-0.5688668489456177,-0.7616395354270935,0.6398172974586487,0.7351795434951782]]]},"scorer_inputs":{"text_states":[[[0.3396981954574585,1.273138165473938,-1.5371414422988892,1.0402002334594727,-1.1433438062667847,-0.6347954869270325,1.180337905883789,0.37527406215667725,0.8895691633224487,2.015190601348877,0.12288182973861694,0.283156156539917],[0.7347974181175232,1.4138669967651367,-0.06916153430938721,-0.2501976788043976,1.7436981201171875,-0.38771310448646545,0.7723278403282166,0.22671592235565186,-0.38377487659454346,-0.7681311964988708,-1.003414511680603,-0.9655314683914185],[0.2913444936275482,-1.1624996662139893,-0.4783642590045929,1.3870203495025635,-0.4341132640838623,0.2929425835609436,0.11547855287790298,-0.17881472408771515,-0.7945829629898071,-0.5694751739501953,0.0549420528113842,-0.7646364569664001],[-1.049290657043457,1.0828652381896973,-0.9056679606437683,-1.0400710105895996,2.1871161460876465,1.1617707014083862,0.6346917748451233,-0.4325243830680847,-1.275227665901184,-0.9807081818580627,0.7780191898345947,-0.23794014751911163],[1.774613618850708,1.5161902904510498,-1.0465848445892334,1.3759161233901978,0.05762922391295433,0.35209375619888306,1.0967378616333008,-0.9882838129997253,-0.5617773532867432,-0.4788392186164856,-0.7523331642150879,-0.4768691956996918],[1.2204824686050415,-0.15345095098018646,-0.5119454264640808,0.11924171447753906,0.023615580052137375,-0.7165514230728149,0.9313072562217712,0.11275327950716019,0.5183308720588684,0.10639318078756332,-0.0708099976181984,0.1868356168270111],[-1.135615348815918,-0.6121562123298645,0.47324639558792114,0.22348645329475403,0.4591689705848694,0.3639809787273407,0.24808230996131897,-0.5714194178581238,-1.8829892873764038,1.6257703304290771,-0.7037963271141052,0.558134913444519]]],"query_states":[[[0.42072793841362,1.2927827835083008,-0.062042392790317535,-0.07447009533643723,1.8144526481628418,0.4126001298427582,-0.5384427905082703,-0.6699069142341614,0.13211850821971893,-0.45546427369117737,0.7100200653076172,-1.4318643808364868],[1.3607994318008423,0.6645160913467407,-0.19414693117141724,-1.3718643188476562,1.7688463926315308,1.2299517393112183,0.29568976163864136,0.8400743007659912,0.3205632269382477,-0.5852316617965698,0.4060876965522766,0.2572997212409973],[0.8778835535049438,0.7318398952484131,0.27381399273872375,1.7188366651535034,-1.4703203439712524,-0.3415771424770355,1.2259759902954102,1.8914308547973633,-0.6524525284767151,-0.8879274725914001,-1.0803742408752441,1.4245182275772095]]],"boundary_states":[[[0.435674786567688,-0.3552878201007843,1.0446969270706177,-0.23349082469940186,1.0340025424957275,1.100759744644165,-2.3040928840637207,-0.6984513401985168],[-0.2731614410877228,0.32742464542388916,1.1776659488677979,1.5370218753814697,-0.5072778463363647,0.9318945407867432,-2.1648144721984863,-0.8701298236846924],[0.18581756949424744,2.110225200653076,-0.3917454183101654,-1.081823468208313,-0.19402827322483063,-0.059499531984329224,0.1265341341495514,-1.1011775732040405],[-1.378901481628418,-1.346583604812622,0.8190063834190369,0.9230098724365234,1.1955806016921997,1.1885589361190796,-0.4205845296382904,-0.736480176448822],[-0.015728741884231567,-0.1459733247756958,-0.5577592849731445,-1.0362745523452759,2.0171844959259033,-0.07644002139568329,-0.32235294580459595,-0.5144805312156677],[-0.42267632484436035,0.21576334536075592,0.9981375932693481,0.6242271065711975,0.5597060322761536,1.5171633958816528,-1.327799677848816,-2.2183334827423096],[-0.9369467496871948,-0.9813311100006104,0.9250056147575378,0.8637326955795288,1.349330186843872,1.1422157287597656,-1.6476943492889404,-0.7070372104644775],[-0.0252586230635643,1.9137977361679077,0.32240030169487,-1.8818881511688232,-0.098744235932827,-0.11144017428159714,-0.28226974606513977,0.4701942205429077]]],"start_logits":[[[-0.4956240952014923,0.6737099289894104,0.5640376210212708,-0.2739996016025543,-0.7868633270263672,0.4668186604976654,-0.34135890007019043,-0.14971965551376343],[0.5836982727050781,1.8898866176605225,0.4864198565483093,0.24729780852794647,-0.9303515553474426,1.4244972467422485,0.4904747009277344,-0.14380036294460297],[-0.20199060440063477,-0.33043450117111206,0.17490944266319275,-0.382266640663147,0.04883630573749542,-0.25426435470581055,-0.3888903558254242,0.07494132220745087]]],"end_logits":[[[-0.24732623994350433,-0.1219557374715805,-0.037788767367601395,-0.22911766171455383,-0.08409157395362854,-0.2427012324333191,-0.27881643176078796,-0.2082749903202057],[-0.23388980329036713,-0.23882760107517242,-0.21627427637577057,0.03876682370901108,0.13328732550144196,-0.26707398891448975,-0.025014901533722878,-0.25488826632499695],[0.1483789086341858,0.2214311957359314,-0.2682836353778839,0.3091442883014679,-0.07662579417228699,0.25217729806900024,0.3251894414424896,-0.34030088782310486]]],"inside_logits":[[[0.14138878881931305,0.5262569785118103,-0.23322194814682007,0.47248613834381104,0.21367579698562622,0.09782277047634125,0.2908124625682831],[0.19397881627082825,0.32151395082473755,0.07189678400754929,0.37444400787353516,0.4646332561969757,0.04348849877715111,0.4586145579814911],[0.13514390587806702,-0.6980278491973877,0.32489508390426636,-0.5097112655639648,-0.2958136796951294,0.03978963941335678,0.1785593181848526]]],"starts":[[0,1,1,1,5,0,0,2]],"ends":[[3,3,5,6,6,6,5,5]],"valid":[[1,1,1,1,1,1,1,0]],"compat":[[0.7352330088615417,0.7417117953300476,0.41302597522735596,0.4041096866130829,0.6236526966094971,0.4858851432800293,0.20160916447639465,0.12127283960580826]]},"scores":[[[-0.40379753708839417,0.7640310525894165,0.8662115335464478,0.877773106098175,-0.05293281376361847,-0.32532328367233276,-0.34477680921554565,-10000.0],[1.2030330896377563,2.5273733139038086,2.4978184700012207,2.7077460289001465,1.493767499923706,1.3854429721832275,1.1557984352111816,-10000.0],[0.6535661220550537,0.10388350486755371,-0.19243544340133667,-0.09690341353416443,1.0134506225585938,0.40546756982803345,0.24792075157165527,-10000.0]]]} diff --git a/test/fixtures/gliner/global/model.safetensors b/test/fixtures/gliner/global/model.safetensors new file mode 100644 index 00000000..e707622c Binary files /dev/null and b/test/fixtures/gliner/global/model.safetensors differ diff --git a/test/fixtures/gliner/local/config.json b/test/fixtures/gliner/local/config.json new file mode 100644 index 00000000..4be44799 --- /dev/null +++ b/test/fixtures/gliner/local/config.json @@ -0,0 +1,92 @@ +{ + "architectures": [ + "BoundaryExtractor" + ], + "architecture": "boundary", + "token_pooling": "first", + "boundary_head": { + "boundary_dim": 8, + "pair_dim": 8, + "boundary_refinement_layers": 1, + "boundary_ffn_multiplier": 2.0, + "start_top_k": 16, + "end_top_k": 16, + "ends_per_start": 8, + "starts_per_end": 8, + "candidate_budget": 128, + "training_candidate_budget": 160, + "max_gold_per_query": 32, + "end_block_size": 256, + "bidirectional_proposals": true, + "use_inside_evidence": true, + "dropout": 0.1, + "export_mode": "auto", + "vectorized_pair_elements": 16777216, + "boundary_negative_weight": 1.0, + "boundary_marginal_loss": "bce", + "loss_reduction": "global", + "boundary_focal_gamma_positive": 0.0, + "boundary_focal_gamma_negative": 2.0, + "boundary_focal_clip": 0.05, + "hard_negatives_per_positive": 5, + "minimum_hard_negatives": 8, + "hard_negative_keep_all_when_absent": false, + "enable_span_content": true, + "content_dim": 4, + "content_soft_max_pool": false, + "enable_rotary_endpoints": false, + "rotary_base": 10000.0, + "boundary_attention_layers": 1, + "boundary_attention_heads": 2, + "boundary_attention_window": 2, + "query_conditioned_inside_weight": false, + "endpoint_difference_features": false, + "reranker_endpoint_compat": true, + "multihead_pair_compat_heads": 8, + "boundary_top_k_alpha": 0.0, + "boundary_top_k_max": 128, + "boundary_top_k_bucket": 8, + "candidate_pool": "shared", + "pool_boundary_top_k": 4, + "pool_size": 8, + "min_pool_per_query": 2, + "candidate_attention_layers": 0, + "candidate_attention_heads": 4, + "query_attention_layers": 0, + "enable_abstention": true, + "abstention_threshold": 0.5, + "proposal_loss_weight": 0.3, + "consistency_loss_weight": 0.1, + "rerank_listwise_weight": 0.3, + "soft_iou_aux_weight": 0.2, + "soft_iou_anneal_steps": 20000, + "abstention_loss_weight": 0.2, + "consistency_warmup_steps": 2000, + "enable_count_head": true, + "count_loss_weight": 0.2, + "adaptive_threshold": false, + "overlap_policy": "flat", + "pair_temperature": 1.0, + "relation_temperature": 1.0, + "record_temperature": 1.0, + "classification_temperature": 1.0, + "negative_query_ratio": 0.5, + "max_negative_queries_per_batch": 64, + "classification_loss_weight": 1.0, + "enable_records": true, + "record_dim": 128, + "record_instance_queries": 32, + "record_anchor_proposal_threshold": 0.2, + "record_anchor_threshold": 0.5, + "record_field_threshold": 0.5, + "record_loss_weight": 1.0, + "enable_relations": true, + "relation_heads_per_type": 32, + "relation_tails_per_type": 32, + "relation_pair_cap": 128, + "relation_loss_weight": 1.0, + "relation_argument_proposal_threshold": 0.0, + "directional_relation_states": false, + "relation_biaffine_content": false + } +} diff --git a/test/fixtures/gliner/local/encoder_config/config.json b/test/fixtures/gliner/local/encoder_config/config.json new file mode 100644 index 00000000..78f1d858 --- /dev/null +++ b/test/fixtures/gliner/local/encoder_config/config.json @@ -0,0 +1,34 @@ +{ + "architectures": [ + "DebertaV2Model" + ], + "attention_probs_dropout_prob": 0.1, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.1, + "hidden_size": 12, + "initializer_range": 0.02, + "intermediate_size": 17, + "layer_norm_eps": 1e-07, + "legacy": true, + "max_position_embeddings": 32, + "max_relative_positions": -1, + "model_type": "deberta-v2", + "norm_rel_ebd": "layer_norm", + "num_attention_heads": 3, + "num_hidden_layers": 2, + "pad_token_id": 0, + "pooler_dropout": 0, + "pooler_hidden_act": "gelu", + "pooler_hidden_size": 12, + "pos_att_type": [ + "c2p", + "p2c" + ], + "position_biased_input": false, + "position_buckets": 8, + "relative_attention": true, + "share_att_key": true, + "transformers_version": "4.57.6", + "type_vocab_size": 0, + "vocab_size": 32 +} diff --git a/test/fixtures/gliner/local/encoder_expected.json b/test/fixtures/gliner/local/encoder_expected.json new file mode 100644 index 00000000..3649329d --- /dev/null +++ b/test/fixtures/gliner/local/encoder_expected.json @@ -0,0 +1 @@ +{"inputs":{"input_ids":[[2,3,4,5,6,7,8,9,10,11,12,13],[9,8,7,6,5,4,3,2,0,0,0,0]],"attention_mask":[[1,1,1,1,1,1,1,1,1,1,1,1],[1,1,1,1,1,1,1,1,0,0,0,0]]},"hidden_state":[[[-0.9749699234962463,1.0197323560714722,0.05689677223563194,2.1088759899139404,-1.130217432975769,-0.030761880800127983,-1.0234105587005615,1.2406774759292603,0.3019302189350128,-1.2038524150848389,-0.39101237058639526,0.026111694052815437],[-0.1356392353773117,0.3916495144367218,-0.035261671990156174,-0.10688740760087967,-0.5984441041946411,0.7537246346473694,-0.07780924439430237,1.6195348501205444,1.4266082048416138,0.12900526821613312,-2.2037739753723145,-1.1627064943313599],[-0.03985588997602463,-0.91127610206604,0.5267403721809387,-1.3032928705215454,0.38113006949424744,-0.9745063185691833,1.6519145965576172,1.1713287830352783,-0.25476643443107605,1.3858968019485474,-0.2349640130996704,-1.3983491659164429],[-1.413928508758545,1.558103084564209,-1.9861857891082764,0.026388416066765785,0.6154791116714478,-0.7946645021438599,1.1889212131500244,0.5364839434623718,0.21700553596019745,-0.7250324487686157,0.5316702127456665,0.24575985968112946],[0.9153928160667419,-2.4441192150115967,0.17886380851268768,1.0267003774642944,0.11110588908195496,-0.4989504814147949,0.9270901679992676,-1.0921049118041992,1.0123569965362549,0.6023404002189636,-0.11611360311508179,-0.6225622296333313],[0.09524653851985931,-1.2467247247695923,-1.5003842115402222,1.4010460376739502,-0.5309679508209229,-1.2952319383621216,0.7376418709754944,-0.04143829643726349,1.7889275550842285,-0.053045790642499924,0.7136768102645874,-0.06874646991491318],[0.7929097414016724,0.5531272292137146,-1.4603670835494995,1.190278172492981,-1.4889410734176636,-0.4162304997444153,0.9769580364227295,-1.4700028896331787,-0.5128310918807983,0.48892858624458313,1.2218220233917236,0.12434897571802139],[-0.35767731070518494,1.5619736909866333,-0.8594222664833069,-1.5427265167236328,-0.2729879319667816,0.4786938726902008,-0.21364636719226837,0.3212680518627167,-0.5572684407234192,-0.0063305930234491825,2.22440242767334,-0.7762786746025085],[1.5762172937393188,-0.4878321588039398,-0.6088244318962097,1.7241843938827515,-0.8398234844207764,-1.0739225149154663,0.9048086404800415,-0.5529155731201172,-0.11920813471078873,-1.5310094356536865,0.2986011803150177,0.7097240090370178],[1.365493655204773,-0.22735002636909485,0.36120671033859253,-1.4573370218276978,1.2294355630874634,-0.07859785109758377,-2.266423463821411,0.31731054186820984,0.6990044116973877,-0.4761694073677063,-0.06250321865081787,0.595930278301239],[1.6368097066879272,-1.1059705018997192,-0.13850119709968567,-0.5714932084083557,-1.335493564605713,-0.2906222641468048,-0.926969051361084,1.3688260316848755,0.7632933259010315,-0.8921260833740234,1.320156216621399,0.1720902919769287],[-1.9462778568267822,-1.238234519958496,-0.5553929805755615,1.4941781759262085,0.15323248505592346,0.24066376686096191,0.28036192059516907,0.4268314838409424,-1.1165692806243896,1.4852033853530884,0.532763659954071,0.24323977530002594]],[[-0.3586321771144867,1.5619277954101562,-0.8614831566810608,-1.5414738655090332,-0.2726227939128876,0.48000845313072205,-0.21346931159496307,0.32184460759162903,-0.557503342628479,-0.005044789984822273,2.2236642837524414,-0.777215838432312],[0.7918386459350586,0.5529963970184326,-1.4622037410736084,1.1913844347000122,-1.4883676767349243,-0.41486555337905884,0.9770107269287109,-1.469229817390442,-0.5129779577255249,0.4901409149169922,1.220903992652893,0.12336893379688263],[0.09427300840616226,-1.2465602159500122,-1.5021576881408691,1.4020707607269287,-0.5305028557777405,-1.2936900854110718,0.7377086281776428,-0.04086124151945114,1.788392186164856,-0.051774684339761734,0.7127991914749146,-0.06969712674617767],[0.9144044518470764,-2.444100856781006,0.17679251730442047,1.0279500484466553,0.11148149520158768,-0.4976246654987335,0.9272586703300476,-1.09149169921875,1.0121020078659058,0.6035878658294678,-0.1168631985783577,-0.6234967708587646],[-1.4144806861877441,1.5575906038284302,-1.9876437187194824,0.027655960991978645,0.6156631112098694,-0.7931388020515442,1.1887720823287964,0.5369158387184143,0.21672514081001282,-0.7235497236251831,0.5307620763778687,0.24472828209400177],[-0.040812697261571884,-0.9112955927848816,0.5246597528457642,-1.3019638061523438,0.3814748525619507,-0.9731408357620239,1.6520392894744873,1.1718659400939941,-0.2549915313720703,1.3870962858200073,-0.2356976419687271,-1.3992341756820679],[-0.13656321167945862,0.3914484977722168,-0.03729788959026337,-0.10556721687316895,-0.5978556871414185,0.754763126373291,-0.07757706195116043,1.6195365190505981,1.4258522987365723,0.13023078441619873,-2.203728199005127,-1.163241982460022],[-0.9757550954818726,1.0194793939590454,0.05484793707728386,2.1097609996795654,-1.1296464204788208,-0.029445305466651917,-1.0230340957641602,1.2410340309143066,0.3016539216041565,-1.2023674249649048,-0.39168596267700195,0.025158490985631943],[-0.597461998462677,0.43132078647613525,-0.9675875306129456,0.3237153887748718,1.1303560733795166,1.6185516119003296,0.9081758260726929,-1.231545329093933,0.3073790669441223,-1.5219855308532715,0.717634916305542,-1.1185530424118042],[-0.597461998462677,0.43132078647613525,-0.9675875306129456,0.3237153887748718,1.1303560733795166,1.6185516119003296,0.9081758260726929,-1.231545329093933,0.3073790669441223,-1.5219855308532715,0.717634916305542,-1.1185530424118042],[-0.597461998462677,0.43132078647613525,-0.9675875306129456,0.3237153887748718,1.1303560733795166,1.6185516119003296,0.9081758260726929,-1.231545329093933,0.3073790669441223,-1.5219855308532715,0.717634916305542,-1.1185530424118042],[-0.597461998462677,0.43132078647613525,-0.9675875306129456,0.3237153887748718,1.1303560733795166,1.6185516119003296,0.9081758260726929,-1.231545329093933,0.3073790669441223,-1.5219855308532715,0.717634916305542,-1.1185530424118042]]]} diff --git a/test/fixtures/gliner/local/expected.json b/test/fixtures/gliner/local/expected.json new file mode 100644 index 00000000..e557f554 --- /dev/null +++ b/test/fixtures/gliner/local/expected.json @@ -0,0 +1 @@ +{"inputs":{"text_states":[[[0.3396981954574585,1.273138165473938,-1.5371414422988892,1.0402002334594727,-1.1433438062667847,-0.6347954869270325,1.180337905883789,0.37527406215667725,0.8895691633224487,2.015190601348877,0.12288182973861694,0.283156156539917],[0.7347974181175232,1.4138669967651367,-0.06916153430938721,-0.2501976788043976,1.7436981201171875,-0.38771310448646545,0.7723278403282166,0.22671592235565186,-0.38377487659454346,-0.7681311964988708,-1.003414511680603,-0.9655314683914185],[0.2913444936275482,-1.1624996662139893,-0.4783642590045929,1.3870203495025635,-0.4341132640838623,0.2929425835609436,0.11547855287790298,-0.17881472408771515,-0.7945829629898071,-0.5694751739501953,0.0549420528113842,-0.7646364569664001],[-1.049290657043457,1.0828652381896973,-0.9056679606437683,-1.0400710105895996,2.1871161460876465,1.1617707014083862,0.6346917748451233,-0.4325243830680847,-1.275227665901184,-0.9807081818580627,0.7780191898345947,-0.23794014751911163],[1.774613618850708,1.5161902904510498,-1.0465848445892334,1.3759161233901978,0.05762922391295433,0.35209375619888306,1.0967378616333008,-0.9882838129997253,-0.5617773532867432,-0.4788392186164856,-0.7523331642150879,-0.4768691956996918],[1.2204824686050415,-0.15345095098018646,-0.5119454264640808,0.11924171447753906,0.023615580052137375,-0.7165514230728149,0.9313072562217712,0.11275327950716019,0.5183308720588684,0.10639318078756332,-0.0708099976181984,0.1868356168270111],[-1.135615348815918,-0.6121562123298645,0.47324639558792114,0.22348645329475403,0.4591689705848694,0.3639809787273407,0.24808230996131897,-0.5714194178581238,-1.8829892873764038,1.6257703304290771,-0.7037963271141052,0.558134913444519]]],"query_states":[[[0.42072793841362,1.2927827835083008,-0.062042392790317535,-0.07447009533643723,1.8144526481628418,0.4126001298427582,-0.5384427905082703,-0.6699069142341614,0.13211850821971893,-0.45546427369117737,0.7100200653076172,-1.4318643808364868],[1.3607994318008423,0.6645160913467407,-0.19414693117141724,-1.3718643188476562,1.7688463926315308,1.2299517393112183,0.29568976163864136,0.8400743007659912,0.3205632269382477,-0.5852316617965698,0.4060876965522766,0.2572997212409973],[0.8778835535049438,0.7318398952484131,0.27381399273872375,1.7188366651535034,-1.4703203439712524,-0.3415771424770355,1.2259759902954102,1.8914308547973633,-0.6524525284767151,-0.8879274725914001,-1.0803742408752441,1.4245182275772095]]]},"outputs":{"boundary_states":[[[0.5543657541275024,-0.6946316957473755,1.1090117692947388,-0.3196151554584503,0.9836046099662781,1.089868426322937,-2.386857748031616,-0.6547727584838867],[-0.12505602836608887,0.24557068943977356,1.2140291929244995,1.504260540008545,-0.4793061316013336,1.0113677978515625,-2.1620192527770996,-0.8438272476196289],[0.18238911032676697,2.1066346168518066,-0.4316194951534271,-1.0760812759399414,-0.24301744997501373,-0.129335418343544,0.21105362474918365,-1.1436104774475098],[-1.348842978477478,-1.361039638519287,0.8845534324645996,0.9332841634750366,1.2417471408843994,1.2107521295547485,-0.4067049026489258,-0.8042908906936646],[-0.06912174820899963,0.04821866750717163,-0.5649689435958862,-0.9744324684143066,2.0869693756103516,-0.07302378118038177,-0.2024613618850708,-0.6368979215621948],[-0.46103525161743164,0.4610676169395447,0.985417902469635,0.6738717555999756,0.6586606502532959,1.5966030359268188,-1.3193405866622925,-2.240709066390991],[-1.1184861660003662,-0.8969690799713135,0.8981121778488159,0.9026665687561035,1.3187808990478516,1.0598702430725098,-1.6596120595932007,-0.7316306233406067],[0.0012756064534187317,1.9650238752365112,0.3999372124671936,-1.9156148433685303,0.012058913707733154,-0.0032021254301071167,-0.4125906229019165,0.5017411708831787]]],"start_logits":[[[-0.6482831239700317,0.6374261975288391,0.5836296677589417,-0.27536410093307495,-0.6846007704734802,0.5613149404525757,-0.30481982231140137,-0.17659202218055725],[0.4975072145462036,1.8872268199920654,0.47561782598495483,0.26484695076942444,-0.872822105884552,1.4976412057876587,0.49400463700294495,-0.12412828207015991],[-0.1879054605960846,-0.32712504267692566,0.18928766250610352,-0.38722938299179077,0.04240357503294945,-0.27273285388946533,-0.3961479961872101,0.05433414503931999]]],"end_logits":[[[-0.23483985662460327,-0.11300341039896011,-0.018318362534046173,-0.2238014191389084,-0.0790729969739914,-0.26191446185112,-0.29262712597846985,-0.2385561764240265],[-0.28257569670677185,-0.23248818516731262,-0.21698512136936188,0.04880421608686447,0.1754920929670334,-0.22486552596092224,-0.03418394550681114,-0.24812285602092743],[0.16538938879966736,0.21505479514598846,-0.2762315273284912,0.30663126707077026,-0.09249747544527054,0.24200713634490967,0.3350890874862671,-0.3351801931858063]]],"inside_logits":[[[0.14138878881931305,0.5262569785118103,-0.23322194814682007,0.47248613834381104,0.21367579698562622,0.09782277047634125,0.2908124625682831],[0.19397881627082825,0.32151395082473755,0.07189678400754929,0.37444400787353516,0.4646332561969757,0.04348849877715111,0.4586145579814911],[0.13514390587806702,-0.6980278491973877,0.32489508390426636,-0.5097112655639648,-0.2958136796951294,0.03978963941335678,0.1785593181848526]]],"null_logits":[[[0.0],[0.0],[0.0]]],"pool_start":[[[0.01973271369934082,-1.1949794292449951,-0.33095377683639526,0.17409272491931915,-0.08674687147140503,0.417887806892395,-0.40595245361328125,-0.7204489707946777],[0.25336021184921265,-1.7361632585525513,0.06955176591873169,0.0656891018152237,-0.7567868232727051,0.2901933193206787,-0.19905252754688263,0.25979870557785034],[0.6327720880508423,-0.5307411551475525,0.09591835737228394,0.28087395429611206,0.06709796190261841,-0.3107140362262726,0.20812644064426422,0.3538324534893036],[-0.007063537836074829,-0.4119642972946167,-0.805838942527771,0.3992084264755249,-0.7320541143417358,0.7741769552230835,-0.8107820749282837,-0.5652262568473816],[0.02349698543548584,-0.1681683361530304,-0.5625001192092896,0.2526845335960388,0.29159116744995117,0.3053928017616272,-0.1800398826599121,-0.9002181887626648],[0.92133629322052,-1.7017991542816162,-0.3539333641529083,0.3486838936805725,-0.36068424582481384,0.47266677021980286,-0.3804474472999573,0.1909562349319458],[0.0800061821937561,-0.9631006121635437,-0.6667594909667969,0.42838162183761597,-0.653972327709198,0.7486647367477417,-0.6650055050849915,-0.7951223850250244],[0.3218591809272766,0.030715562403202057,-0.28869688510894775,0.7885398864746094,-0.09742344915866852,-0.32081612944602966,-0.07769119739532471,-0.1758747398853302]]],"pool_end":[[[-0.535184383392334,0.12630656361579895,-0.2958766520023346,0.9383572340011597,0.24714873731136322,0.6724323034286499,0.47279682755470276,0.6263011693954468],[-0.8006899356842041,0.07200932502746582,-0.41277965903282166,0.7090768218040466,0.475746750831604,0.45147380232810974,0.05292518064379692,0.5249415636062622],[0.6143132448196411,-0.1769571602344513,0.6494200825691223,0.008416205644607544,-0.5634390115737915,-0.7757945656776428,-0.15646545588970184,0.9808950424194336],[0.4135632812976837,-0.6896292567253113,-1.6605653762817383,0.01303565502166748,-0.3898886442184448,1.5887058973312378,-0.5022797584533691,0.4603557884693146],[0.6591869592666626,-0.242180734872818,-0.31260207295417786,-0.3302041292190552,-0.8187612295150757,0.43586793541908264,-0.4738417863845825,0.38021525740623474],[0.29227858781814575,-0.30225396156311035,-1.0902901887893677,0.18019182980060577,-0.2732589840888977,0.9124890565872192,-0.4606233835220337,1.1758697032928467],[-0.11628314852714539,-0.3240511417388916,-1.2885268926620483,0.20421501994132996,-0.25647372007369995,1.4082460403442383,-0.29634588956832886,0.3295486867427826],[-0.06060993671417236,-0.6633170247077942,0.9955673813819885,0.616875946521759,-0.5671885013580322,-0.7189795970916748,0.6791592836380005,0.7461681365966797]]]},"scorer_inputs":{"text_states":[[[0.3396981954574585,1.273138165473938,-1.5371414422988892,1.0402002334594727,-1.1433438062667847,-0.6347954869270325,1.180337905883789,0.37527406215667725,0.8895691633224487,2.015190601348877,0.12288182973861694,0.283156156539917],[0.7347974181175232,1.4138669967651367,-0.06916153430938721,-0.2501976788043976,1.7436981201171875,-0.38771310448646545,0.7723278403282166,0.22671592235565186,-0.38377487659454346,-0.7681311964988708,-1.003414511680603,-0.9655314683914185],[0.2913444936275482,-1.1624996662139893,-0.4783642590045929,1.3870203495025635,-0.4341132640838623,0.2929425835609436,0.11547855287790298,-0.17881472408771515,-0.7945829629898071,-0.5694751739501953,0.0549420528113842,-0.7646364569664001],[-1.049290657043457,1.0828652381896973,-0.9056679606437683,-1.0400710105895996,2.1871161460876465,1.1617707014083862,0.6346917748451233,-0.4325243830680847,-1.275227665901184,-0.9807081818580627,0.7780191898345947,-0.23794014751911163],[1.774613618850708,1.5161902904510498,-1.0465848445892334,1.3759161233901978,0.05762922391295433,0.35209375619888306,1.0967378616333008,-0.9882838129997253,-0.5617773532867432,-0.4788392186164856,-0.7523331642150879,-0.4768691956996918],[1.2204824686050415,-0.15345095098018646,-0.5119454264640808,0.11924171447753906,0.023615580052137375,-0.7165514230728149,0.9313072562217712,0.11275327950716019,0.5183308720588684,0.10639318078756332,-0.0708099976181984,0.1868356168270111],[-1.135615348815918,-0.6121562123298645,0.47324639558792114,0.22348645329475403,0.4591689705848694,0.3639809787273407,0.24808230996131897,-0.5714194178581238,-1.8829892873764038,1.6257703304290771,-0.7037963271141052,0.558134913444519]]],"query_states":[[[0.42072793841362,1.2927827835083008,-0.062042392790317535,-0.07447009533643723,1.8144526481628418,0.4126001298427582,-0.5384427905082703,-0.6699069142341614,0.13211850821971893,-0.45546427369117737,0.7100200653076172,-1.4318643808364868],[1.3607994318008423,0.6645160913467407,-0.19414693117141724,-1.3718643188476562,1.7688463926315308,1.2299517393112183,0.29568976163864136,0.8400743007659912,0.3205632269382477,-0.5852316617965698,0.4060876965522766,0.2572997212409973],[0.8778835535049438,0.7318398952484131,0.27381399273872375,1.7188366651535034,-1.4703203439712524,-0.3415771424770355,1.2259759902954102,1.8914308547973633,-0.6524525284767151,-0.8879274725914001,-1.0803742408752441,1.4245182275772095]]],"boundary_states":[[[0.5543657541275024,-0.6946316957473755,1.1090117692947388,-0.3196151554584503,0.9836046099662781,1.089868426322937,-2.386857748031616,-0.6547727584838867],[-0.12505602836608887,0.24557068943977356,1.2140291929244995,1.504260540008545,-0.4793061316013336,1.0113677978515625,-2.1620192527770996,-0.8438272476196289],[0.18238911032676697,2.1066346168518066,-0.4316194951534271,-1.0760812759399414,-0.24301744997501373,-0.129335418343544,0.21105362474918365,-1.1436104774475098],[-1.348842978477478,-1.361039638519287,0.8845534324645996,0.9332841634750366,1.2417471408843994,1.2107521295547485,-0.4067049026489258,-0.8042908906936646],[-0.06912174820899963,0.04821866750717163,-0.5649689435958862,-0.9744324684143066,2.0869693756103516,-0.07302378118038177,-0.2024613618850708,-0.6368979215621948],[-0.46103525161743164,0.4610676169395447,0.985417902469635,0.6738717555999756,0.6586606502532959,1.5966030359268188,-1.3193405866622925,-2.240709066390991],[-1.1184861660003662,-0.8969690799713135,0.8981121778488159,0.9026665687561035,1.3187808990478516,1.0598702430725098,-1.6596120595932007,-0.7316306233406067],[0.0012756064534187317,1.9650238752365112,0.3999372124671936,-1.9156148433685303,0.012058913707733154,-0.0032021254301071167,-0.4125906229019165,0.5017411708831787]]],"start_logits":[[[-0.6482831239700317,0.6374261975288391,0.5836296677589417,-0.27536410093307495,-0.6846007704734802,0.5613149404525757,-0.30481982231140137,-0.17659202218055725],[0.4975072145462036,1.8872268199920654,0.47561782598495483,0.26484695076942444,-0.872822105884552,1.4976412057876587,0.49400463700294495,-0.12412828207015991],[-0.1879054605960846,-0.32712504267692566,0.18928766250610352,-0.38722938299179077,0.04240357503294945,-0.27273285388946533,-0.3961479961872101,0.05433414503931999]]],"end_logits":[[[-0.23483985662460327,-0.11300341039896011,-0.018318362534046173,-0.2238014191389084,-0.0790729969739914,-0.26191446185112,-0.29262712597846985,-0.2385561764240265],[-0.28257569670677185,-0.23248818516731262,-0.21698512136936188,0.04880421608686447,0.1754920929670334,-0.22486552596092224,-0.03418394550681114,-0.24812285602092743],[0.16538938879966736,0.21505479514598846,-0.2762315273284912,0.30663126707077026,-0.09249747544527054,0.24200713634490967,0.3350890874862671,-0.3351801931858063]]],"inside_logits":[[[0.14138878881931305,0.5262569785118103,-0.23322194814682007,0.47248613834381104,0.21367579698562622,0.09782277047634125,0.2908124625682831],[0.19397881627082825,0.32151395082473755,0.07189678400754929,0.37444400787353516,0.4646332561969757,0.04348849877715111,0.4586145579814911],[0.13514390587806702,-0.6980278491973877,0.32489508390426636,-0.5097112655639648,-0.2958136796951294,0.03978963941335678,0.1785593181848526]]],"starts":[[0,1,1,5,1,0,2,0]],"ends":[[3,3,6,6,5,6,5,5]],"valid":[[1,1,1,1,1,1,1,0]],"compat":[[0.6908627152442932,0.7647802233695984,0.425785630941391,0.6736619472503662,0.49624529480934143,0.4739565849304199,0.10950657725334167,0.1781979501247406]]},"scores":[[[-0.5826656818389893,0.7399387359619141,0.8322502970695496,0.08617706596851349,0.8079578876495361,-0.5206392407417297,0.7217500805854797,-10000.0],[1.1015548706054688,2.5385539531707764,2.6989235877990723,1.595247507095337,2.544283151626587,1.2673475742340088,1.1250475645065308,-10000.0],[0.6889083385467529,0.11803781986236572,-0.08281180262565613,0.9682692289352417,-0.1483369767665863,0.4464244246482849,0.33433201909065247,-10000.0]]]} diff --git a/test/fixtures/gliner/local/model.safetensors b/test/fixtures/gliner/local/model.safetensors new file mode 100644 index 00000000..8ff10e7a Binary files /dev/null and b/test/fixtures/gliner/local/model.safetensors differ diff --git a/test/fixtures/gliner/ner_reference.json b/test/fixtures/gliner/ner_reference.json new file mode 100644 index 00000000..c703b56d --- /dev/null +++ b/test/fixtures/gliner/ner_reference.json @@ -0,0 +1,238 @@ +{ + "model": "fastino/gliner2.5-base-v1", + "revision": "78cea040597df251eedefa9d7ee2a756af39fe64", + "reference": "gliner2 2.0.0, torch 2.14.0, CPU float32", + "examples": [ + { + "text": "Apple CEO Tim Cook announced the iPhone 15 in Cupertino.", + "labels": [ + "company", + "person", + "product", + "location" + ], + "prediction": { + "entities": { + "company": [ + { + "text": "Apple", + "confidence": 0.9979098439216614, + "start": 0, + "end": 5 + } + ], + "person": [ + { + "text": "Tim Cook", + "confidence": 0.995536208152771, + "start": 10, + "end": 18 + } + ], + "product": [ + { + "text": "iPhone 15", + "confidence": 0.9923769235610962, + "start": 33, + "end": 42 + } + ], + "location": [ + { + "text": "Cupertino", + "confidence": 0.9973949193954468, + "start": 46, + "end": 55 + } + ] + } + } + }, + { + "text": "It was quiet and nothing happened.", + "labels": [ + "company", + "person", + "product", + "location" + ], + "prediction": { + "entities": { + "company": [], + "person": [], + "product": [], + "location": [] + } + } + }, + { + "text": "İpek met José in São Paulo. 👩🏽‍💻 Alice emailed bob@example.com.", + "labels": [ + "person", + "location", + "email" + ], + "prediction": { + "entities": { + "person": [ + { + "text": "Alice", + "confidence": 0.9852306246757507, + "start": 33, + "end": 38 + }, + { + "text": "José", + "confidence": 0.9819398522377014, + "start": 9, + "end": 13 + }, + { + "text": "İpek", + "confidence": 0.5640807747840881, + "start": 0, + "end": 4 + } + ], + "location": [ + { + "text": "São Paulo", + "confidence": 0.9983048439025879, + "start": 17, + "end": 26 + } + ], + "email": [ + { + "text": "bob@example.com", + "confidence": 0.9971128702163696, + "start": 47, + "end": 62 + } + ] + } + } + }, + { + "text": "The University of California, Berkeley is in California. New York University is in New York City.", + "labels": [ + "organization", + "location" + ], + "prediction": { + "entities": { + "organization": [ + { + "text": "New York University", + "confidence": 0.9389233589172363, + "start": 57, + "end": 76 + }, + { + "text": "Berkeley", + "confidence": 0.6230217814445496, + "start": 30, + "end": 38 + }, + { + "text": "The University of California", + "confidence": 0.5583738088607788, + "start": 0, + "end": 28 + } + ], + "location": [ + { + "text": "California", + "confidence": 0.9944834113121033, + "start": 45, + "end": 55 + }, + { + "text": "New York City", + "confidence": 0.9925059676170349, + "start": 83, + "end": 96 + }, + { + "text": "California", + "confidence": 0.9822803139686584, + "start": 18, + "end": 28 + } + ] + } + } + }, + { + "text": "The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. The committee discussed the agenda and reviewed the previous meeting notes. Microsoft CEO Satya Nadella met Tim Cook in Seattle.", + "labels": [ + "company", + "person", + "location" + ], + "prediction": { + "entities": { + "company": [ + { + "text": "Microsoft", + "confidence": 0.996722400188446, + "start": 1368, + "end": 1377 + } + ], + "person": [ + { + "text": "Tim Cook", + "confidence": 0.9962313771247864, + "start": 1400, + "end": 1408 + }, + { + "text": "Satya Nadella", + "confidence": 0.9937824606895447, + "start": 1382, + "end": 1395 + } + ], + "location": [ + { + "text": "Seattle", + "confidence": 0.9983269572257996, + "start": 1412, + "end": 1419 + } + ] + } + } + }, + { + "text": "( Source : ANI )", + "labels": [ + "person", + "location", + "corporation", + "product", + "creative-work", + "group" + ], + "prediction": { + "entities": { + "person": [ + { + "text": "ANI", + "confidence": 0.7916450500488281, + "start": 11, + "end": 14 + } + ], + "location": [], + "corporation": [], + "product": [], + "creative-work": [], + "group": [] + } + } + } + ] +} diff --git a/test/fixtures/gliner/no_refinement/config.json b/test/fixtures/gliner/no_refinement/config.json new file mode 100644 index 00000000..a34fd420 --- /dev/null +++ b/test/fixtures/gliner/no_refinement/config.json @@ -0,0 +1,92 @@ +{ + "architectures": [ + "BoundaryExtractor" + ], + "architecture": "boundary", + "token_pooling": "first", + "boundary_head": { + "boundary_dim": 8, + "pair_dim": 8, + "boundary_refinement_layers": 0, + "boundary_ffn_multiplier": 2.0, + "start_top_k": 16, + "end_top_k": 16, + "ends_per_start": 8, + "starts_per_end": 8, + "candidate_budget": 128, + "training_candidate_budget": 160, + "max_gold_per_query": 32, + "end_block_size": 256, + "bidirectional_proposals": true, + "use_inside_evidence": true, + "dropout": 0.1, + "export_mode": "auto", + "vectorized_pair_elements": 16777216, + "boundary_negative_weight": 1.0, + "boundary_marginal_loss": "bce", + "loss_reduction": "global", + "boundary_focal_gamma_positive": 0.0, + "boundary_focal_gamma_negative": 2.0, + "boundary_focal_clip": 0.05, + "hard_negatives_per_positive": 5, + "minimum_hard_negatives": 8, + "hard_negative_keep_all_when_absent": false, + "enable_span_content": true, + "content_dim": 4, + "content_soft_max_pool": false, + "enable_rotary_endpoints": false, + "rotary_base": 10000.0, + "boundary_attention_layers": 0, + "boundary_attention_heads": 2, + "boundary_attention_window": 0, + "query_conditioned_inside_weight": false, + "endpoint_difference_features": false, + "reranker_endpoint_compat": true, + "multihead_pair_compat_heads": 8, + "boundary_top_k_alpha": 0.0, + "boundary_top_k_max": 128, + "boundary_top_k_bucket": 8, + "candidate_pool": "shared", + "pool_boundary_top_k": 4, + "pool_size": 8, + "min_pool_per_query": 2, + "candidate_attention_layers": 0, + "candidate_attention_heads": 4, + "query_attention_layers": 0, + "enable_abstention": true, + "abstention_threshold": 0.5, + "proposal_loss_weight": 0.3, + "consistency_loss_weight": 0.1, + "rerank_listwise_weight": 0.3, + "soft_iou_aux_weight": 0.2, + "soft_iou_anneal_steps": 20000, + "abstention_loss_weight": 0.2, + "consistency_warmup_steps": 2000, + "enable_count_head": true, + "count_loss_weight": 0.2, + "adaptive_threshold": false, + "overlap_policy": "flat", + "pair_temperature": 1.0, + "relation_temperature": 1.0, + "record_temperature": 1.0, + "classification_temperature": 1.0, + "negative_query_ratio": 0.5, + "max_negative_queries_per_batch": 64, + "classification_loss_weight": 1.0, + "enable_records": true, + "record_dim": 128, + "record_instance_queries": 32, + "record_anchor_proposal_threshold": 0.2, + "record_anchor_threshold": 0.5, + "record_field_threshold": 0.5, + "record_loss_weight": 1.0, + "enable_relations": true, + "relation_heads_per_type": 32, + "relation_tails_per_type": 32, + "relation_pair_cap": 128, + "relation_loss_weight": 1.0, + "relation_argument_proposal_threshold": 0.0, + "directional_relation_states": false, + "relation_biaffine_content": false + } +} diff --git a/test/fixtures/gliner/no_refinement/expected.json b/test/fixtures/gliner/no_refinement/expected.json new file mode 100644 index 00000000..61a273e7 --- /dev/null +++ b/test/fixtures/gliner/no_refinement/expected.json @@ -0,0 +1 @@ +{"inputs":{"text_states":[[[-0.89991295337677,0.4791645407676697,-0.24811258912086487,-1.3093348741531372,-0.2919580638408661,0.45643654465675354,-0.3147018849849701,-0.41328999400138855,-0.17061950266361237,-0.6140995025634766,0.53551185131073,-1.6644023656845093],[-0.15447334945201874,-1.6013476848602295,-0.059470854699611664,-0.9928666353225708,0.48605024814605713,-0.4023412764072418,0.11986448615789413,-0.17284131050109863,0.8866234421730042,-0.8611426949501038,-0.2726478576660156,0.9804211258888245],[0.4526977837085724,0.41128507256507874,-2.4095988273620605,-0.5115848183631897,-0.5345802903175354,0.645367443561554,-2.9890658855438232,0.18370741605758667,1.526895523071289,-1.5928043127059937,0.6899217367172241,-0.47910887002944946],[-0.9380512237548828,-0.6385068893432617,0.2194686383008957,-0.439242959022522,0.182944193482399,-0.28223374485969543,-1.7268664836883545,-0.500217616558075,-0.18189166486263275,-1.237966775894165,0.514220118522644,-0.15104348957538605],[-0.49643412232398987,0.14655134081840515,0.985339879989624,0.40662944316864014,0.40576910972595215,0.25104594230651855,0.1281542032957077,-0.1980288028717041,-0.834316611289978,0.7485802173614502,-2.0932791233062744,0.27361416816711426],[-1.2403943538665771,0.24919238686561584,-1.1415528059005737,0.78214031457901,-1.0222961902618408,-0.7083863019943237,0.0595916211605072,-0.686333179473877,0.8713390231132507,-0.5227583050727844,1.240007996559143,-0.905766487121582],[0.7680287957191467,1.622213363647461,0.08158040046691895,0.20281589031219482,0.33024170994758606,-0.9533718824386597,1.5734566450119019,1.8697383403778076,-1.0638948678970337,-0.22725839912891388,0.25006136298179626,1.161847710609436]]],"query_states":[[[-0.11422377079725266,-0.05629456788301468,0.8497495651245117,-0.8599119782447815,-0.6105663776397705,1.0629346370697021,1.2221823930740356,0.7718911170959473,-1.2797164916992188,-1.5433486700057983,-0.6020247936248779,0.321399986743927],[-0.060615554451942444,-1.170371651649475,-2.773613929748535,-0.0298247579485178,-0.9166154265403748,0.4702746868133545,1.8777929544448853,0.5223742127418518,-1.5005147457122803,-0.2404831200838089,-0.48548775911331177,-0.06616523861885071],[0.9090389609336853,0.29768696427345276,-0.4455963373184204,1.88334321975708,-0.6739510893821716,-0.9192121624946594,1.212047815322876,-1.3463493585586548,-1.0458097457885742,-0.8419854044914246,-0.5115393996238708,-0.09544137865304947]]]},"outputs":{"boundary_states":[[[-1.8617631196975708,1.2227681875228882,0.04745284095406532,1.5066418647766113,-0.6594765782356262,0.3662765622138977,-0.2834624946117401,-0.33843719959259033],[-0.9982035756111145,0.9203038215637207,0.021941613405942917,1.9233506917953491,-0.7904142141342163,0.09155687689781189,0.1706259697675705,-1.3391609191894531],[1.0349376201629639,0.7226118445396423,-1.2603429555892944,-0.862286388874054,1.6290838718414307,-0.6853060126304626,0.3362111747264862,-0.9149091839790344],[-1.1463432312011719,-1.1197313070297241,0.3534020483493805,1.7380131483078003,0.04911056160926819,1.1911243200302124,-0.14723235368728638,-0.918343186378479],[-1.0889256000518799,-0.7391351461410522,0.483184814453125,-0.6978680491447449,1.5140292644500732,1.4501978158950806,-1.0631271600723267,0.14164434373378754],[0.7139455080032349,1.9051191806793213,0.3262629210948944,-1.7318700551986694,-0.4083448648452759,-0.057037293910980225,-0.7643734812736511,0.016298318281769753],[-0.4227040410041809,-0.6240066289901733,0.3993285298347473,0.8959555625915527,0.800061047077179,0.6567660570144653,-2.2569401264190674,0.5515395402908325],[-0.7994303107261658,-0.5049539804458618,-1.4332786798477173,0.08441617339849472,1.8952529430389404,1.1585052013397217,-0.3195914924144745,-0.08091980963945389]]],"start_logits":[[[-0.03419455885887146,0.28550541400909424,-0.2796443700790405,0.3283366858959198,-0.19521768391132355,-0.07055336982011795,-0.1741439253091812,-0.5065072774887085],[-0.2201099693775177,0.13224463164806366,-0.06358881294727325,-0.23888006806373596,-0.5756077766418457,0.5616352558135986,-0.42427128553390503,-0.8910267353057861],[0.014287570491433144,0.05423668026924133,-0.04447522759437561,0.1917112171649933,0.10021188855171204,-0.1560019999742508,0.04757707566022873,0.06888562440872192]]],"end_logits":[[[0.09803540259599686,0.13671836256980896,-0.14295972883701324,0.10699134320020676,0.05763520300388336,0.262997567653656,0.05438602715730667,-0.18524807691574097],[-0.05687681958079338,-0.07781114429235458,-0.020855272188782692,-0.4843590259552002,-0.23484276235103607,0.6738273501396179,-0.5490736365318298,-0.5801032781600952],[0.07692018896341324,0.16919909417629242,-0.252502977848053,0.46449559926986694,0.15919138491153717,-0.2280358523130417,0.22560830414295197,0.07518421858549118]]],"inside_logits":[[[-0.19842468202114105,-0.5580665469169617,-0.1554490327835083,-0.22552064061164856,-0.09711020439863205,0.0878102108836174,-0.19905145466327667],[-0.19563154876232147,0.0558164082467556,0.7565726637840271,0.37804996967315674,-0.18222364783287048,-0.054141923785209656,-0.38786834478378296],[-0.3943777084350586,0.38990938663482666,0.09474824368953705,-0.0009076485293917358,-0.04401547461748123,0.117103211581707,-0.020098255947232246]]],"null_logits":[[[0.0],[0.0],[0.0]]],"pool_start":[[[0.7863291501998901,0.34011560678482056,-0.09461970627307892,0.5731335282325745,-0.402275025844574,0.9183982014656067,-0.7834705114364624,-0.8572399616241455],[0.7547451257705688,0.1407146155834198,-0.06705810129642487,0.8116952776908875,-0.26734915375709534,0.47609320282936096,-0.6977919936180115,-0.5419049263000488],[-0.8783164620399475,0.04503973573446274,-0.22664792835712433,0.30836057662963867,-0.20475731790065765,-0.8183746337890625,-0.17269912362098694,0.9312818050384521],[0.9958206415176392,-0.11268576979637146,0.3783947229385376,1.0419126749038696,-0.5295219421386719,0.7820128202438354,-0.5363039970397949,-0.26253262162208557],[0.04450015723705292,-0.025100700557231903,0.44415906071662903,0.5792513489723206,-0.12050661444664001,0.883765459060669,-0.062252938747406006,0.3078513741493225],[-0.6005741357803345,0.16324961185455322,-0.08062459528446198,-0.6782263517379761,0.7422529458999634,-0.10377264022827148,-0.06375880539417267,0.3211195766925812],[0.8625158071517944,0.01582030951976776,-0.5306178331375122,0.11103355139493942,-0.1687450408935547,0.9531468152999878,-0.4430171251296997,-0.34168896079063416],[0.2525477409362793,0.31104350090026855,-0.21691878139972687,0.5709354281425476,-0.8237012624740601,0.2389921396970749,-0.7905511856079102,-0.030709266662597656]]],"pool_end":[[[0.20017844438552856,-0.9105241298675537,0.5067005753517151,-0.208455428481102,0.34463655948638916,0.8419100046157837,-0.41018134355545044,0.3261953592300415],[0.1855366826057434,-0.9199000000953674,0.5311906933784485,-0.13315649330615997,0.11834390461444855,1.2398862838745117,-0.30726197361946106,-0.3424519896507263],[-0.30654752254486084,-0.7567541599273682,0.08360825479030609,1.0629184246063232,-0.051765695214271545,-0.46291980147361755,-0.28253641724586487,0.0240384042263031],[0.3165682256221771,-0.2948687672615051,-0.03365997225046158,-0.7316427230834961,0.21398872137069702,0.6069812178611755,0.17922183871269226,-0.6238462328910828],[-0.24961654841899872,-0.21872997283935547,-0.3410085439682007,-0.10256163775920868,0.7060184478759766,-0.7078277468681335,0.6407865881919861,0.5947461128234863],[-0.29880189895629883,-0.4957456886768341,0.6148008108139038,1.1756634712219238,0.3181438744068146,-0.3092899024486542,0.26090729236602783,1.145919919013977],[-0.13444294035434723,-0.48136261105537415,0.12064729630947113,-0.312005877494812,-0.34497761726379395,-0.564429759979248,0.485093355178833,0.5168750286102295],[0.1528315544128418,-0.7347306609153748,0.03286682069301605,0.2593191862106323,0.10985235869884491,-0.6413516402244568,-0.14382176101207733,-0.1573278158903122]]]},"scorer_inputs":{"text_states":[[[-0.89991295337677,0.4791645407676697,-0.24811258912086487,-1.3093348741531372,-0.2919580638408661,0.45643654465675354,-0.3147018849849701,-0.41328999400138855,-0.17061950266361237,-0.6140995025634766,0.53551185131073,-1.6644023656845093],[-0.15447334945201874,-1.6013476848602295,-0.059470854699611664,-0.9928666353225708,0.48605024814605713,-0.4023412764072418,0.11986448615789413,-0.17284131050109863,0.8866234421730042,-0.8611426949501038,-0.2726478576660156,0.9804211258888245],[0.4526977837085724,0.41128507256507874,-2.4095988273620605,-0.5115848183631897,-0.5345802903175354,0.645367443561554,-2.9890658855438232,0.18370741605758667,1.526895523071289,-1.5928043127059937,0.6899217367172241,-0.47910887002944946],[-0.9380512237548828,-0.6385068893432617,0.2194686383008957,-0.439242959022522,0.182944193482399,-0.28223374485969543,-1.7268664836883545,-0.500217616558075,-0.18189166486263275,-1.237966775894165,0.514220118522644,-0.15104348957538605],[-0.49643412232398987,0.14655134081840515,0.985339879989624,0.40662944316864014,0.40576910972595215,0.25104594230651855,0.1281542032957077,-0.1980288028717041,-0.834316611289978,0.7485802173614502,-2.0932791233062744,0.27361416816711426],[-1.2403943538665771,0.24919238686561584,-1.1415528059005737,0.78214031457901,-1.0222961902618408,-0.7083863019943237,0.0595916211605072,-0.686333179473877,0.8713390231132507,-0.5227583050727844,1.240007996559143,-0.905766487121582],[0.7680287957191467,1.622213363647461,0.08158040046691895,0.20281589031219482,0.33024170994758606,-0.9533718824386597,1.5734566450119019,1.8697383403778076,-1.0638948678970337,-0.22725839912891388,0.25006136298179626,1.161847710609436]]],"query_states":[[[-0.11422377079725266,-0.05629456788301468,0.8497495651245117,-0.8599119782447815,-0.6105663776397705,1.0629346370697021,1.2221823930740356,0.7718911170959473,-1.2797164916992188,-1.5433486700057983,-0.6020247936248779,0.321399986743927],[-0.060615554451942444,-1.170371651649475,-2.773613929748535,-0.0298247579485178,-0.9166154265403748,0.4702746868133545,1.8777929544448853,0.5223742127418518,-1.5005147457122803,-0.2404831200838089,-0.48548775911331177,-0.06616523861885071],[0.9090389609336853,0.29768696427345276,-0.4455963373184204,1.88334321975708,-0.6739510893821716,-0.9192121624946594,1.212047815322876,-1.3463493585586548,-1.0458097457885742,-0.8419854044914246,-0.5115393996238708,-0.09544137865304947]]],"boundary_states":[[[-1.8617631196975708,1.2227681875228882,0.04745284095406532,1.5066418647766113,-0.6594765782356262,0.3662765622138977,-0.2834624946117401,-0.33843719959259033],[-0.9982035756111145,0.9203038215637207,0.021941613405942917,1.9233506917953491,-0.7904142141342163,0.09155687689781189,0.1706259697675705,-1.3391609191894531],[1.0349376201629639,0.7226118445396423,-1.2603429555892944,-0.862286388874054,1.6290838718414307,-0.6853060126304626,0.3362111747264862,-0.9149091839790344],[-1.1463432312011719,-1.1197313070297241,0.3534020483493805,1.7380131483078003,0.04911056160926819,1.1911243200302124,-0.14723235368728638,-0.918343186378479],[-1.0889256000518799,-0.7391351461410522,0.483184814453125,-0.6978680491447449,1.5140292644500732,1.4501978158950806,-1.0631271600723267,0.14164434373378754],[0.7139455080032349,1.9051191806793213,0.3262629210948944,-1.7318700551986694,-0.4083448648452759,-0.057037293910980225,-0.7643734812736511,0.016298318281769753],[-0.4227040410041809,-0.6240066289901733,0.3993285298347473,0.8959555625915527,0.800061047077179,0.6567660570144653,-2.2569401264190674,0.5515395402908325],[-0.7994303107261658,-0.5049539804458618,-1.4332786798477173,0.08441617339849472,1.8952529430389404,1.1585052013397217,-0.3195914924144745,-0.08091980963945389]]],"start_logits":[[[-0.03419455885887146,0.28550541400909424,-0.2796443700790405,0.3283366858959198,-0.19521768391132355,-0.07055336982011795,-0.1741439253091812,-0.5065072774887085],[-0.2201099693775177,0.13224463164806366,-0.06358881294727325,-0.23888006806373596,-0.5756077766418457,0.5616352558135986,-0.42427128553390503,-0.8910267353057861],[0.014287570491433144,0.05423668026924133,-0.04447522759437561,0.1917112171649933,0.10021188855171204,-0.1560019999742508,0.04757707566022873,0.06888562440872192]]],"end_logits":[[[0.09803540259599686,0.13671836256980896,-0.14295972883701324,0.10699134320020676,0.05763520300388336,0.262997567653656,0.05438602715730667,-0.18524807691574097],[-0.05687681958079338,-0.07781114429235458,-0.020855272188782692,-0.4843590259552002,-0.23484276235103607,0.6738273501396179,-0.5490736365318298,-0.5801032781600952],[0.07692018896341324,0.16919909417629242,-0.252502977848053,0.46449559926986694,0.15919138491153717,-0.2280358523130417,0.22560830414295197,0.07518421858549118]]],"inside_logits":[[[-0.19842468202114105,-0.5580665469169617,-0.1554490327835083,-0.22552064061164856,-0.09711020439863205,0.0878102108836174,-0.19905145466327667],[-0.19563154876232147,0.0558164082467556,0.7565726637840271,0.37804996967315674,-0.18222364783287048,-0.054141923785209656,-0.38786834478378296],[-0.3943777084350586,0.38990938663482666,0.09474824368953705,-0.0009076485293917358,-0.04401547461748123,0.117103211581707,-0.020098255947232246]]],"starts":[[1,1,3,4,5,3,4,1]],"ends":[[3,5,5,5,6,6,6,6]],"valid":[[1,1,1,1,1,1,1,0]],"compat":[[0.017889071255922318,-0.14763322472572327,0.12897154688835144,0.3458012044429779,0.050065137445926666,-0.35837626457214355,-0.1588772088289261,-0.43332639336586]]},"scores":[[[0.6662091016769409,0.5585938692092896,1.0786570310592651,1.2464543581008911,0.6953462958335876,1.2058172225952148,0.9891802668571472,-10000.0],[0.6442533731460571,1.5116736888885498,0.7383575439453125,0.7995238304138184,0.9620802998542786,-0.09749714285135269,-0.29160410165786743,-10000.0],[1.2381231784820557,0.2618332505226135,-0.007547348737716675,-0.20521345734596252,-0.09224306792020798,0.12184228748083115,0.018929380923509598,-10000.0]]]} diff --git a/test/fixtures/gliner/no_refinement/model.safetensors b/test/fixtures/gliner/no_refinement/model.safetensors new file mode 100644 index 00000000..3247713f Binary files /dev/null and b/test/fixtures/gliner/no_refinement/model.safetensors differ