onnx/cmd/extract/main.go


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

package main

import (
	"fmt"
	"log"

	"go.jknobloc.com/x/onnx"
	"go.jknobloc.com/x/shelf"
	"go.jknobloc.com/x/tensor"
)

func main() {
	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)
	}

	t := tensor.NewDense[float32](shape, data)

	eot := t.Select(0, 50256).Contiguous()

	fmt.Println(eot)
}