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)
}