mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 13:38:14 -05:00
models: add gliner-small-v2.1 mlx support
Added native GLiNER entity extraction to the MLX runner through `/api/extract` and the Go client, returning entity labels, scores, and Unicode offsets. This includes tokenization, DeBERTa encoding, span scoring, model import and metadata, an export script, and documentation. Initial support targets float32 DeBERTa GLiNER `markerV0` models. The build and affected tests pass. I verified Python checkpoint parity and HTTP extraction, including concurrent requests and input validation, on MLX CPU.
This commit is contained in:
Bruce MacDonald
@@ -435,6 +435,15 @@ func (c *Client) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse,
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// Extract returns scored entity spans from a local extraction model.
|
||||
func (c *Client) Extract(ctx context.Context, req *ExtractRequest) (*ExtractResponse, error) {
|
||||
var resp ExtractResponse
|
||||
if err := c.do(ctx, http.MethodPost, "/api/extract", req, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// Embeddings generates an embedding from a model.
|
||||
func (c *Client) Embeddings(ctx context.Context, req *EmbeddingRequest) (*EmbeddingResponse, error) {
|
||||
var resp EmbeddingResponse
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// ExtractRequest finds entities of the supplied types in text.
|
||||
type ExtractRequest struct {
|
||||
Model string `json:"model"`
|
||||
Input string `json:"input"`
|
||||
Labels []string `json:"labels"`
|
||||
Threshold *float32 `json:"threshold,omitempty"`
|
||||
KeepAlive *Duration `json:"keep_alive,omitempty"`
|
||||
}
|
||||
|
||||
// Entity offsets count Unicode code points; End is exclusive.
|
||||
type Entity struct {
|
||||
Text string `json:"text"`
|
||||
Label string `json:"label"`
|
||||
Start int `json:"start"`
|
||||
End int `json:"end"`
|
||||
Score float32 `json:"score"`
|
||||
}
|
||||
|
||||
type ExtractResponse struct {
|
||||
Model string `json:"model"`
|
||||
Entities []Entity `json:"entities"`
|
||||
TotalDuration time.Duration `json:"total_duration,omitempty"`
|
||||
LoadDuration time.Duration `json:"load_duration,omitempty"`
|
||||
PromptEvalCount int `json:"prompt_eval_count,omitempty"`
|
||||
}
|
||||
|
||||
func (r ExtractRequest) ScoreThreshold() float32 {
|
||||
if r.Threshold == nil {
|
||||
return 0.5
|
||||
}
|
||||
return *r.Threshold
|
||||
}
|
||||
|
||||
// Validate bounds request work before tokenization or model loading.
|
||||
func (r ExtractRequest) Validate() error {
|
||||
if !utf8.ValidString(r.Input) {
|
||||
return fmt.Errorf("input must be valid UTF-8")
|
||||
}
|
||||
if len(r.Input) > 1<<20 {
|
||||
return fmt.Errorf("input exceeds 1 MiB")
|
||||
}
|
||||
if len(r.Labels) == 0 || len(r.Labels) > 128 {
|
||||
return fmt.Errorf("labels must contain between 1 and 128 entity types")
|
||||
}
|
||||
seen := make(map[string]bool, len(r.Labels))
|
||||
for _, label := range r.Labels {
|
||||
if !utf8.ValidString(label) || strings.TrimSpace(label) == "" || len(label) > 256 {
|
||||
return fmt.Errorf("labels must be nonempty UTF-8 strings of at most 256 bytes")
|
||||
}
|
||||
if seen[label] {
|
||||
return fmt.Errorf("duplicate label: %q", label)
|
||||
}
|
||||
seen[label] = true
|
||||
}
|
||||
t := r.ScoreThreshold()
|
||||
if math.IsNaN(float64(t)) || t < 0 || t > 1 {
|
||||
return fmt.Errorf("threshold must be between 0 and 1")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestExtractValidation(t *testing.T) {
|
||||
valid := ExtractRequest{Input: "", Labels: []string{"person"}}
|
||||
if err := valid.Validate(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if valid.ScoreThreshold() != .5 {
|
||||
t.Fatal("default threshold")
|
||||
}
|
||||
for _, r := range []ExtractRequest{
|
||||
{Labels: nil}, {Labels: []string{" "}}, {Labels: []string{"person", "person"}},
|
||||
{Labels: []string{strings.Repeat("x", 257)}}, {Labels: make([]string, 129)},
|
||||
{Input: strings.Repeat("x", 1<<20+1), Labels: valid.Labels},
|
||||
{Input: string([]byte{0xff}), Labels: valid.Labels},
|
||||
} {
|
||||
if err := r.Validate(); err == nil {
|
||||
t.Fatal("accepted invalid request")
|
||||
}
|
||||
}
|
||||
for _, threshold := range []float32{-1, 1.1, float32(math.NaN()), float32(math.Inf(1))} {
|
||||
r := valid
|
||||
r.Threshold = &threshold
|
||||
if err := r.Validate(); err == nil {
|
||||
t.Fatalf("accepted %v", threshold)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -41,6 +41,16 @@ type Classification struct {
|
||||
// quantization from the user's requested type, rejecting requests that are not
|
||||
// allowed for the kind.
|
||||
func Classify(inv Inventory, requested string) (Classification, error) {
|
||||
if inv.Config.Architecture() == "GLiNER" && (requested != "" || detectKind(inv) != SourceFloat) {
|
||||
return Classification{}, fmt.Errorf("GLiNER currently requires unquantized floating-point weights")
|
||||
}
|
||||
if inv.Config.Architecture() == "GLiNER" {
|
||||
for _, name := range sortedTensorNames(inv) {
|
||||
if inv.Tensors[name].Dtype != "F32" {
|
||||
return Classification{}, fmt.Errorf("GLiNER currently requires float32 weights; re-export with scripts/export_gliner.py")
|
||||
}
|
||||
}
|
||||
}
|
||||
requested, err := normalizeRequested(requested)
|
||||
if err != nil {
|
||||
return Classification{}, err
|
||||
|
||||
@@ -108,6 +108,21 @@ func TestClassify(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGLiNERRequiresFloat32(t *testing.T) {
|
||||
cfg := sourceModelConfig{Architectures: []string{"GLiNER"}}
|
||||
inv := newInventory(cfg, map[string]string{"weight": "F32"})
|
||||
if _, err := Classify(inv, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := Classify(inv, "int4"); err == nil {
|
||||
t.Fatal("accepted quantization")
|
||||
}
|
||||
inv = newInventory(cfg, map[string]string{"weight": "F16"})
|
||||
if _, err := Classify(inv, ""); err == nil {
|
||||
t.Fatal("accepted float16")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyErrors(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -97,6 +97,9 @@ func readChatTemplateStrict(modelDir string) (string, error) {
|
||||
}
|
||||
|
||||
func inferSafetensorsCapabilitiesFromConfig(cfg sourceModelConfig, chatTemplate, parserName string) []string {
|
||||
if cfg.Architecture() == "GLiNER" {
|
||||
return []string{"extraction"}
|
||||
}
|
||||
capabilities := []string{"completion"}
|
||||
|
||||
caps := detectCapabilitiesFromConfig(cfg, chatTemplate)
|
||||
|
||||
@@ -101,6 +101,11 @@ func TestInferSafetensorsConfigFamilies(t *testing.T) {
|
||||
wantRenderer string
|
||||
wantCaps []string
|
||||
}{
|
||||
{
|
||||
name: "gliner extraction",
|
||||
config: `{"architectures":["GLiNER"],"model_type":"gliner"}`,
|
||||
wantCaps: []string{"extraction"},
|
||||
},
|
||||
{
|
||||
name: "qwen3",
|
||||
config: `{"architectures":["Qwen3ForCausalLM"]}`,
|
||||
|
||||
@@ -4,8 +4,19 @@ import (
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/mlxrunner"
|
||||
)
|
||||
|
||||
func TestGLiNERArchitectureRegistration(t *testing.T) {
|
||||
if !mlxrunner.SupportsArchitecture("GLiNER") {
|
||||
t.Fatal("GLiNER must be accepted by model import validation")
|
||||
}
|
||||
if mlxrunner.SupportsDraftArchitecture("GLiNER") {
|
||||
t.Fatal("GLiNER must not be accepted as a draft model")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateMLXModelRejectsUnsupportedArchitecture(t *testing.T) {
|
||||
err := validateMLXSource(sourceModelConfig{Architectures: []string{"UnsupportedForCausalLM"}}, false, MLXValidationOptions{})
|
||||
if !errors.Is(err, ErrUnsupportedMLXArchitecture) {
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
package mlxrunner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlx/mlxthread"
|
||||
)
|
||||
|
||||
func (c *Client) Extract(ctx context.Context, req api.ExtractRequest) (*api.ExtractResponse, error) {
|
||||
body, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/v1/extract", c.port), strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
resp, err := c.client.Do(r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
return nil, api.StatusError{StatusCode: resp.StatusCode, ErrorMessage: strings.TrimSpace(string(body))}
|
||||
}
|
||||
var out api.ExtractResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &out, nil
|
||||
}
|
||||
|
||||
func (r *Runner) extractHandler(w http.ResponseWriter, req *http.Request) {
|
||||
if r.Extractor == nil {
|
||||
http.Error(w, "model does not support extraction", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var input api.ExtractRequest
|
||||
if err := json.NewDecoder(http.MaxBytesReader(w, req.Body, 2<<20)).Decode(&input); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if err := input.Validate(); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
out, err := mlxthread.Call(req.Context(), r.mlxThread, func() (*api.ExtractResponse, error) {
|
||||
var out *api.ExtractResponse
|
||||
var err error
|
||||
mlx.Scoped(func() {
|
||||
out, err = r.Extractor.Extract(req.Context(), input)
|
||||
})
|
||||
return out, err
|
||||
})
|
||||
if err != nil {
|
||||
status := http.StatusInternalServerError
|
||||
var se api.StatusError
|
||||
if errors.As(err, &se) {
|
||||
status = se.StatusCode
|
||||
}
|
||||
http.Error(w, err.Error(), status)
|
||||
return
|
||||
}
|
||||
if err := json.NewEncoder(w).Encode(out); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/dflash"
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/gemma4"
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/glimmer"
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/gliner"
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/glm4_moe_lite"
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/laguna"
|
||||
_ "github.com/ollama/ollama/mlxrunner/model/llama"
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/mlx"
|
||||
)
|
||||
|
||||
// Extractor runs a full encoder pass and returns scored spans in the input.
|
||||
// It has no vocabulary projection or autoregressive KV cache.
|
||||
type Extractor interface {
|
||||
LoadWeights(map[string]*mlx.Array) error
|
||||
MaxContextLength() int
|
||||
Extract(context.Context, api.ExtractRequest) (*api.ExtractResponse, error)
|
||||
}
|
||||
|
||||
var extractors = make(map[string]func(*Root) (Extractor, error))
|
||||
|
||||
func RegisterExtractor(arch string, fn func(*Root) (Extractor, error)) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if _, exists := extractors[arch]; exists {
|
||||
panic(fmt.Sprintf("extractor architecture %q already registered", arch))
|
||||
}
|
||||
extractors[arch] = fn
|
||||
}
|
||||
|
||||
// NewExtractor returns nil for a model registered on the generation path.
|
||||
func NewExtractor(root *Root) (Extractor, error) {
|
||||
data, err := root.Manifest.ReadConfig("config.json")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var cfg struct {
|
||||
Architectures []string `json:"architectures"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &cfg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(cfg.Architectures) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
mu.Lock()
|
||||
fn := extractors[cfg.Architectures[0]]
|
||||
mu.Unlock()
|
||||
if fn == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return fn(root)
|
||||
}
|
||||
@@ -0,0 +1,357 @@
|
||||
// Package gliner implements DeBERTa-v3 GLiNER span extraction on MLX.
|
||||
package gliner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"slices"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlxrunner/model"
|
||||
"github.com/ollama/ollama/mlxrunner/tokenizer"
|
||||
)
|
||||
|
||||
func init() { model.RegisterExtractor("GLiNER", newModel) }
|
||||
|
||||
type encoderConfig struct {
|
||||
ModelType string `json:"model_type"`
|
||||
HiddenSize int `json:"hidden_size"`
|
||||
EmbeddingSize int `json:"embedding_size"`
|
||||
IntermediateSize int `json:"intermediate_size"`
|
||||
NumHiddenLayers int `json:"num_hidden_layers"`
|
||||
NumAttentionHeads int `json:"num_attention_heads"`
|
||||
MaxPositions int `json:"max_position_embeddings"`
|
||||
MaxRelativePositions int `json:"max_relative_positions"`
|
||||
PositionBuckets int `json:"position_buckets"`
|
||||
VocabSize int `json:"vocab_size"`
|
||||
LayerNormEps float32 `json:"layer_norm_eps"`
|
||||
HiddenAct string `json:"hidden_act"`
|
||||
RelativeAttention bool `json:"relative_attention"`
|
||||
PositionBiasedInput bool `json:"position_biased_input"`
|
||||
ShareAttKey bool `json:"share_att_key"`
|
||||
NormRelEbd string `json:"norm_rel_ebd"`
|
||||
PosAttType []string `json:"pos_att_type"`
|
||||
TypeVocabSize int `json:"type_vocab_size"`
|
||||
ConvKernelSize int `json:"conv_kernel_size"`
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
Encoder encoderConfig `json:"encoder_config"`
|
||||
HiddenSize int `json:"hidden_size"`
|
||||
MaxWidth int `json:"max_width"`
|
||||
MaxLength int `json:"max_len"`
|
||||
SpanMode string `json:"span_mode"`
|
||||
SubtokenPooling string `json:"subtoken_pooling"`
|
||||
WordsSplitter string `json:"words_splitter_type"`
|
||||
HasRNN bool `json:"has_rnn"`
|
||||
FuseLayers bool `json:"fuse_layers"`
|
||||
EmbedEntToken bool `json:"embed_ent_token"`
|
||||
LabelsEncoder string `json:"labels_encoder"`
|
||||
PostFusionSchema string `json:"post_fusion_schema"`
|
||||
ClassTokenIndex int32 `json:"class_token_index"`
|
||||
EntToken string `json:"ent_token"`
|
||||
SepToken string `json:"sep_token"`
|
||||
}
|
||||
|
||||
type Model struct {
|
||||
Weights map[string]*mlx.Array
|
||||
cfg Config
|
||||
tok *tokenizer.Unigram
|
||||
cls, sep, ent, textSep int32
|
||||
}
|
||||
|
||||
func newModel(root *model.Root) (model.Extractor, error) {
|
||||
if root.QuantType() != "" {
|
||||
return nil, fmt.Errorf("GLiNER requires unquantized weights")
|
||||
}
|
||||
config, err := root.Manifest.ReadConfig("config.json")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tok, err := root.Manifest.ReadConfig("tokenizer.json")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return New(config, tok)
|
||||
}
|
||||
|
||||
// New constructs a model from the self-contained export produced by
|
||||
// scripts/export_gliner.py. Weight loading remains on the MLX thread.
|
||||
func New(config, tok []byte) (*Model, error) {
|
||||
var cfg Config
|
||||
if err := json.Unmarshal(config, &cfg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := cfg.validate(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
t, err := tokenizer.LoadUnigram(tok)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("GLiNER tokenizer: %w", err)
|
||||
}
|
||||
m := &Model{cfg: cfg, tok: t}
|
||||
if t.VocabSize() != cfg.Encoder.VocabSize {
|
||||
return nil, fmt.Errorf("GLiNER tokenizer vocabulary does not match encoder")
|
||||
}
|
||||
for token, target := range map[string]*int32{"[CLS]": &m.cls, "[SEP]": &m.sep, cfg.EntToken: &m.ent, cfg.SepToken: &m.textSep} {
|
||||
id, ok := t.TokenID(token)
|
||||
if !ok || id < 0 || int(id) >= cfg.Encoder.VocabSize {
|
||||
return nil, fmt.Errorf("missing or invalid GLiNER token %q", token)
|
||||
}
|
||||
*target = id
|
||||
}
|
||||
if m.ent != cfg.ClassTokenIndex {
|
||||
return nil, fmt.Errorf("class_token_index does not match tokenizer")
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
func (c Config) validate() error {
|
||||
e := c.Encoder
|
||||
if c.SpanMode != "markerV0" || c.SubtokenPooling != "first" || c.WordsSplitter != "whitespace" || !c.EmbedEntToken || c.FuseLayers || c.LabelsEncoder != "" || c.PostFusionSchema != "" {
|
||||
return fmt.Errorf("unsupported GLiNER variant: requires uni-encoder markerV0 spans, first-subtoken pooling, whitespace splitting, and entity-token embeddings")
|
||||
}
|
||||
if c.HiddenSize <= 0 || c.HiddenSize > 4096 || c.HiddenSize%2 != 0 || c.MaxLength <= 0 || c.MaxLength > 512 || c.MaxWidth <= 0 || c.MaxWidth > 128 || c.EntToken == "" || c.SepToken == "" || c.EntToken == c.SepToken {
|
||||
return fmt.Errorf("invalid GLiNER dimensions or special tokens")
|
||||
}
|
||||
if e.ModelType != "deberta-v2" || !e.RelativeAttention || e.PositionBiasedInput || !e.ShareAttKey || e.NormRelEbd != "layer_norm" || e.HiddenAct != "gelu" || e.TypeVocabSize != 0 || e.ConvKernelSize != 0 || (e.EmbeddingSize != 0 && e.EmbeddingSize != e.HiddenSize) || len(e.PosAttType) != 2 || !slices.Contains(e.PosAttType, "p2c") || !slices.Contains(e.PosAttType, "c2p") {
|
||||
return fmt.Errorf("unsupported GLiNER encoder: requires DeBERTa-v3 with shared relative attention and no convolution or absolute-position embeddings")
|
||||
}
|
||||
if e.HiddenSize <= 0 || e.HiddenSize > 4096 || e.NumAttentionHeads <= 0 || e.HiddenSize%e.NumAttentionHeads != 0 || e.NumHiddenLayers <= 0 || e.NumHiddenLayers > 48 || e.IntermediateSize <= 0 || e.IntermediateSize > 32768 || e.MaxPositions < 4 || e.MaxPositions > 512 || e.PositionBuckets < 4 || e.PositionBuckets > e.MaxPositions || e.VocabSize <= 0 || e.LayerNormEps <= 0 {
|
||||
return fmt.Errorf("invalid GLiNER encoder dimensions")
|
||||
}
|
||||
if e.MaxRelativePositions > 0 && e.MaxRelativePositions <= e.PositionBuckets/2+1 {
|
||||
return fmt.Errorf("invalid max_relative_positions")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
const encoderPrefix = "token_rep_layer.bert_layer.model."
|
||||
const spanPrefix = "span_rep_layer.span_rep_layer."
|
||||
|
||||
func (m *Model) LoadWeights(weights map[string]*mlx.Array) error {
|
||||
c, e := m.cfg, m.cfg.Encoder
|
||||
expected := make(map[string][]int)
|
||||
weight := func(name string, shape ...int) { expected[name] = shape }
|
||||
linear := func(name string, in, out int) { weight(name+".weight", out, in); weight(name+".bias", out) }
|
||||
norm := func(name string, width int) { weight(name+".weight", width); weight(name+".bias", width) }
|
||||
weight(encoderPrefix+"embeddings.word_embeddings.weight", e.VocabSize, e.HiddenSize)
|
||||
norm(encoderPrefix+"embeddings.LayerNorm", e.HiddenSize)
|
||||
weight(encoderPrefix+"encoder.rel_embeddings.weight", 2*e.PositionBuckets, e.HiddenSize)
|
||||
norm(encoderPrefix+"encoder.LayerNorm", e.HiddenSize)
|
||||
for i := range e.NumHiddenLayers {
|
||||
p := fmt.Sprintf("%sencoder.layer.%d.", encoderPrefix, i)
|
||||
for _, proj := range []string{"query_proj", "key_proj", "value_proj"} {
|
||||
linear(p+"attention.self."+proj, e.HiddenSize, e.HiddenSize)
|
||||
}
|
||||
linear(p+"attention.output.dense", e.HiddenSize, e.HiddenSize)
|
||||
norm(p+"attention.output.LayerNorm", e.HiddenSize)
|
||||
linear(p+"intermediate.dense", e.HiddenSize, e.IntermediateSize)
|
||||
linear(p+"output.dense", e.IntermediateSize, e.HiddenSize)
|
||||
norm(p+"output.LayerNorm", e.HiddenSize)
|
||||
}
|
||||
if e.HiddenSize != c.HiddenSize {
|
||||
linear("token_rep_layer.projection", e.HiddenSize, c.HiddenSize)
|
||||
}
|
||||
if c.HasRNN {
|
||||
for _, suffix := range []string{"", "_reverse"} {
|
||||
weight("rnn.lstm.weight_ih_l0"+suffix, 2*c.HiddenSize, c.HiddenSize)
|
||||
weight("rnn.lstm.weight_hh_l0"+suffix, 2*c.HiddenSize, c.HiddenSize/2)
|
||||
weight("rnn.lstm.bias_ih_l0"+suffix, 2*c.HiddenSize)
|
||||
weight("rnn.lstm.bias_hh_l0"+suffix, 2*c.HiddenSize)
|
||||
}
|
||||
}
|
||||
for _, name := range []string{spanPrefix + "project_start", spanPrefix + "project_end", spanPrefix + "out_project", "prompt_rep_layer"} {
|
||||
in := c.HiddenSize
|
||||
if name == spanPrefix+"out_project" {
|
||||
in *= 2
|
||||
}
|
||||
linear(name+".0", in, 4*c.HiddenSize)
|
||||
linear(name+".3", 4*c.HiddenSize, c.HiddenSize)
|
||||
}
|
||||
for name, shape := range expected {
|
||||
a := weights[name]
|
||||
if a == nil {
|
||||
return fmt.Errorf("missing GLiNER tensor %s", name)
|
||||
}
|
||||
if !slices.Equal(a.Dims(), shape) {
|
||||
return fmt.Errorf("GLiNER tensor %s has shape %v, expected %v", name, a.Dims(), shape)
|
||||
}
|
||||
if a.DType() != mlx.DTypeFloat32 {
|
||||
return fmt.Errorf("GLiNER currently requires float32 weights: %s has dtype %v", name, a.DType())
|
||||
}
|
||||
}
|
||||
for name := range weights {
|
||||
if _, ok := expected[name]; !ok {
|
||||
return fmt.Errorf("unsupported GLiNER tensor %s", name)
|
||||
}
|
||||
}
|
||||
m.Weights = weights
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *Model) MaxContextLength() int { return m.cfg.Encoder.MaxPositions }
|
||||
|
||||
func (m *Model) linear(x *mlx.Array, name string) *mlx.Array {
|
||||
return m.Weights[name+".bias"].Addmm(x, m.Weights[name+".weight"].Transpose(1, 0), 1, 1)
|
||||
}
|
||||
|
||||
func (m *Model) norm(x *mlx.Array, name string) *mlx.Array {
|
||||
n := mlx.LayerNorm{Weight: m.Weights[name+".weight"], Bias: m.Weights[name+".bias"]}
|
||||
return n.Forward(x, m.cfg.Encoder.LayerNormEps)
|
||||
}
|
||||
|
||||
func relu(x *mlx.Array) *mlx.Array { return mlx.Maximum(x, mlx.FromValue(float32(0))) }
|
||||
func (m *Model) project(x *mlx.Array, name string) *mlx.Array {
|
||||
return m.linear(relu(m.linear(x, name+".0")), name+".3")
|
||||
}
|
||||
|
||||
// encode returns projected subword embeddings. Each request is unpadded,
|
||||
// so all tokens attend bidirectionally to all other tokens.
|
||||
func (m *Model) encode(ctx context.Context, ids []int32) (*mlx.Array, error) {
|
||||
e := m.cfg.Encoder
|
||||
n, heads, dim := len(ids), e.NumAttentionHeads, e.HiddenSize/e.NumAttentionHeads
|
||||
h := m.Weights[encoderPrefix+"embeddings.word_embeddings.weight"].TakeAxis(mlx.FromValues(ids, n), 0)
|
||||
h = m.norm(h, encoderPrefix+"embeddings.LayerNorm")
|
||||
rel := m.norm(m.Weights[encoderPrefix+"encoder.rel_embeddings.weight"], encoderPrefix+"encoder.LayerNorm")
|
||||
maxPos := e.MaxRelativePositions
|
||||
if maxPos <= 0 {
|
||||
maxPos = e.MaxPositions
|
||||
}
|
||||
c2p, p2c := relativePositions(n, e.PositionBuckets, maxPos)
|
||||
ci, pi := mlx.FromValues(c2p, 1, n, n), mlx.FromValues(p2c, 1, n, n)
|
||||
scale := mlx.FromValue(float32(math.Sqrt(float64(dim * 3))))
|
||||
reshapeHeads := func(x *mlx.Array, length int) *mlx.Array { return x.Reshape(length, heads, dim).Transpose(1, 0, 2) }
|
||||
for i := range e.NumHiddenLayers {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
p := fmt.Sprintf("%sencoder.layer.%d.", encoderPrefix, i)
|
||||
q := reshapeHeads(m.linear(h, p+"attention.self.query_proj"), n)
|
||||
k := reshapeHeads(m.linear(h, p+"attention.self.key_proj"), n)
|
||||
v := reshapeHeads(m.linear(h, p+"attention.self.value_proj"), n)
|
||||
pq := reshapeHeads(m.linear(rel, p+"attention.self.query_proj"), 2*e.PositionBuckets)
|
||||
pk := reshapeHeads(m.linear(rel, p+"attention.self.key_proj"), 2*e.PositionBuckets)
|
||||
scores := q.Matmul(k.Transpose(0, 2, 1).Divide(scale))
|
||||
contentPosition := q.Matmul(pk.Transpose(0, 2, 1)).TakeAlongAxis(ci, 2).Divide(scale)
|
||||
positionContent := k.Matmul(pq.Transpose(0, 2, 1)).TakeAlongAxis(pi, 2).Transpose(0, 2, 1).Divide(scale)
|
||||
scores = scores.Add(contentPosition).Add(positionContent)
|
||||
att := mlx.SoftmaxAxis(scores, -1, true).Matmul(v).Transpose(1, 0, 2).Reshape(n, e.HiddenSize)
|
||||
att = m.norm(h.Add(m.linear(att, p+"attention.output.dense")), p+"attention.output.LayerNorm")
|
||||
h = m.norm(att.Add(m.linear(mlx.GELU(m.linear(att, p+"intermediate.dense")), p+"output.dense")), p+"output.LayerNorm")
|
||||
}
|
||||
if e.HiddenSize != m.cfg.HiddenSize {
|
||||
h = m.linear(h, "token_rep_layer.projection")
|
||||
}
|
||||
return h, nil
|
||||
}
|
||||
|
||||
func relativePositions(n, buckets, maxPosition int) ([]int32, []int32) {
|
||||
a, b := make([]int32, n*n), make([]int32, n*n)
|
||||
mid := buckets / 2
|
||||
for i := range n {
|
||||
for j := range n {
|
||||
d := i - j
|
||||
abs := int(math.Abs(float64(d)))
|
||||
if abs > mid {
|
||||
v := int(math.Ceil(math.Log(float64(abs)/float64(mid))/math.Log(float64(maxPosition-1)/float64(mid))*float64(mid-1))) + mid
|
||||
if d < 0 {
|
||||
d = -v
|
||||
} else {
|
||||
d = v
|
||||
}
|
||||
}
|
||||
a[i*n+j] = int32(min(max(d+buckets, 0), 2*buckets-1))
|
||||
b[i*n+j] = int32(min(max(-d+buckets, 0), 2*buckets-1))
|
||||
}
|
||||
}
|
||||
return a, b
|
||||
}
|
||||
|
||||
func (m *Model) recurrent(ctx context.Context, words *mlx.Array) (*mlx.Array, error) {
|
||||
var directions []*mlx.Array
|
||||
n, width := words.Dim(0), m.cfg.HiddenSize/2
|
||||
for direction, suffix := range []string{"", "_reverse"} {
|
||||
x := m.Weights["rnn.lstm.bias_ih_l0"+suffix].Addmm(words, m.Weights["rnn.lstm.weight_ih_l0"+suffix].Transpose(1, 0), 1, 1)
|
||||
h, c := mlx.Zeros(mlx.DTypeFloat32, 1, width), mlx.Zeros(mlx.DTypeFloat32, 1, width)
|
||||
rows := make([]*mlx.Array, n)
|
||||
for step := range n {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
i := step
|
||||
if direction == 1 {
|
||||
i = n - 1 - step
|
||||
}
|
||||
gates := m.Weights["rnn.lstm.bias_hh_l0"+suffix].Addmm(h, m.Weights["rnn.lstm.weight_hh_l0"+suffix].Transpose(1, 0), 1, 1).Add(x.Slice(mlx.Slice(i, i+1), mlx.Slice()))
|
||||
gate := func(g int) *mlx.Array { return gates.Slice(mlx.Slice(), mlx.Slice(g*width, (g+1)*width)) }
|
||||
c = gate(1).Sigmoid().Multiply(c).Add(gate(0).Sigmoid().Multiply(gate(2).Tanh()))
|
||||
h = gate(3).Sigmoid().Multiply(c.Tanh())
|
||||
rows[i] = h
|
||||
}
|
||||
directions = append(directions, rows[0].Concatenate(0, rows[1:]...))
|
||||
}
|
||||
return directions[0].Concatenate(1, directions[1]), nil
|
||||
}
|
||||
|
||||
func badInput(format string, args ...any) error {
|
||||
return api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: fmt.Sprintf(format, args...)}
|
||||
}
|
||||
|
||||
func (m *Model) Extract(ctx context.Context, req api.ExtractRequest) (*api.ExtractResponse, error) {
|
||||
if err := req.Validate(); err != nil {
|
||||
return nil, badInput("%s", err)
|
||||
}
|
||||
p, err := m.prepare(req.Input, req.Labels)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := &api.ExtractResponse{Entities: []api.Entity{}, PromptEvalCount: len(p.ids)}
|
||||
if len(p.words) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
probabilities := mlx.ScopedArrays(func() []*mlx.Array {
|
||||
var logits *mlx.Array
|
||||
logits, err = m.forward(ctx, p)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return []*mlx.Array{logits.Sigmoid()}
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
probs := probabilities[0].Floats()
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out.Entities = decode(req.Input, req.Labels, p.words, p.spans, probs, req.ScoreThreshold())
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (m *Model) forward(ctx context.Context, p *prepared) (*mlx.Array, error) {
|
||||
h, err := m.encode(ctx, p.ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
words := h.TakeAxis(mlx.FromValues(p.wordIndices, len(p.wordIndices)), 0)
|
||||
if m.cfg.HasRNN {
|
||||
words, err = m.recurrent(ctx, words)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
labels := m.project(h.TakeAxis(mlx.FromValues(p.labelIndices, len(p.labelIndices)), 0), "prompt_rep_layer")
|
||||
starts, ends := make([]int32, len(p.spans)), make([]int32, len(p.spans))
|
||||
for i, span := range p.spans {
|
||||
starts[i], ends[i] = int32(span[0]), int32(span[1])
|
||||
}
|
||||
start := m.project(words, spanPrefix+"project_start").TakeAxis(mlx.FromValues(starts, len(starts)), 0)
|
||||
end := m.project(words, spanPrefix+"project_end").TakeAxis(mlx.FromValues(ends, len(ends)), 0)
|
||||
spans := m.project(relu(start.Concatenate(1, end)), spanPrefix+"out_project")
|
||||
return spans.Matmul(labels.Transpose(1, 0)), nil
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
package gliner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"slices"
|
||||
"testing"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/mlx"
|
||||
"github.com/ollama/ollama/mlx/mlxtest"
|
||||
)
|
||||
|
||||
func readJSON(t *testing.T, path string, v any) {
|
||||
t.Helper()
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err = json.Unmarshal(b, v); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func closeFloats(t *mlxtest.T, got, want []float32, tolerance float64) {
|
||||
t.Helper()
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("length %d, expected %d", len(got), len(want))
|
||||
}
|
||||
var maxError float64
|
||||
for i, v := range got {
|
||||
d := math.Abs(float64(v - want[i]))
|
||||
if math.IsNaN(float64(v)) || d > tolerance {
|
||||
t.Fatalf("value %d = %g, expected %g (error %g)", i, v, want[i], d)
|
||||
}
|
||||
maxError = max(maxError, d)
|
||||
}
|
||||
t.Logf("maximum absolute error: %.8g", maxError)
|
||||
}
|
||||
|
||||
func loadTestWeights(t *mlxtest.T, m *Model, path string) {
|
||||
t.Helper()
|
||||
w := make(map[string]*mlx.Array)
|
||||
for k, v := range mlx.Load(path) {
|
||||
w[k] = v
|
||||
}
|
||||
if err := m.LoadWeights(w); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
weights := mlx.NewScope()
|
||||
weights.Attach(mlx.Collect(m)...)
|
||||
t.Cleanup(weights.Close)
|
||||
}
|
||||
|
||||
// The fixture is generated by the actual PyTorch GLiNER implementation with
|
||||
// small random weights. It exercises logarithmic relative positions, both
|
||||
// LSTM directions, encoder projection, and all span-scoring layers offline.
|
||||
func TestReferenceForward(t *testing.T) {
|
||||
var cfg Config
|
||||
readJSON(t, "testdata/config.json", &cfg)
|
||||
if err := cfg.validate(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var ref struct {
|
||||
IDs []int32 `json:"ids"`
|
||||
WordIndices []int32 `json:"word_indices"`
|
||||
LabelIndices []int32 `json:"label_indices"`
|
||||
Spans [][2]int `json:"spans"`
|
||||
Encoded, Logits []float32
|
||||
}
|
||||
readJSON(t, "testdata/reference.json", &ref)
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
mlx.Scoped(func() {
|
||||
m := &Model{cfg: cfg}
|
||||
loadTestWeights(t, m, "testdata/model.safetensors")
|
||||
h, err := m.encode(context.Background(), ref.IDs)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
closeFloats(t, h.Floats(), ref.Encoded, 2e-5)
|
||||
p := &prepared{ids: ref.IDs, wordIndices: ref.WordIndices, labelIndices: ref.LabelIndices, spans: ref.Spans}
|
||||
logits, err := m.forward(context.Background(), p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
closeFloats(t, logits.Floats(), ref.Logits, 2e-5)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
if _, err = m.forward(ctx, p); err != context.Canceled {
|
||||
t.Fatalf("cancellation: %v", err)
|
||||
}
|
||||
delete(m.Weights, "prompt_rep_layer.3.bias")
|
||||
if err = m.LoadWeights(m.Weights); err == nil {
|
||||
t.Fatal("accepted missing weight")
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
// Run against scripts/export_gliner.py --reference output. The real checkpoint
|
||||
// is deliberately opt-in so normal tests do not download model weights.
|
||||
func TestCheckpointParity(t *testing.T) {
|
||||
dir := os.Getenv("OLLAMA_GLINER_TEST_MODEL")
|
||||
if dir == "" {
|
||||
t.Skip("set OLLAMA_GLINER_TEST_MODEL to an export with --reference")
|
||||
}
|
||||
config, err := os.ReadFile(filepath.Join(dir, "config.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tok, err := os.ReadFile(filepath.Join(dir, "tokenizer.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m, err := New(config, tok)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var cases []struct {
|
||||
Input string
|
||||
Labels []string
|
||||
IDs []int32 `json:"ids"`
|
||||
Encoded []float32
|
||||
Entities []api.Entity
|
||||
}
|
||||
readJSON(t, filepath.Join(dir, "reference.json"), &cases)
|
||||
mlxtest.Run(t, func(t *mlxtest.T) {
|
||||
loadTestWeights(t, m, filepath.Join(dir, "model.safetensors"))
|
||||
for _, c := range cases {
|
||||
mlx.Scoped(func() {
|
||||
t.Log(c.Input)
|
||||
p, err := m.prepare(c.Input, c.Labels)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Equal(p.ids, c.IDs) {
|
||||
t.Fatalf("tokens:\n got %v\nwant %v", p.ids, c.IDs)
|
||||
}
|
||||
h, err := m.encode(context.Background(), p.ids)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
closeFloats(t, h.Floats(), c.Encoded, 5e-4)
|
||||
out, err := m.Extract(context.Background(), api.ExtractRequest{Input: c.Input, Labels: c.Labels})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out.Entities) != len(c.Entities) {
|
||||
t.Fatalf("entities: got %+v, expected %+v", out.Entities, c.Entities)
|
||||
}
|
||||
for i, got := range out.Entities {
|
||||
want := c.Entities[i]
|
||||
if math.Abs(float64(got.Score-want.Score)) > 1e-4 {
|
||||
t.Errorf("score: %v vs %v", got, want)
|
||||
}
|
||||
got.Score = want.Score
|
||||
if got != want {
|
||||
t.Errorf("entity: %+v vs %+v", got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSplitAndDecode(t *testing.T) {
|
||||
if words := splitWords("a\u001cb\u001dc\u001ed\u001fe"); len(words) != 5 {
|
||||
t.Fatalf("Python whitespace: %+v", words)
|
||||
}
|
||||
text := "👋 Élodie met Jean-Luc in São Paulo."
|
||||
words := splitWords(text)
|
||||
var tokens []string
|
||||
for _, w := range words {
|
||||
tokens = append(tokens, w.text)
|
||||
}
|
||||
if !reflect.DeepEqual(tokens, []string{"👋", "Élodie", "met", "Jean-Luc", "in", "São", "Paulo", "."}) {
|
||||
t.Fatal(tokens)
|
||||
}
|
||||
out := decode(text, []string{"person", "city"}, words, [][2]int{{1, 1}, {3, 3}, {5, 5}, {5, 6}}, []float32{.9, .1, .8, .1, .1, .7, .1, .95}, .5)
|
||||
want := []api.Entity{{Text: "Élodie", Label: "person", Start: 2, End: 8, Score: .9}, {Text: "Jean-Luc", Label: "person", Start: 13, End: 21, Score: .8}, {Text: "São Paulo", Label: "city", Start: 25, End: 34, Score: .95}}
|
||||
if !reflect.DeepEqual(out, want) {
|
||||
t.Fatalf("got %+v, want %+v", out, want)
|
||||
}
|
||||
if out := decode(text, []string{"person"}, words, [][2]int{{1, 1}}, []float32{.5}, .5); len(out) != 0 {
|
||||
t.Fatal("threshold is strict")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package gliner
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"slices"
|
||||
"unicode"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
type word struct {
|
||||
text string
|
||||
start, end, byteStart, byteEnd int
|
||||
}
|
||||
type prepared struct {
|
||||
ids, wordIndices, labelIndices []int32
|
||||
words []word
|
||||
spans [][2]int
|
||||
}
|
||||
|
||||
// splitWords matches Python's Unicode \w+(?:[-_]\w+)*|\S and retains
|
||||
// both byte offsets for slicing and code-point offsets for the public API.
|
||||
func splitWords(text string) []word {
|
||||
runes := []rune(text)
|
||||
byteOffsets := make([]int, 0, len(runes)+1)
|
||||
for i := range text {
|
||||
byteOffsets = append(byteOffsets, i)
|
||||
}
|
||||
byteOffsets = append(byteOffsets, len(text))
|
||||
isWord := func(r rune) bool { return unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' }
|
||||
var words []word
|
||||
for i := 0; i < len(runes); {
|
||||
// Python's Unicode whitespace also includes the four ASCII record
|
||||
// separators, which Go's unicode.IsSpace deliberately excludes.
|
||||
if unicode.IsSpace(runes[i]) || (runes[i] >= '\u001c' && runes[i] <= '\u001f') {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
start := i
|
||||
if isWord(runes[i]) {
|
||||
i++
|
||||
for i < len(runes) {
|
||||
if isWord(runes[i]) {
|
||||
i++
|
||||
continue
|
||||
}
|
||||
if runes[i] == '-' && i+1 < len(runes) && isWord(runes[i+1]) {
|
||||
i += 2
|
||||
continue
|
||||
}
|
||||
break
|
||||
}
|
||||
} else {
|
||||
i++
|
||||
}
|
||||
words = append(words, word{text[byteOffsets[start]:byteOffsets[i]], start, i, byteOffsets[start], byteOffsets[i]})
|
||||
}
|
||||
return words
|
||||
}
|
||||
|
||||
func (m *Model) prepare(text string, labels []string) (*prepared, error) {
|
||||
p := &prepared{words: splitWords(text), ids: []int32{m.cls}}
|
||||
if len(p.words) > m.cfg.MaxLength {
|
||||
return nil, badInput("input has %d words; model limit is %d (split the input into smaller chunks)", len(p.words), m.cfg.MaxLength)
|
||||
}
|
||||
for _, label := range labels {
|
||||
if m.tok.HasAddedToken(label) {
|
||||
return nil, badInput("label %q contains a reserved token", label)
|
||||
}
|
||||
p.labelIndices = append(p.labelIndices, int32(len(p.ids)))
|
||||
p.ids = append(p.ids, m.ent)
|
||||
ids := m.tok.EncodeWord(label)
|
||||
if len(ids) == 0 {
|
||||
return nil, badInput("label %q is empty after normalization", label)
|
||||
}
|
||||
p.ids = append(p.ids, ids...)
|
||||
}
|
||||
p.ids = append(p.ids, m.textSep)
|
||||
for i, w := range p.words {
|
||||
ids := m.tok.EncodeWord(w.text)
|
||||
if len(ids) == 0 {
|
||||
return nil, badInput("word at offset %d is empty after normalization", w.start)
|
||||
}
|
||||
p.wordIndices = append(p.wordIndices, int32(len(p.ids)))
|
||||
p.ids = append(p.ids, ids...)
|
||||
for width := 0; width < m.cfg.MaxWidth && i+width < len(p.words); width++ {
|
||||
p.spans = append(p.spans, [2]int{i, i + width})
|
||||
}
|
||||
}
|
||||
p.ids = append(p.ids, m.sep)
|
||||
if len(p.ids) > m.MaxContextLength() {
|
||||
return nil, badInput("input and labels require %d tokens; model limit is %d (split the input or use fewer labels)", len(p.ids), m.MaxContextLength())
|
||||
}
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// decode performs GLiNER's default flat, single-label greedy decoding.
|
||||
func decode(text string, labels []string, words []word, spans [][2]int, scores []float32, threshold float32) []api.Entity {
|
||||
type candidate struct {
|
||||
start, end, label int
|
||||
score float32
|
||||
}
|
||||
var candidates []candidate
|
||||
for i, span := range spans {
|
||||
for j := range labels {
|
||||
score := scores[i*len(labels)+j]
|
||||
if score > threshold {
|
||||
candidates = append(candidates, candidate{span[0], span[1], j, score})
|
||||
}
|
||||
}
|
||||
}
|
||||
slices.SortStableFunc(candidates, func(a, b candidate) int { return cmp.Compare(b.score, a.score) })
|
||||
occupied := make([]bool, len(words))
|
||||
entities := make([]api.Entity, 0)
|
||||
for _, c := range candidates {
|
||||
if slices.Contains(occupied[c.start:c.end+1], true) {
|
||||
continue
|
||||
}
|
||||
for i := c.start; i <= c.end; i++ {
|
||||
occupied[i] = true
|
||||
}
|
||||
a, b := words[c.start], words[c.end]
|
||||
entities = append(entities, api.Entity{Text: text[a.byteStart:b.byteEnd], Label: labels[c.label], Start: a.start, End: b.end, Score: c.score})
|
||||
}
|
||||
slices.SortFunc(entities, func(a, b api.Entity) int { return cmp.Compare(a.Start, b.Start) })
|
||||
return entities
|
||||
}
|
||||
+178
@@ -0,0 +1,178 @@
|
||||
{
|
||||
"return_dict": true,
|
||||
"output_hidden_states": false,
|
||||
"output_attentions": false,
|
||||
"torchscript": false,
|
||||
"torch_dtype": null,
|
||||
"use_bfloat16": false,
|
||||
"tf_legacy_loss": false,
|
||||
"pruned_heads": {},
|
||||
"tie_word_embeddings": true,
|
||||
"chunk_size_feed_forward": 0,
|
||||
"is_encoder_decoder": false,
|
||||
"is_decoder": false,
|
||||
"cross_attention_hidden_size": null,
|
||||
"add_cross_attention": false,
|
||||
"tie_encoder_decoder": false,
|
||||
"max_length": 20,
|
||||
"min_length": 0,
|
||||
"do_sample": false,
|
||||
"early_stopping": false,
|
||||
"num_beams": 1,
|
||||
"num_beam_groups": 1,
|
||||
"diversity_penalty": 0.0,
|
||||
"temperature": 1.0,
|
||||
"top_k": 50,
|
||||
"top_p": 1.0,
|
||||
"typical_p": 1.0,
|
||||
"repetition_penalty": 1.0,
|
||||
"length_penalty": 1.0,
|
||||
"no_repeat_ngram_size": 0,
|
||||
"encoder_no_repeat_ngram_size": 0,
|
||||
"bad_words_ids": null,
|
||||
"num_return_sequences": 1,
|
||||
"output_scores": false,
|
||||
"return_dict_in_generate": false,
|
||||
"forced_bos_token_id": null,
|
||||
"forced_eos_token_id": null,
|
||||
"remove_invalid_values": false,
|
||||
"exponential_decay_length_penalty": null,
|
||||
"suppress_tokens": null,
|
||||
"begin_suppress_tokens": null,
|
||||
"architectures": null,
|
||||
"finetuning_task": null,
|
||||
"id2label": {
|
||||
"0": "LABEL_0",
|
||||
"1": "LABEL_1"
|
||||
},
|
||||
"label2id": {
|
||||
"LABEL_0": 0,
|
||||
"LABEL_1": 1
|
||||
},
|
||||
"tokenizer_class": null,
|
||||
"prefix": null,
|
||||
"bos_token_id": null,
|
||||
"pad_token_id": null,
|
||||
"eos_token_id": null,
|
||||
"sep_token_id": null,
|
||||
"decoder_start_token_id": null,
|
||||
"task_specific_params": null,
|
||||
"problem_type": null,
|
||||
"_name_or_path": "",
|
||||
"_attn_implementation_autoset": false,
|
||||
"transformers_version": "4.51.3",
|
||||
"encoder_config": {
|
||||
"return_dict": true,
|
||||
"output_hidden_states": false,
|
||||
"output_attentions": false,
|
||||
"torchscript": false,
|
||||
"torch_dtype": null,
|
||||
"use_bfloat16": false,
|
||||
"tf_legacy_loss": false,
|
||||
"pruned_heads": {},
|
||||
"tie_word_embeddings": true,
|
||||
"chunk_size_feed_forward": 0,
|
||||
"is_encoder_decoder": false,
|
||||
"is_decoder": false,
|
||||
"cross_attention_hidden_size": null,
|
||||
"add_cross_attention": false,
|
||||
"tie_encoder_decoder": false,
|
||||
"max_length": 20,
|
||||
"min_length": 0,
|
||||
"do_sample": false,
|
||||
"early_stopping": false,
|
||||
"num_beams": 1,
|
||||
"num_beam_groups": 1,
|
||||
"diversity_penalty": 0.0,
|
||||
"temperature": 1.0,
|
||||
"top_k": 50,
|
||||
"top_p": 1.0,
|
||||
"typical_p": 1.0,
|
||||
"repetition_penalty": 1.0,
|
||||
"length_penalty": 1.0,
|
||||
"no_repeat_ngram_size": 0,
|
||||
"encoder_no_repeat_ngram_size": 0,
|
||||
"bad_words_ids": null,
|
||||
"num_return_sequences": 1,
|
||||
"output_scores": false,
|
||||
"return_dict_in_generate": false,
|
||||
"forced_bos_token_id": null,
|
||||
"forced_eos_token_id": null,
|
||||
"remove_invalid_values": false,
|
||||
"exponential_decay_length_penalty": null,
|
||||
"suppress_tokens": null,
|
||||
"begin_suppress_tokens": null,
|
||||
"architectures": null,
|
||||
"finetuning_task": null,
|
||||
"id2label": {
|
||||
"0": "LABEL_0",
|
||||
"1": "LABEL_1"
|
||||
},
|
||||
"label2id": {
|
||||
"LABEL_0": 0,
|
||||
"LABEL_1": 1
|
||||
},
|
||||
"tokenizer_class": null,
|
||||
"prefix": null,
|
||||
"bos_token_id": null,
|
||||
"pad_token_id": 0,
|
||||
"eos_token_id": null,
|
||||
"sep_token_id": null,
|
||||
"decoder_start_token_id": null,
|
||||
"task_specific_params": null,
|
||||
"problem_type": null,
|
||||
"_name_or_path": "",
|
||||
"_attn_implementation_autoset": false,
|
||||
"model_type": "deberta-v2",
|
||||
"position_buckets": 8,
|
||||
"share_att_key": true,
|
||||
"norm_rel_ebd": "layer_norm",
|
||||
"hidden_size": 16,
|
||||
"num_hidden_layers": 2,
|
||||
"num_attention_heads": 4,
|
||||
"intermediate_size": 32,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_dropout_prob": 0.1,
|
||||
"attention_probs_dropout_prob": 0.1,
|
||||
"max_position_embeddings": 32,
|
||||
"type_vocab_size": 0,
|
||||
"initializer_range": 0.02,
|
||||
"relative_attention": true,
|
||||
"max_relative_positions": -1,
|
||||
"position_biased_input": false,
|
||||
"pos_att_type": [
|
||||
"p2c",
|
||||
"c2p"
|
||||
],
|
||||
"vocab_size": 32,
|
||||
"layer_norm_eps": 1e-07,
|
||||
"pooler_hidden_size": 16,
|
||||
"pooler_dropout": 0,
|
||||
"pooler_hidden_act": "gelu",
|
||||
"legacy": true
|
||||
},
|
||||
"labels_encoder_config": null,
|
||||
"model_name": "microsoft/deberta-v3-small",
|
||||
"labels_encoder": null,
|
||||
"name": "span level gliner",
|
||||
"max_width": 3,
|
||||
"hidden_size": 8,
|
||||
"dropout": 0.4,
|
||||
"fine_tune": true,
|
||||
"subtoken_pooling": "first",
|
||||
"span_mode": "markerV0",
|
||||
"post_fusion_schema": "",
|
||||
"num_post_fusion_layers": 1,
|
||||
"vocab_size": -1,
|
||||
"max_neg_type_ratio": 1,
|
||||
"max_types": 25,
|
||||
"max_len": 32,
|
||||
"words_splitter_type": "whitespace",
|
||||
"has_rnn": true,
|
||||
"fuse_layers": false,
|
||||
"class_token_index": 4,
|
||||
"embed_ent_token": true,
|
||||
"ent_token": "<<ENT>>",
|
||||
"sep_token": "<<SEP>>",
|
||||
"model_type": "gliner"
|
||||
}
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
"""Regenerate the small offline fixture with the export script's reference environment."""
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from gliner.config import GLiNERConfig
|
||||
from gliner.modeling.base import SpanModel
|
||||
from safetensors.torch import save_file
|
||||
|
||||
torch.manual_seed(42)
|
||||
torch.set_num_threads(1)
|
||||
out = Path(__file__).parent
|
||||
config = GLiNERConfig(
|
||||
hidden_size=8, max_width=3, max_len=32, class_token_index=4,
|
||||
encoder_config=dict(model_type="deberta-v2", hidden_size=16, num_hidden_layers=2,
|
||||
num_attention_heads=4, intermediate_size=32, max_position_embeddings=32,
|
||||
vocab_size=32, relative_attention=True, position_buckets=8,
|
||||
max_relative_positions=-1, position_biased_input=False, type_vocab_size=0,
|
||||
share_att_key=True, norm_rel_ebd="layer_norm", pos_att_type=["p2c", "c2p"],
|
||||
layer_norm_eps=1e-7, hidden_act="gelu")
|
||||
)
|
||||
model = SpanModel(config, False).eval()
|
||||
ids = torch.tensor([[1, 4, 7, 4, 9, 5] + list(range(10, 28)) + [2]])
|
||||
word_indices = list(range(6, 24))
|
||||
label_indices = [1, 3]
|
||||
spans = [[s, s + k] for s in range(18) for k in range(3) if s + k < 18]
|
||||
with torch.inference_mode():
|
||||
encoded = model.token_rep_layer(ids, torch.ones_like(ids))
|
||||
words = model.rnn(encoded[:, word_indices], torch.ones(1, 18))
|
||||
start = model.span_rep_layer.span_rep_layer.project_start(words)
|
||||
end = model.span_rep_layer.span_rep_layer.project_end(words)
|
||||
joined = torch.cat([start[:, [s for s, e in spans]], end[:, [e for s, e in spans]]], -1).relu()
|
||||
projected = model.span_rep_layer.span_rep_layer.out_project(joined)
|
||||
labels = model.prompt_rep_layer(encoded[:, label_indices])
|
||||
logits = projected @ labels.transpose(-1, -2)
|
||||
save_file(model.state_dict(), str(out / "model.safetensors"))
|
||||
(out / "config.json").write_text(json.dumps(config.to_dict(), indent=2) + "\n")
|
||||
(out / "reference.json").write_text(json.dumps(dict(ids=ids[0].tolist(), word_indices=word_indices,
|
||||
label_indices=label_indices, spans=spans, encoded=encoded.flatten().tolist(),
|
||||
logits=logits.flatten().tolist()), indent=2) + "\n")
|
||||
Binary file not shown.
+565
@@ -0,0 +1,565 @@
|
||||
{
|
||||
"ids": [
|
||||
1,
|
||||
4,
|
||||
7,
|
||||
4,
|
||||
9,
|
||||
5,
|
||||
10,
|
||||
11,
|
||||
12,
|
||||
13,
|
||||
14,
|
||||
15,
|
||||
16,
|
||||
17,
|
||||
18,
|
||||
19,
|
||||
20,
|
||||
21,
|
||||
22,
|
||||
23,
|
||||
24,
|
||||
25,
|
||||
26,
|
||||
27,
|
||||
2
|
||||
],
|
||||
"word_indices": [
|
||||
6,
|
||||
7,
|
||||
8,
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
12,
|
||||
13,
|
||||
14,
|
||||
15,
|
||||
16,
|
||||
17,
|
||||
18,
|
||||
19,
|
||||
20,
|
||||
21,
|
||||
22,
|
||||
23
|
||||
],
|
||||
"label_indices": [
|
||||
1,
|
||||
3
|
||||
],
|
||||
"spans": [
|
||||
[
|
||||
0,
|
||||
0
|
||||
],
|
||||
[
|
||||
0,
|
||||
1
|
||||
],
|
||||
[
|
||||
0,
|
||||
2
|
||||
],
|
||||
[
|
||||
1,
|
||||
1
|
||||
],
|
||||
[
|
||||
1,
|
||||
2
|
||||
],
|
||||
[
|
||||
1,
|
||||
3
|
||||
],
|
||||
[
|
||||
2,
|
||||
2
|
||||
],
|
||||
[
|
||||
2,
|
||||
3
|
||||
],
|
||||
[
|
||||
2,
|
||||
4
|
||||
],
|
||||
[
|
||||
3,
|
||||
3
|
||||
],
|
||||
[
|
||||
3,
|
||||
4
|
||||
],
|
||||
[
|
||||
3,
|
||||
5
|
||||
],
|
||||
[
|
||||
4,
|
||||
4
|
||||
],
|
||||
[
|
||||
4,
|
||||
5
|
||||
],
|
||||
[
|
||||
4,
|
||||
6
|
||||
],
|
||||
[
|
||||
5,
|
||||
5
|
||||
],
|
||||
[
|
||||
5,
|
||||
6
|
||||
],
|
||||
[
|
||||
5,
|
||||
7
|
||||
],
|
||||
[
|
||||
6,
|
||||
6
|
||||
],
|
||||
[
|
||||
6,
|
||||
7
|
||||
],
|
||||
[
|
||||
6,
|
||||
8
|
||||
],
|
||||
[
|
||||
7,
|
||||
7
|
||||
],
|
||||
[
|
||||
7,
|
||||
8
|
||||
],
|
||||
[
|
||||
7,
|
||||
9
|
||||
],
|
||||
[
|
||||
8,
|
||||
8
|
||||
],
|
||||
[
|
||||
8,
|
||||
9
|
||||
],
|
||||
[
|
||||
8,
|
||||
10
|
||||
],
|
||||
[
|
||||
9,
|
||||
9
|
||||
],
|
||||
[
|
||||
9,
|
||||
10
|
||||
],
|
||||
[
|
||||
9,
|
||||
11
|
||||
],
|
||||
[
|
||||
10,
|
||||
10
|
||||
],
|
||||
[
|
||||
10,
|
||||
11
|
||||
],
|
||||
[
|
||||
10,
|
||||
12
|
||||
],
|
||||
[
|
||||
11,
|
||||
11
|
||||
],
|
||||
[
|
||||
11,
|
||||
12
|
||||
],
|
||||
[
|
||||
11,
|
||||
13
|
||||
],
|
||||
[
|
||||
12,
|
||||
12
|
||||
],
|
||||
[
|
||||
12,
|
||||
13
|
||||
],
|
||||
[
|
||||
12,
|
||||
14
|
||||
],
|
||||
[
|
||||
13,
|
||||
13
|
||||
],
|
||||
[
|
||||
13,
|
||||
14
|
||||
],
|
||||
[
|
||||
13,
|
||||
15
|
||||
],
|
||||
[
|
||||
14,
|
||||
14
|
||||
],
|
||||
[
|
||||
14,
|
||||
15
|
||||
],
|
||||
[
|
||||
14,
|
||||
16
|
||||
],
|
||||
[
|
||||
15,
|
||||
15
|
||||
],
|
||||
[
|
||||
15,
|
||||
16
|
||||
],
|
||||
[
|
||||
15,
|
||||
17
|
||||
],
|
||||
[
|
||||
16,
|
||||
16
|
||||
],
|
||||
[
|
||||
16,
|
||||
17
|
||||
],
|
||||
[
|
||||
17,
|
||||
17
|
||||
]
|
||||
],
|
||||
"encoded": [
|
||||
0.13803166151046753,
|
||||
0.9183173179626465,
|
||||
0.05566143989562988,
|
||||
0.11429533362388611,
|
||||
-0.09891319274902344,
|
||||
-1.0339157581329346,
|
||||
-0.2648378312587738,
|
||||
0.19268861413002014,
|
||||
-0.7558888792991638,
|
||||
-0.33493104577064514,
|
||||
0.6075437068939209,
|
||||
-0.76958167552948,
|
||||
0.7482602596282959,
|
||||
-0.984503984451294,
|
||||
0.7428271770477295,
|
||||
-0.6775073409080505,
|
||||
-0.08831340074539185,
|
||||
-0.4241906702518463,
|
||||
-0.7314355373382568,
|
||||
0.6786075830459595,
|
||||
-0.3390544056892395,
|
||||
-0.057516396045684814,
|
||||
0.1924310326576233,
|
||||
0.6061497926712036,
|
||||
-0.7558885812759399,
|
||||
-0.33492812514305115,
|
||||
0.607545018196106,
|
||||
-0.7695777416229248,
|
||||
0.7482624053955078,
|
||||
-0.9845048189163208,
|
||||
0.7428231239318848,
|
||||
-0.6775059103965759,
|
||||
-0.977482259273529,
|
||||
0.09945341944694519,
|
||||
0.3061619699001312,
|
||||
0.011966943740844727,
|
||||
-0.5748891830444336,
|
||||
0.15055745840072632,
|
||||
0.03589993715286255,
|
||||
0.3831486999988556,
|
||||
0.2973214387893677,
|
||||
0.2701311409473419,
|
||||
-1.2943135499954224,
|
||||
-0.9672036170959473,
|
||||
1.0692132711410522,
|
||||
-0.42943495512008667,
|
||||
0.4046257436275482,
|
||||
0.38249942660331726,
|
||||
-0.9524793028831482,
|
||||
-1.4186633825302124,
|
||||
-0.032049745321273804,
|
||||
0.2794858515262604,
|
||||
0.6271830797195435,
|
||||
-0.6069919466972351,
|
||||
0.011258870363235474,
|
||||
-0.6641677618026733,
|
||||
0.1800948977470398,
|
||||
0.27628976106643677,
|
||||
-0.6181120872497559,
|
||||
-0.5388696193695068,
|
||||
-0.45457589626312256,
|
||||
0.2241964042186737,
|
||||
1.196892261505127,
|
||||
0.914732813835144,
|
||||
-0.9574040174484253,
|
||||
0.6381024122238159,
|
||||
0.7895979881286621,
|
||||
-0.716254711151123,
|
||||
-0.9984464645385742,
|
||||
0.09159636497497559,
|
||||
-0.96076899766922,
|
||||
-0.154495507478714,
|
||||
0.10557758808135986,
|
||||
-0.1123606264591217,
|
||||
-0.32288941740989685,
|
||||
-1.3604029417037964,
|
||||
0.6144088506698608,
|
||||
-0.30522650480270386,
|
||||
-0.5696865916252136,
|
||||
-0.2840988039970398,
|
||||
0.34857386350631714,
|
||||
0.48727354407310486,
|
||||
0.7820385694503784,
|
||||
0.5528427362442017,
|
||||
0.16006138920783997,
|
||||
0.09681862592697144,
|
||||
-0.8344127535820007,
|
||||
0.20780569314956665,
|
||||
-0.3221020996570587,
|
||||
1.1946192979812622,
|
||||
0.12379211187362671,
|
||||
0.30996057391166687,
|
||||
-0.9909920692443848,
|
||||
0.2575644850730896,
|
||||
-0.08421467244625092,
|
||||
0.0898049920797348,
|
||||
0.14613556861877441,
|
||||
0.1796131730079651,
|
||||
-1.2461665868759155,
|
||||
-1.3771171569824219,
|
||||
-0.7152254581451416,
|
||||
0.4048614501953125,
|
||||
0.06484302133321762,
|
||||
0.01850736141204834,
|
||||
-0.21829503774642944,
|
||||
-0.8247443437576294,
|
||||
-0.9160616397857666,
|
||||
-0.3780176043510437,
|
||||
-0.3856428861618042,
|
||||
-0.7090206146240234,
|
||||
-0.6045150756835938,
|
||||
-0.5446706414222717,
|
||||
-0.3015395402908325,
|
||||
0.08231842517852783,
|
||||
-0.15918944776058197,
|
||||
-0.40952157974243164,
|
||||
0.4064802825450897,
|
||||
-0.5617918968200684,
|
||||
-0.11836183071136475,
|
||||
-0.06997661292552948,
|
||||
0.3433116674423218,
|
||||
-0.010303735733032227,
|
||||
-0.29681113362312317,
|
||||
0.0468025803565979,
|
||||
-1.5556714534759521,
|
||||
-0.16321563720703125,
|
||||
0.27187541127204895,
|
||||
0.19061493873596191,
|
||||
-0.08666950464248657,
|
||||
-0.7716259956359863,
|
||||
-0.5424197912216187,
|
||||
-0.7344834804534912,
|
||||
-0.8484213352203369,
|
||||
-0.2844278812408447,
|
||||
-0.2613758146762848,
|
||||
0.9237504005432129,
|
||||
-0.7714951038360596,
|
||||
-0.6685887575149536,
|
||||
0.31593775749206543,
|
||||
0.13643702864646912,
|
||||
1.2499351501464844,
|
||||
-0.761989951133728,
|
||||
0.1697283834218979,
|
||||
-0.6484171152114868,
|
||||
0.09897387027740479,
|
||||
0.49846914410591125,
|
||||
0.4258536696434021,
|
||||
0.2018621861934662,
|
||||
0.15373334288597107,
|
||||
-1.371038556098938,
|
||||
-0.09506925940513611,
|
||||
-0.1308162957429886,
|
||||
-0.3378608822822571,
|
||||
-0.025023728609085083,
|
||||
0.8757535219192505,
|
||||
0.051070839166641235,
|
||||
0.009554579854011536,
|
||||
0.7214887738227844,
|
||||
-0.6653670072555542,
|
||||
0.471973180770874,
|
||||
-0.343339204788208,
|
||||
0.3723837435245514,
|
||||
0.1323729157447815,
|
||||
-0.3552990257740021,
|
||||
0.08084544539451599,
|
||||
-0.9168650507926941,
|
||||
0.4032351076602936,
|
||||
-0.08859285712242126,
|
||||
0.9973082542419434,
|
||||
0.4259945750236511,
|
||||
-0.5897647142410278,
|
||||
-1.0792120695114136,
|
||||
-0.4280335605144501,
|
||||
-0.770277202129364,
|
||||
-0.009631067514419556,
|
||||
0.6412805914878845,
|
||||
-0.3160144090652466,
|
||||
0.8820979595184326,
|
||||
0.6573563814163208,
|
||||
-0.10854104161262512,
|
||||
-1.2993168830871582,
|
||||
0.780213475227356,
|
||||
0.11156290769577026,
|
||||
-0.010875821113586426,
|
||||
-0.48892661929130554,
|
||||
-0.3976599872112274,
|
||||
0.23569948971271515,
|
||||
-0.2610653042793274,
|
||||
0.2739239037036896,
|
||||
-0.5665829181671143,
|
||||
0.45181649923324585,
|
||||
-0.7946584820747375,
|
||||
-0.15467257797718048,
|
||||
0.4027685821056366,
|
||||
-0.24998189508914948,
|
||||
-0.7914183139801025,
|
||||
0.310392290353775,
|
||||
0.4627382755279541,
|
||||
0.6305574774742126,
|
||||
0.022901952266693115
|
||||
],
|
||||
"logits": [
|
||||
0.02887692116200924,
|
||||
0.02887662500143051,
|
||||
0.029962953180074692,
|
||||
0.029962651431560516,
|
||||
0.0287387203425169,
|
||||
0.028738420456647873,
|
||||
0.030104132369160652,
|
||||
0.030103830620646477,
|
||||
0.028792517259716988,
|
||||
0.02879221737384796,
|
||||
0.03180895373225212,
|
||||
0.03180864453315735,
|
||||
0.028943020850419998,
|
||||
0.028942719101905823,
|
||||
0.03068733774125576,
|
||||
0.030687034130096436,
|
||||
0.030114680528640747,
|
||||
0.030114375054836273,
|
||||
0.030524827539920807,
|
||||
0.030524522066116333,
|
||||
0.029952164739370346,
|
||||
0.02995186299085617,
|
||||
0.02882479690015316,
|
||||
0.02882450260221958,
|
||||
0.030079608783125877,
|
||||
0.030079307034611702,
|
||||
0.02895224466919899,
|
||||
0.028951944783329964,
|
||||
0.029965810477733612,
|
||||
0.029965512454509735,
|
||||
0.028896866366267204,
|
||||
0.028896566480398178,
|
||||
0.029910428449511528,
|
||||
0.029910128563642502,
|
||||
0.029356351122260094,
|
||||
0.029356054961681366,
|
||||
0.029780646786093712,
|
||||
0.029780343174934387,
|
||||
0.029210662469267845,
|
||||
0.02921035885810852,
|
||||
0.02979191206395626,
|
||||
0.02979160286486149,
|
||||
0.029107315465807915,
|
||||
0.02910701185464859,
|
||||
0.02987082488834858,
|
||||
0.029870517551898956,
|
||||
0.02972847782075405,
|
||||
0.029728174209594727,
|
||||
0.029894860461354256,
|
||||
0.02989455871284008,
|
||||
0.02974036894738674,
|
||||
0.029740065336227417,
|
||||
0.030545493587851524,
|
||||
0.03054518811404705,
|
||||
0.029773803427815437,
|
||||
0.02977350354194641,
|
||||
0.03057892993092537,
|
||||
0.030578630045056343,
|
||||
0.030010119080543518,
|
||||
0.030009811744093895,
|
||||
0.030615832656621933,
|
||||
0.030615532770752907,
|
||||
0.030161835253238678,
|
||||
0.030161531642079353,
|
||||
0.03114919550716877,
|
||||
0.031148886308073997,
|
||||
0.029874255880713463,
|
||||
0.02987395040690899,
|
||||
0.030491715297102928,
|
||||
0.030491413548588753,
|
||||
0.030022364109754562,
|
||||
0.030022060498595238,
|
||||
0.030642934143543243,
|
||||
0.03064263053238392,
|
||||
0.030173592269420624,
|
||||
0.03017328679561615,
|
||||
0.030222628265619278,
|
||||
0.030222326517105103,
|
||||
0.03023531660437584,
|
||||
0.030235014855861664,
|
||||
0.030284350737929344,
|
||||
0.03028404898941517,
|
||||
0.03081112541258335,
|
||||
0.030810818076133728,
|
||||
0.03019564226269722,
|
||||
0.030195333063602448,
|
||||
0.030763816088438034,
|
||||
0.03076351247727871,
|
||||
0.02912849932909012,
|
||||
0.02912820689380169,
|
||||
0.03076179511845112,
|
||||
0.030761489644646645,
|
||||
0.029167894273996353,
|
||||
0.029167594388127327,
|
||||
0.030040323734283447,
|
||||
0.030040020123124123,
|
||||
0.029166216030716896,
|
||||
0.029165921732783318,
|
||||
0.03003864921629429,
|
||||
0.030038349330425262,
|
||||
0.029819760471582413,
|
||||
0.029819458723068237
|
||||
]
|
||||
}
|
||||
@@ -97,12 +97,13 @@ func RegisterDraft(arch string, fn func(root *Root, target Model) (DraftModel, e
|
||||
draftRegistry[arch] = fn
|
||||
}
|
||||
|
||||
// SupportsArchitecture reports whether a target model constructor is registered.
|
||||
// SupportsArchitecture reports whether a generation or extraction model constructor is registered.
|
||||
func SupportsArchitecture(arch string) bool {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
_, ok := registry[arch]
|
||||
return ok
|
||||
_, generation := registry[arch]
|
||||
_, extraction := extractors[arch]
|
||||
return generation || extraction
|
||||
}
|
||||
|
||||
// SupportsDraftArchitecture reports whether a draft model constructor is registered.
|
||||
|
||||
@@ -43,6 +43,7 @@ type Request struct {
|
||||
|
||||
type Runner struct {
|
||||
Model model.Model
|
||||
Extractor model.Extractor
|
||||
weights *mlx.Scope
|
||||
Tokenizer *tokenizer.Tokenizer
|
||||
Requests chan Request
|
||||
@@ -79,6 +80,32 @@ func (r *Runner) loadModel(modelName string) (weights []*mlx.Array, err error) {
|
||||
return nil
|
||||
}
|
||||
|
||||
extractor, e := model.NewExtractor(root)
|
||||
if e != nil {
|
||||
err = e
|
||||
return nil
|
||||
}
|
||||
if extractor != nil {
|
||||
if root.Draft != nil {
|
||||
err = errors.New("extraction models do not support draft models")
|
||||
return nil
|
||||
}
|
||||
tensors, e := loadTensorsFromManifest(root)
|
||||
if e != nil {
|
||||
err = e
|
||||
return nil
|
||||
}
|
||||
if mlx.MetalIsAvailable() {
|
||||
mlx.Eval(slices.Collect(maps.Values(tensors))...)
|
||||
}
|
||||
if err = extractor.LoadWeights(tensors); err != nil {
|
||||
return nil
|
||||
}
|
||||
r.Extractor = extractor
|
||||
r.contextLength = extractor.MaxContextLength()
|
||||
return mlx.Collect(extractor)
|
||||
}
|
||||
|
||||
m, e := model.New(root)
|
||||
if e != nil {
|
||||
err = e
|
||||
|
||||
@@ -87,6 +87,7 @@ func Execute(args []string) error {
|
||||
)
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("POST /v1/extract", runner.extractHandler)
|
||||
mux.HandleFunc("GET /v1/status", func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewEncoder(w).Encode(statusResponse{
|
||||
Status: 0,
|
||||
@@ -118,6 +119,10 @@ func Execute(args []string) error {
|
||||
})
|
||||
|
||||
mux.HandleFunc("POST /v1/completions", func(w http.ResponseWriter, r *http.Request) {
|
||||
if runner.Extractor != nil {
|
||||
http.Error(w, "extraction models do not support completion", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
request := Request{Responses: make(chan CompletionResponse)}
|
||||
|
||||
if err := json.NewDecoder(r.Body).Decode(&request.CompletionRequest); err != nil {
|
||||
@@ -189,6 +194,10 @@ func Execute(args []string) error {
|
||||
})
|
||||
|
||||
mux.HandleFunc("POST /v1/tokenize", func(w http.ResponseWriter, r *http.Request) {
|
||||
if runner.Tokenizer == nil {
|
||||
http.Error(w, "tokenization is not available for this model", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var b bytes.Buffer
|
||||
if _, err := io.Copy(&b, r.Body); err != nil {
|
||||
slog.Error("Failed to read request body", "error", err)
|
||||
|
||||
@@ -0,0 +1,260 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Unigram encodes pre-split words with the DeBERTa SentencePiece tokenizer.
|
||||
// It keeps the checkpoint's precompiled Unicode normalization table rather
|
||||
// than approximating it with a different Unicode normalization form.
|
||||
type Unigram struct {
|
||||
pieces map[string]int32
|
||||
scores []float64
|
||||
added map[string]int32
|
||||
maxPiece int
|
||||
unknown int32
|
||||
unknownScore float64
|
||||
chars []uint32
|
||||
replacements string
|
||||
}
|
||||
|
||||
func LoadUnigram(data []byte) (*Unigram, error) {
|
||||
var raw struct {
|
||||
Model struct {
|
||||
Type string `json:"type"`
|
||||
Vocab []json.RawMessage `json:"vocab"`
|
||||
Unknown int32 `json:"unk_id"`
|
||||
ByteFallback bool `json:"byte_fallback"`
|
||||
} `json:"model"`
|
||||
Added []struct {
|
||||
ID int32 `json:"id"`
|
||||
Content string `json:"content"`
|
||||
} `json:"added_tokens"`
|
||||
Normalizer struct {
|
||||
Type string `json:"type"`
|
||||
Normalizers []json.RawMessage `json:"normalizers"`
|
||||
} `json:"normalizer"`
|
||||
PreTokenizer json.RawMessage `json:"pre_tokenizer"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if raw.Model.Type != "Unigram" || raw.Model.ByteFallback {
|
||||
return nil, fmt.Errorf("expected Unigram tokenizer without byte fallback")
|
||||
}
|
||||
if len(raw.Model.Vocab) == 0 || raw.Model.Unknown < 0 || int(raw.Model.Unknown) >= len(raw.Model.Vocab) {
|
||||
return nil, fmt.Errorf("invalid Unigram vocabulary")
|
||||
}
|
||||
u := &Unigram{pieces: make(map[string]int32), added: make(map[string]int32), unknown: raw.Model.Unknown}
|
||||
for i, entry := range raw.Model.Vocab {
|
||||
var pair []json.RawMessage
|
||||
if err := json.Unmarshal(entry, &pair); err != nil || len(pair) != 2 {
|
||||
return nil, fmt.Errorf("invalid Unigram piece %d", i)
|
||||
}
|
||||
var piece string
|
||||
var score float64
|
||||
if err := json.Unmarshal(pair[0], &piece); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := json.Unmarshal(pair[1], &score); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
u.pieces[piece] = int32(i)
|
||||
u.scores = append(u.scores, score)
|
||||
u.maxPiece = max(u.maxPiece, len(piece))
|
||||
u.unknownScore = min(u.unknownScore, score)
|
||||
}
|
||||
u.unknownScore -= 10
|
||||
for _, a := range raw.Added {
|
||||
u.added[a.Content] = a.ID
|
||||
}
|
||||
// Restrict the accepted pipeline to the one implemented here. Other
|
||||
// tokenizer families must never silently use DeBERTa normalization.
|
||||
if raw.Normalizer.Type != "Sequence" || len(raw.Normalizer.Normalizers) != 3 {
|
||||
return nil, fmt.Errorf("unsupported Unigram normalizer")
|
||||
}
|
||||
for i, n := range raw.Normalizer.Normalizers {
|
||||
var cfg struct {
|
||||
Type string `json:"type"`
|
||||
Left bool `json:"strip_left"`
|
||||
Right bool `json:"strip_right"`
|
||||
Chars string `json:"precompiled_charsmap"`
|
||||
Pattern struct {
|
||||
Regex string `json:"Regex"`
|
||||
} `json:"pattern"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(n, &cfg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch {
|
||||
case i == 0 && cfg.Type == "Strip" && cfg.Left && cfg.Right:
|
||||
case i == 1 && cfg.Type == "Precompiled":
|
||||
b, err := base64.StdEncoding.DecodeString(cfg.Chars)
|
||||
if err != nil || len(b) < 4 {
|
||||
return nil, fmt.Errorf("invalid precompiled normalization table")
|
||||
}
|
||||
n := int(binary.LittleEndian.Uint32(b))
|
||||
if n < 1024 || n%1024 != 0 || n >= len(b)-4 {
|
||||
return nil, fmt.Errorf("invalid precompiled normalization trie")
|
||||
}
|
||||
for j := 4; j < 4+n; j += 4 {
|
||||
u.chars = append(u.chars, binary.LittleEndian.Uint32(b[j:]))
|
||||
}
|
||||
u.replacements = string(b[4+n:])
|
||||
case i == 2 && cfg.Type == "Replace" && cfg.Pattern.Regex == " {2,}" && cfg.Content == " ":
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported Unigram normalizer step %d (%s)", i, cfg.Type)
|
||||
}
|
||||
}
|
||||
var pre struct {
|
||||
Type string `json:"type"`
|
||||
Items []struct {
|
||||
Type string `json:"type"`
|
||||
Replacement string `json:"replacement"`
|
||||
Prepend string `json:"prepend_scheme"`
|
||||
Split bool `json:"split"`
|
||||
} `json:"pretokenizers"`
|
||||
}
|
||||
if err := json.Unmarshal(raw.PreTokenizer, &pre); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if pre.Type != "Sequence" || len(pre.Items) != 1 || pre.Items[0].Type != "Metaspace" || pre.Items[0].Replacement != "▁" || pre.Items[0].Prepend != "always" || !pre.Items[0].Split {
|
||||
return nil, fmt.Errorf("unsupported Unigram pre-tokenizer")
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func (u *Unigram) TokenID(token string) (int32, bool) {
|
||||
if id, ok := u.added[token]; ok {
|
||||
return id, true
|
||||
}
|
||||
id, ok := u.pieces[token]
|
||||
return id, ok
|
||||
}
|
||||
|
||||
func (u *Unigram) VocabSize() int {
|
||||
n := len(u.scores)
|
||||
for _, id := range u.added {
|
||||
n = max(n, int(id)+1)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (u *Unigram) HasAddedToken(s string) bool {
|
||||
for token := range u.added {
|
||||
if strings.Contains(s, token) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// normalize applies the longest matching rule in the SentencePiece Darts
|
||||
// trie. Every lookup is bounded because the table comes from model data.
|
||||
func (u *Unigram) normalize(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
var out strings.Builder
|
||||
offset := func(v uint32) uint32 { return (v >> 10) << ((v & (1 << 9)) >> 6) }
|
||||
for len(s) > 0 {
|
||||
pos := offset(u.chars[0])
|
||||
n, value := 0, 0
|
||||
for i := 0; i < len(s); i++ {
|
||||
pos ^= uint32(s[i])
|
||||
if int(pos) >= len(u.chars) {
|
||||
break
|
||||
}
|
||||
unit := u.chars[pos]
|
||||
if unit&0x800000ff != uint32(s[i]) {
|
||||
break
|
||||
}
|
||||
pos ^= offset(unit)
|
||||
if int(pos) >= len(u.chars) {
|
||||
break
|
||||
}
|
||||
if unit&256 != 0 {
|
||||
n, value = i+1, int(u.chars[pos]&0x7fffffff)
|
||||
}
|
||||
}
|
||||
if n > 0 && value < len(u.replacements) {
|
||||
if end := strings.IndexByte(u.replacements[value:], 0); end >= 0 {
|
||||
out.WriteString(u.replacements[value : value+end])
|
||||
s = s[n:]
|
||||
continue
|
||||
}
|
||||
}
|
||||
_, n = utf8.DecodeRuneInString(s)
|
||||
out.WriteString(s[:n])
|
||||
s = s[n:]
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
|
||||
// EncodeWord returns pieces for one word in an is_split_into_words input.
|
||||
// Special markers are handled by the caller; user text cannot create markers.
|
||||
func (u *Unigram) EncodeWord(s string) []int32 {
|
||||
var ids []int32
|
||||
for word := range strings.SplitSeq(u.normalize(s), " ") {
|
||||
if word == "" {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(word, "▁") {
|
||||
word = "▁" + word
|
||||
}
|
||||
ids = append(ids, u.segment(word)...)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func (u *Unigram) segment(s string) []int32 {
|
||||
type path struct {
|
||||
score float64
|
||||
previous int
|
||||
id int32
|
||||
}
|
||||
best := make([]path, len(s)+1)
|
||||
for i := 1; i < len(best); i++ {
|
||||
best[i].score = math.Inf(-1)
|
||||
}
|
||||
for start := 0; start < len(s); {
|
||||
_, size := utf8.DecodeRuneInString(s[start:])
|
||||
hasRune := false
|
||||
for end := start + 1; end <= min(len(s), start+u.maxPiece); end++ {
|
||||
id, ok := u.pieces[s[start:end]]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if end-start == size {
|
||||
hasRune = true
|
||||
}
|
||||
score := best[start].score + u.scores[id]
|
||||
if score > best[end].score {
|
||||
best[end] = path{score, start, id}
|
||||
}
|
||||
}
|
||||
if !hasRune {
|
||||
end := start + size
|
||||
score := best[start].score + u.unknownScore
|
||||
if score > best[end].score {
|
||||
best[end] = path{score, start, u.unknown}
|
||||
}
|
||||
}
|
||||
start += size
|
||||
}
|
||||
var ids []int32
|
||||
for i := len(s); i > 0; i = best[i].previous {
|
||||
id := best[i].id
|
||||
if id != u.unknown || len(ids) == 0 || ids[len(ids)-1] != id {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
slices.Reverse(ids)
|
||||
return ids
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestUnigramBestPath(t *testing.T) {
|
||||
u := &Unigram{pieces: map[string]int32{"a": 0, "b": 1, "ab": 2, "c": 3}, scores: []float64{-1, -1, -3, -1}, maxPiece: 2, unknown: 4, unknownScore: -10}
|
||||
if got := u.segment("abc"); !slices.Equal(got, []int32{0, 1, 3}) {
|
||||
t.Fatal(got)
|
||||
}
|
||||
if got := u.segment("a😀🤖b"); !slices.Equal(got, []int32{0, 4, 1}) {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrecompiledNormalization(t *testing.T) {
|
||||
u := &Unigram{chars: make([]uint32, 256), replacements: "e\x00"}
|
||||
// A small Darts trie whose sole rule maps UTF-8 é to e.
|
||||
u.chars[0] = 1 << 10
|
||||
u.chars[194] = 0xc3 | 2<<10
|
||||
u.chars[105] = 0xa9 | 2<<10 | 256
|
||||
u.chars[107] = 0x80000000
|
||||
if got := u.normalize(" café! "); got != "cafe!" {
|
||||
t.Fatal(got)
|
||||
}
|
||||
u.chars[0] = 0x7ffffc00 // corrupt offsets must not read past the table
|
||||
if got := u.normalize("é"); got != "é" {
|
||||
t.Fatal(got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnigramRejectsUnsupportedTokenizer(t *testing.T) {
|
||||
for _, data := range []string{`{}`, `{"model":{"type":"BPE"}}`, `{"model":{"type":"Unigram","byte_fallback":true}}`} {
|
||||
if _, err := LoadUnigram([]byte(data)); err == nil {
|
||||
t.Fatal("accepted", data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDebertaTokenizerParity(t *testing.T) {
|
||||
dir := os.Getenv("OLLAMA_GLINER_TEST_MODEL")
|
||||
if dir == "" {
|
||||
t.Skip("set OLLAMA_GLINER_TEST_MODEL to a GLiNER export with --reference")
|
||||
}
|
||||
data, err := os.ReadFile(filepath.Join(dir, "tokenizer.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
u, err := LoadUnigram(data)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err = os.ReadFile(filepath.Join(dir, "tokenizer_reference.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var cases []struct {
|
||||
Word string
|
||||
IDs []int32
|
||||
}
|
||||
if err = json.Unmarshal(data, &cases); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := u.EncodeWord(c.Word); !slices.Equal(got, c.IDs) {
|
||||
t.Errorf("%q: got %v, want %v", c.Word, got, c.IDs)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Export a GLiNER markerV0 checkpoint for Ollama's native MLX runner.
|
||||
|
||||
Reference environment: gliner==0.2.13 transformers==4.51.3 sentencepiece.
|
||||
Python is only needed for export and reference generation, not serving.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from gliner import GLiNER
|
||||
from safetensors.torch import save_file
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("model", help="Hugging Face model ID or local GLiNER directory")
|
||||
parser.add_argument("output", type=Path)
|
||||
parser.add_argument("--revision", help="Hugging Face commit or revision")
|
||||
parser.add_argument("--reference", action="store_true", help="write parity cases for Go integration tests")
|
||||
args = parser.parse_args()
|
||||
args.output.mkdir(parents=True, exist_ok=True)
|
||||
kwargs = {"revision": args.revision} if args.revision else {}
|
||||
model = GLiNER.from_pretrained(args.model, map_location="cpu", **kwargs).float().eval()
|
||||
config = model.config.to_dict()
|
||||
encoder = model.model.token_rep_layer.bert_layer.model.config.to_dict()
|
||||
if (config["span_mode"] != "markerV0" or config.get("labels_encoder")
|
||||
or config.get("fuse_layers") or config.get("post_fusion_schema")
|
||||
or encoder["model_type"] != "deberta-v2"
|
||||
or encoder.get("conv_kernel_size", 0) != 0):
|
||||
raise ValueError("This exporter supports uni-encoder DeBERTa GLiNER markerV0 models only")
|
||||
tokenizer = model.data_processor.transformer_tokenizer
|
||||
if not tokenizer.is_fast:
|
||||
raise ValueError("A fast tokenizer with tokenizer.json is required")
|
||||
config.update(architectures=["GLiNER"], model_type="gliner", encoder_config=encoder, torch_dtype="float32")
|
||||
(args.output / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
||||
tokenizer.save_pretrained(args.output)
|
||||
save_file({k: v.detach().float().contiguous() for k, v in model.model.state_dict().items()},
|
||||
str(args.output / "model.safetensors"))
|
||||
if args.reference:
|
||||
cases = []
|
||||
examples = [
|
||||
("John works at Google in Paris.", ["person", "organization", "location"]),
|
||||
("Élodie met José in São Paulo. 東京 is in Japan. 👋", ["person", "city", "country"]),
|
||||
("Alice joined Acme Corp in New York on January 2, 2024.", ["person", "organization", "location", "date"]),
|
||||
("Dr. Jean-Luc Picard paid $42.50 for café crème.", ["person", "amount", "food"]),
|
||||
("A\u0308nne works at Ｇｏｏｇｌｅ. مرحبا بالعالم 😀", ["person", "organization"]),
|
||||
]
|
||||
for text, labels in examples:
|
||||
words = [w[0] for w in model.data_processor.words_splitter(text)]
|
||||
batch = model.data_processor.tokenize_inputs([words], {label: i + 1 for i, label in enumerate(labels)})
|
||||
with torch.inference_mode():
|
||||
encoded = model.model.token_rep_layer(batch["input_ids"], batch["attention_mask"])[0]
|
||||
entities = model.predict_entities(text, labels)
|
||||
cases.append({"input": text, "labels": labels, "ids": batch["input_ids"][0].tolist(),
|
||||
"encoded": encoded.flatten().tolist(), "entities": entities})
|
||||
(args.output / "reference.json").write_text(json.dumps(cases) + "\n")
|
||||
words = ["John", "Jean-Luc", "Élodie", "A\u0308nne", "Ｇｏｏｇｌｅ", "東京", "😀🤖", "مرحبا",
|
||||
"São Paulo", "person", " a b ", "▁test", "\u200btest", "\u00a0space", "fiancée"]
|
||||
token_cases = [{"word": w, "ids": tokenizer([w], is_split_into_words=True, add_special_tokens=False)["input_ids"]}
|
||||
for w in words]
|
||||
(args.output / "tokenizer_reference.json").write_text(json.dumps(token_cases) + "\n")
|
||||
print(f"Exported to {args.output}. Import with: ollama create --experimental gliner -f <Modelfile>")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,75 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/types/model"
|
||||
)
|
||||
|
||||
type entityExtractor interface {
|
||||
Extract(context.Context, api.ExtractRequest) (*api.ExtractResponse, error)
|
||||
}
|
||||
|
||||
func (s *Server) ExtractHandler(c *gin.Context) {
|
||||
start := time.Now()
|
||||
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 2<<20)
|
||||
var req api.ExtractRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := req.Validate(); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
ref, err := parseAndValidateModelRef(req.Model)
|
||||
if err != nil {
|
||||
writeModelRefParseError(c, err, http.StatusNotFound, fmt.Sprintf("model '%s' not found", req.Model))
|
||||
return
|
||||
}
|
||||
if ref.Source == modelSourceCloud {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "entity extraction requires a local model"})
|
||||
return
|
||||
}
|
||||
name, err := getExistingName(ref.Name)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Model)})
|
||||
return
|
||||
}
|
||||
m, err := GetModel(name.String())
|
||||
if err != nil {
|
||||
handleScheduleError(c, req.Model, err)
|
||||
return
|
||||
}
|
||||
r, _, _, err := s.scheduleRunner(c.Request.Context(), m, []model.Capability{model.CapabilityExtraction}, nil, req.KeepAlive, nil)
|
||||
if err != nil {
|
||||
handleScheduleError(c, req.Model, err)
|
||||
return
|
||||
}
|
||||
loaded := time.Now()
|
||||
extractor, ok := r.(entityExtractor)
|
||||
if !ok {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "model runner does not support extraction"})
|
||||
return
|
||||
}
|
||||
resp, err := extractor.Extract(c.Request.Context(), req)
|
||||
if err != nil {
|
||||
var status api.StatusError
|
||||
if errors.As(err, &status) {
|
||||
c.JSON(status.StatusCode, gin.H{"error": status.ErrorMessage})
|
||||
} else {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
}
|
||||
return
|
||||
}
|
||||
resp.Model = req.Model
|
||||
resp.TotalDuration = time.Since(start)
|
||||
resp.LoadDuration = loaded.Sub(start)
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/ollama/ollama/api"
|
||||
"github.com/ollama/ollama/manifest"
|
||||
"github.com/ollama/ollama/types/model"
|
||||
)
|
||||
|
||||
func TestExtractionModelEndpoints(t *testing.T) {
|
||||
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
||||
config, err := manifest.NewLayer(strings.NewReader(`{"model_format":"safetensors","capabilities":["extraction"],"file_type":"F32"}`), "application/vnd.docker.container.image.v1+json")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
encoder, err := manifest.NewLayer(strings.NewReader(`{"architectures":["GLiNER"],"model_type":"gliner","hidden_size":512,"encoder_config":{"hidden_size":768,"num_hidden_layers":6,"max_position_embeddings":512}}`), "application/vnd.ollama.image.json")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
encoder.Name = "config.json"
|
||||
name := model.ParseName("gliner")
|
||||
if err := manifest.WriteManifest(name, config, []manifest.Layer{encoder}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := GetModelInfo(api.ShowRequest{Model: name.String(), Verbose: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(resp.Capabilities) != 1 || resp.Capabilities[0] != model.CapabilityExtraction {
|
||||
t.Fatalf("capabilities: %v", resp.Capabilities)
|
||||
}
|
||||
if resp.ModelInfo["general.architecture"] != "gliner" || resp.ModelInfo["gliner.context_length"] != 512 || resp.ModelInfo["gliner.block_count"] != 6 {
|
||||
t.Fatalf("model info: %v", resp.ModelInfo)
|
||||
}
|
||||
if resp.Details.QuantizationLevel != "F32" || !strings.Contains(resp.Modelfile, "FROM gliner:latest\n") {
|
||||
t.Fatalf("details: %+v; Modelfile: %s", resp.Details, resp.Modelfile)
|
||||
}
|
||||
// Neither embedding endpoint should try to read a nonexistent GGUF layer
|
||||
// or load a runner for an extraction-only model.
|
||||
s := &Server{}
|
||||
for _, endpoint := range []struct {
|
||||
path string
|
||||
handler gin.HandlerFunc
|
||||
}{
|
||||
{"/api/embed", s.EmbedHandler},
|
||||
{"/api/embeddings", s.EmbeddingsHandler},
|
||||
} {
|
||||
t.Run(endpoint.path, func(t *testing.T) {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, endpoint.path, strings.NewReader(`{"model":"gliner","input":"text","prompt":"text"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
endpoint.handler(c)
|
||||
if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "use /api/extract") {
|
||||
t.Fatalf("%d %s", w.Code, w.Body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractHandlerRejectsInvalidInputBeforeLoading(t *testing.T) {
|
||||
for _, body := range []string{
|
||||
`{`, `{}`, `{"model":"gliner","labels":[]}`,
|
||||
`{"model":"gliner","labels":["person","person"]}`,
|
||||
`{"model":"gliner","labels":["person"],"threshold":2}`,
|
||||
`{"model":"gliner","labels":["person"],"input":["text"]}`,
|
||||
} {
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/extract", strings.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
(&Server{}).ExtractHandler(c)
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("%s: %d %s", body, w.Code, w.Body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractionCapability(t *testing.T) {
|
||||
m := &Model{Config: model.ConfigV2{ModelFormat: "safetensors", Capabilities: []string{"extraction"}}}
|
||||
if err := m.CheckCapabilities(model.CapabilityExtraction); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := m.CheckCapabilities(model.CapabilityCompletion); !errors.Is(err, errCapabilityCompletion) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m = &Model{Config: model.ConfigV2{ModelFormat: "safetensors", Capabilities: []string{"completion"}}}
|
||||
if err := m.CheckCapabilities(model.CapabilityExtraction); !errors.Is(err, errCapabilityExtraction) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -46,6 +46,7 @@ var (
|
||||
errCapabilityVision = errors.New("vision")
|
||||
errCapabilityAudio = errors.New("audio")
|
||||
errCapabilityEmbedding = errors.New("embedding")
|
||||
errCapabilityExtraction = errors.New("extraction")
|
||||
errCapabilityThinking = errors.New("thinking")
|
||||
errCapabilityImage = errors.New("image generation")
|
||||
errInsecureProtocol = errors.New("insecure protocol http")
|
||||
@@ -531,6 +532,7 @@ func (m *Model) CheckCapabilities(want ...model.Capability) error {
|
||||
model.CapabilityVision: errCapabilityVision,
|
||||
model.CapabilityAudio: errCapabilityAudio,
|
||||
model.CapabilityEmbedding: errCapabilityEmbedding,
|
||||
model.CapabilityExtraction: errCapabilityExtraction,
|
||||
model.CapabilityThinking: errCapabilityThinking,
|
||||
model.CapabilityImage: errCapabilityImage,
|
||||
}
|
||||
|
||||
+15
-4
@@ -890,6 +890,11 @@ func (s *Server) EmbedHandler(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if slices.Contains(m.Capabilities(), model.CapabilityExtraction) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "extraction models do not support embeddings; use /api/extract"})
|
||||
return
|
||||
}
|
||||
|
||||
r, m, opts, err := s.scheduleRunner(c.Request.Context(), m, []model.Capability{}, req.Options, req.KeepAlive, nil)
|
||||
if err != nil {
|
||||
handleScheduleError(c, req.Model, err)
|
||||
@@ -1102,6 +1107,11 @@ func (s *Server) EmbeddingsHandler(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if slices.Contains(m.Capabilities(), model.CapabilityExtraction) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "extraction models do not support embeddings; use /api/extract"})
|
||||
return
|
||||
}
|
||||
|
||||
r, m, _, err := s.scheduleRunner(c.Request.Context(), m, []model.Capability{}, req.Options, req.KeepAlive, nil)
|
||||
if err != nil {
|
||||
handleScheduleError(c, req.Model, err)
|
||||
@@ -1457,8 +1467,8 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
|
||||
QuantizationLevel: m.Config.FileType,
|
||||
}
|
||||
|
||||
// For safetensors LLM models, populate details from config.json.
|
||||
if m.Config.ModelFormat == "safetensors" && slices.Contains(m.Config.Capabilities, "completion") {
|
||||
// For safetensors language models, populate details from config.json.
|
||||
if m.Config.ModelFormat == "safetensors" && (slices.Contains(m.Config.Capabilities, "completion") || slices.Contains(m.Config.Capabilities, "extraction")) {
|
||||
if info, err := getSafetensorsLLMInfo(name); err == nil {
|
||||
if arch, ok := info["general.architecture"].(string); ok && arch != "" {
|
||||
modelDetails.Family = arch
|
||||
@@ -1569,8 +1579,8 @@ func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// For safetensors LLM models, populate ModelInfo from config.json.
|
||||
if m.Config.ModelFormat == "safetensors" && slices.Contains(m.Config.Capabilities, "completion") {
|
||||
// Safetensors language models have no GGUF layer to introspect.
|
||||
if m.Config.ModelFormat == "safetensors" && (slices.Contains(m.Config.Capabilities, "completion") || slices.Contains(m.Config.Capabilities, "extraction")) {
|
||||
if info, err := getSafetensorsLLMInfo(name); err == nil {
|
||||
resp.ModelInfo = info
|
||||
}
|
||||
@@ -1927,6 +1937,7 @@ func (s *Server) GenerateRoutes() (http.Handler, error) {
|
||||
r.POST("/api/generate", s.withInferenceRequestLogging("/api/generate", s.GenerateHandler)...)
|
||||
r.POST("/api/chat", s.withInferenceRequestLogging("/api/chat", s.ChatHandler)...)
|
||||
r.POST("/api/embed", s.EmbedHandler)
|
||||
r.POST("/api/extract", s.ExtractHandler)
|
||||
r.POST("/api/embeddings", s.EmbeddingsHandler)
|
||||
|
||||
// Inference (OpenAI compatibility)
|
||||
|
||||
+19
-13
@@ -22,18 +22,19 @@ func canonicalQuantType(quantType string) string {
|
||||
|
||||
// modelConfig represents the HuggingFace config.json structure
|
||||
type modelConfig struct {
|
||||
Architectures []string `json:"architectures"`
|
||||
ModelType string `json:"model_type"`
|
||||
HiddenSize int `json:"hidden_size"`
|
||||
NumHiddenLayers int `json:"num_hidden_layers"`
|
||||
MaxPositionEmbeddings int `json:"max_position_embeddings"`
|
||||
IntermediateSize int `json:"intermediate_size"`
|
||||
NumAttentionHeads int `json:"num_attention_heads"`
|
||||
NumKeyValueHeads int `json:"num_key_value_heads"`
|
||||
VocabSize int `json:"vocab_size"`
|
||||
RMSNormEps float64 `json:"rms_norm_eps"`
|
||||
RopeTheta float64 `json:"rope_theta"`
|
||||
TorchDtype string `json:"torch_dtype"`
|
||||
Architectures []string `json:"architectures"`
|
||||
ModelType string `json:"model_type"`
|
||||
HiddenSize int `json:"hidden_size"`
|
||||
NumHiddenLayers int `json:"num_hidden_layers"`
|
||||
MaxPositionEmbeddings int `json:"max_position_embeddings"`
|
||||
IntermediateSize int `json:"intermediate_size"`
|
||||
NumAttentionHeads int `json:"num_attention_heads"`
|
||||
NumKeyValueHeads int `json:"num_key_value_heads"`
|
||||
VocabSize int `json:"vocab_size"`
|
||||
RMSNormEps float64 `json:"rms_norm_eps"`
|
||||
RopeTheta float64 `json:"rope_theta"`
|
||||
TorchDtype string `json:"torch_dtype"`
|
||||
EncoderConfig *modelConfig `json:"encoder_config"`
|
||||
TextConfig *struct {
|
||||
HiddenSize int `json:"hidden_size"`
|
||||
MaxPositionEmbeddings int `json:"max_position_embeddings"`
|
||||
@@ -41,7 +42,7 @@ type modelConfig struct {
|
||||
} `json:"text_config"`
|
||||
}
|
||||
|
||||
// getSafetensorsLLMInfo extracts model information from safetensors LLM models.
|
||||
// getSafetensorsLLMInfo extracts model information from safetensors language models.
|
||||
// It reads the config.json layer and returns a map compatible with GGML's KV format.
|
||||
func getSafetensorsLLMInfo(name model.Name) (map[string]any, error) {
|
||||
mf, err := manifest.ParseNamedManifest(name)
|
||||
@@ -88,6 +89,11 @@ func buildModelInfo(config modelConfig, totalTensorBytes, tensorCount int64) map
|
||||
arch = strings.TrimSuffix(arch, "forcausallm")
|
||||
arch = strings.TrimSuffix(arch, "forconditionalgeneration")
|
||||
}
|
||||
// GLiNER wraps the transformer dimensions in encoder_config. Keep the
|
||||
// architecture name while reporting the encoder's actual context limit.
|
||||
if arch == "gliner" && config.EncoderConfig != nil {
|
||||
config = *config.EncoderConfig
|
||||
}
|
||||
|
||||
// Use text_config values if they exist (for multimodal models)
|
||||
hiddenSize := config.HiddenSize
|
||||
|
||||
@@ -8,6 +8,7 @@ const (
|
||||
CapabilityInsert = Capability("insert")
|
||||
CapabilityVision = Capability("vision")
|
||||
CapabilityEmbedding = Capability("embedding")
|
||||
CapabilityExtraction = Capability("extraction")
|
||||
CapabilityThinking = Capability("thinking")
|
||||
CapabilityImage = Capability("image")
|
||||
CapabilityAudio = Capability("audio")
|
||||
|
||||
Reference in New Issue
Block a user