Support for GLiNER models

[Simon H.]

What would be needed to support GLiNER models? As far as I understand the Go graphs for XLA are compiled ahead of execution, but for GLiNER models we would need some form of dynamic compilation. There are multiple ways to that, padding to maximum shape (pretty wasteful), XLA dynamic shapes, but that would require upstram support, or multiple compiled graphs which are again not as efficient.

Are there other options? What would be the prefered way? Is it already in the planning or are some of these requirements already met?

[Jan Pfeifer]

hi Simon, using onnx-gomlx it either already supports, or it's close to -- there are a bunch of PRs coming in (in onnx-gomlx, gomlx and go-xla).

AJ is the one that probably can give you a more detailed answer, since he is working on it -- both for AMD64 but also for ARM64 (Linux and Mac) with support for Neon (on the Go backend), as well as a CoreML based backend.

cheers 

Jan

[Simon H.]

Oh nice, could we add an example using a GLiNER model? I'm currently looking into using onnx-community/NuNER_Zero for a project, but there are probably also other smaller models which would lead to a similar example. Other kind of related question: Is there a functionality to convert safetensor models to onnx in the go libraries somewhere?

[Simon H.]

Wait, nevermind, we just added safetensor support, so before converting it I could just use the safetensor model itself.

[Jan Pfeifer]

Yes, it should be straight forward to reimplement the model in GoMLX :)

[AJ Roetker]

FWIW I have done this a couple different ways now on some unmerged branches in onnx-gomlx and gomlx, we'll need to add either dynamic shape support at some part of the pipeline, I'm working on adding it into the gomlx/simplego graph side of things and a bit on the xla side, but also abstracting some of it out during graph compilation time during onnx-gomlx conversion! both work but were definitely some work getting the python and gomlx to match, REBEL is a bit simpler and works using both backend as it is depending on your needs for what it's worth! (I was specifically trying to get the official GLiNER 1 onnx model running and did get it working!)

[Simon H.]

That's great news, thanks. Are there active plans to implement dynamic shape support in the near future? Can I somehow help?
What do you mean by REBEL? Do you mean eg. the Babelscape/rebel models?

[AJ Roetker]

https://github.com/ajroetker/gomlx/tree/dynamic-shapes-v2-005/examples/gliner

That's my current mega branch that I've implemented dynamic shapes in the go backend using a bunch of learnings, I'm sure it's not finished and Jan and I only went over the original design doc and it uses a two pass approach to handle some stuff that I'm not sure Jan will like or not but this is to say that I'm actively trying to get it working for gliner (you can see my not quite working gliner example in that link too).

I also had a different approach that implemented removing the dynamism from the gomlx layer in https://github.com/ajroetker/onnx-gomlx/tree/dynamic-tensor-sizes by using the onnx exports, in the end both approaches will likely be used as dynamic shapes in gomlx would be a win I imagine but also a quicker win would be removing the dynamism at onnx compilation time. I'm specifically targeting multi model support for being able to download models and have them run using whatever runners you want (on your pc/mac or in k8s with tpus) 

https://github.com/antflydb/termite/blob/main/e2e/rebel_test.go
https://github.com/antflydb/termite/blob/main/pkg/termite/lib/ner/rebel.go yeah rebel is another ner model that doesn't use dynamism in it's onnx graphs and so "just works" today with gomlx (I have a pipeline that uses gomlx seq2seq.go with a backend_gomlx.go in the termite repo)

The attention layer stuff your working on is a huge thing, super glad you're working on that. We could also probably split up the working for dyamic shapes in gomlx vs dynamic shape extraction/removal in onnx-gomlx. Happy to hop on a call to talk any of this stuff as Jan knows!

[Jan Pfeifer]

About the dynamic shapes support: while I'm itching to review AJ's code and start integrating it, and have the go backend (`simplego`) implement it, I wanted to cut a new release of go-xla, GoMLX, ONNX-GoMLX, go-huggingface with all the recent updates, before starting it. I'm hoping this new release will include this awesome new attention+generation that is under review :)  

--
You received this message because you are subscribed to the Google Groups "gomlx-discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to gomlx-discus...@googlegroups.com.
To view this discussion visit https://groups.google.com/d/msgid/gomlx-discuss/94180018-c708-479e-808a-0f4c7209eb24n%40googlegroups.com.
For more options, visit https://groups.google.com/d/optout.