From c51a83cfdb80933049543af41b795483a6032946 Mon Sep 17 00:00:00 2001
From: Jonas Knobloch <jonas.knobloch@t-online.de>
Date: Tue, 14 Jul 2026 19:04:11 +0200
Subject: Refactor ONNX extraction

---
 onnx/cmd/extract/main.go | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

(limited to 'onnx/cmd')

diff --git a/onnx/cmd/extract/main.go b/onnx/cmd/extract/main.go
index 32b8c5d..a2c281d 100644
--- a/onnx/cmd/extract/main.go
+++ b/onnx/cmd/extract/main.go
@@ -10,7 +10,15 @@ import (
 )
 
 func main() {
-	data, shape, err := onnx.ExtractInitializer(shelf.Abs("models/gpt2/model.onnx"), "transformer.wte.weight")
+	var model *onnx.Model
+
+	if m, err := onnx.NewModel(shelf.Abs("models/gpt2/model.onnx")); err != nil {
+		log.Fatal(err)
+	} else {
+		model = m
+	}
+
+	data, shape, err := model.ExtractInitializer("transformer.wte.weight")
 
 	if err != nil {
 		log.Fatal(err)
-- 
cgit v1.3.1