Refactor ONNX extraction

1 files changed, 9 insertions, 1 deletions

diff --git a/research/sander/extract.go b/research/sander/extract.go
index c72fe7f..c9a2841 100644
--- a/research/sander/extract.go
+++ b/research/sander/extract.go
@@ -14,7 +14,15 @@ func (e *Experiment) Extract(db *sql.DB) error {
 	var data []float32
 	var shape []int
 
-	if d, s, err := onnx.ExtractInitializer(e.model, "transformer.wte.weight"); err != nil {
+	var model *onnx.Model
+
+	if m, err := onnx.NewModel(e.model); err != nil {
+		model = m
+	} else {
+		return err
+	}
+
+	if d, s, err := model.ExtractInitializer("transformer.wte.weight"); err != nil {
 		return err
 	} else {
 		data = d