1 files changed, 9 insertions, 9 deletions
diff --git a/gpt2/cmd/gpt2/main.go b/gpt2/cmd/gpt2/main.go
index ff84f42..96dabeb 100644
--- a/gpt2/cmd/gpt2/main.go
+++ b/gpt2/cmd/gpt2/main.go
@@ -4,24 +4,24 @@ import (
 	"fmt"
 	"gpt2"
 	"log"
-
-	ort "github.com/yalue/onnxruntime_go"
 )
 
 func main() {
-	ort.SetSharedLibraryPath("lib/onnxruntime-osx-arm64-1.22.0/lib/libonnxruntime.1.22.0.dylib")
+	prompt := []int64{464, 2068, 7586, 21831}
+
+	m := gpt2.NewModel("models/base/model.onnx")
 
-	if err := ort.InitializeEnvironment(); err != nil {
+	if err := m.Init(); err != nil {
 		log.Fatal(err)
 	}
 
-	defer ort.DestroyEnvironment()
-
-	prompt := []int64{464, 2068, 7586, 21831}
-
-	if out, err := gpt2.Generate("scripts/onnx-gpt2/model.onnx", prompt, 5, nil); err != nil {
+	if out, err := m.Generate(prompt, 5, nil); err != nil {
 		log.Fatal(err)
 	} else {
 		fmt.Printf("\n%v\n", out)
 	}
+
+	if err := m.Destroy(); err != nil {
+		log.Fatal(err)
+	}
 }