From 019a25a6082a8cde372304e34dcd0ae3d5e875ed Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Thu, 13 Nov 2025 19:25:27 +0100 Subject: Refactor into workspace --- go.mod | 5 - go.sum | 2 - go.work | 3 + gpt2/cmd/gpt2/main.go | 27 ++++++ gpt2/go.mod | 5 + gpt2/go.sum | 2 + gpt2/main.go | 232 ++++++++++++++++++++++++++++++++++++++++++++++ gpt2/scripts/conv.py | 15 +++ main.go | 250 -------------------------------------------------- scripts/conv.py | 15 --- 10 files changed, 284 insertions(+), 272 deletions(-) delete mode 100644 go.mod delete mode 100644 go.sum create mode 100644 go.work create mode 100644 gpt2/cmd/gpt2/main.go create mode 100644 gpt2/go.mod create mode 100644 gpt2/go.sum create mode 100644 gpt2/main.go create mode 100644 gpt2/scripts/conv.py delete mode 100644 main.go delete mode 100644 scripts/conv.py diff --git a/go.mod b/go.mod deleted file mode 100644 index 77fe489..0000000 --- a/go.mod +++ /dev/null @@ -1,5 +0,0 @@ -module bpc - -go 1.24 - -require github.com/yalue/onnxruntime_go v1.22.0 diff --git a/go.sum b/go.sum deleted file mode 100644 index f2c6460..0000000 --- a/go.sum +++ /dev/null @@ -1,2 +0,0 @@ -github.com/yalue/onnxruntime_go v1.22.0 h1:SzqOfFRRrLRRAFR5VoSxABjTiQSAi8Y4ETYKrMFK1jk= -github.com/yalue/onnxruntime_go v1.22.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4= diff --git a/go.work b/go.work new file mode 100644 index 0000000..bddb430 --- /dev/null +++ b/go.work @@ -0,0 +1,3 @@ +go 1.24.7 + +use ./gpt2 diff --git a/gpt2/cmd/gpt2/main.go b/gpt2/cmd/gpt2/main.go new file mode 100644 index 0000000..ff84f42 --- /dev/null +++ b/gpt2/cmd/gpt2/main.go @@ -0,0 +1,27 @@ +package main + +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") + + if err := ort.InitializeEnvironment(); 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 { + log.Fatal(err) + } else { + fmt.Printf("\n%v\n", out) + } +} diff --git a/gpt2/go.mod b/gpt2/go.mod new file mode 100644 index 0000000..2ca372e --- /dev/null +++ b/gpt2/go.mod @@ -0,0 +1,5 @@ +module gpt2 + +go 1.24 + +require github.com/yalue/onnxruntime_go v1.22.0 diff --git a/gpt2/go.sum b/gpt2/go.sum new file mode 100644 index 0000000..f2c6460 --- /dev/null +++ b/gpt2/go.sum @@ -0,0 +1,2 @@ +github.com/yalue/onnxruntime_go v1.22.0 h1:SzqOfFRRrLRRAFR5VoSxABjTiQSAi8Y4ETYKrMFK1jk= +github.com/yalue/onnxruntime_go v1.22.0/go.mod h1:b4X26A8pekNb1ACJ58wAXgNKeUCGEAQ9dmACut9Sm/4= diff --git a/gpt2/main.go b/gpt2/main.go new file mode 100644 index 0000000..ee9e19f --- /dev/null +++ b/gpt2/main.go @@ -0,0 +1,232 @@ +package gpt2 + +import ( + "errors" + "fmt" + "log" + "math" + "sort" + + ort "github.com/yalue/onnxruntime_go" +) + +const ( + vocabSize = 50257 + nLayers = 12 + nHeads = 12 + headDim = 64 +) + +func Generate(model string, prompt []int64, steps int64, logits *[][]float32) ([]int64, error) { + if len(prompt) == 0 { + return nil, errors.New("empty prompt") + } + + context := int64(len(prompt)) + + cacheNames, cacheValues := emptyCache() + + token := prompt[0] + + out := make([]int64, 0, steps+1) + + for step := range context + steps { + _, _, outputs, err := forward(model, token, step, cacheNames, cacheValues) + + if err != nil { + return nil, err + } + + l := outputs[0].(*ort.Tensor[float32]).GetData() + + if logits != nil { + *logits = append(*logits, l) + } + + idx, p := topK(softmax(l), 5) + + fmt.Printf("\n%d\n\n", token) + + for i, t := range idx { + fmt.Printf("%.4f %.4f [%d]\n", l[t], p[i], t) + } + + if step < context-1 { + token = prompt[step+1] + } else { + token = int64(idx[0]) // choose best token + out = append(out, token) + } + + cacheValues = outputs[1:] + } + + return out[:steps], nil +} + +func forward(model string, token int64, position int64, cacheNames []string, cacheValues []ort.Value) (*ort.Tensor[float32], []string, []ort.Value, error) { + inputNames, inputs, _ := initInputs(token, position) + outputNames, outputs, logits, _ := initOutputs(position) + + inputNames = append(inputNames, cacheNames...) + inputs = append(inputs, cacheValues...) + + session, err := ort.NewAdvancedSession( + model, + inputNames, + outputNames, + inputs, + outputs, + nil, + ) + + if err != nil { + log.Fatal(err) + } + + defer session.Destroy() + + if err := session.Run(); err != nil { + log.Fatal(err) + } + + return logits, outputNames, outputs, nil +} + +func emptyCache() ([]string, []ort.Value) { + names := make([]string, 0, 2*nLayers) + values := make([]ort.Value, 0, 2*nLayers) + shape := []int64{1, int64(nHeads), 0, int64(headDim)} + + for i := range nLayers { + kName := fmt.Sprintf("past_key_values.%d.key", i) + vName := fmt.Sprintf("past_key_values.%d.value", i) + + kTensor, _ := ort.NewEmptyTensor[float32](shape) + vTensor, _ := ort.NewEmptyTensor[float32](shape) + + names = append(names, kName, vName) + values = append(values, ort.Value(kTensor), ort.Value(vTensor)) + } + + return names, values +} + +func initInputs(token, position int64) ([]string, []ort.Value, error) { + inputNames := []string{"input_ids", "position_ids", "attention_mask"} + + var tokens *ort.Tensor[int64] + var positions *ort.Tensor[int64] + var attentionMask *ort.Tensor[int64] + + if t, err := ort.NewTensor[int64]([]int64{1, 1}, []int64{token}); err != nil { + return nil, nil, err + } else { + tokens = t + } + + if p, err := ort.NewTensor[int64]([]int64{1, 1}, []int64{position}); err != nil { + return nil, nil, err + } else { + positions = p + } + + maskData := make([]int64, position+1) + maskShape := []int64{1, position + 1} + + for i := range maskData { + maskData[i] = 1 + } + + if m, err := ort.NewTensor[int64](maskShape, maskData); err != nil { + return nil, nil, err + } else { + attentionMask = m + } + + inputs := []ort.Value{ort.Value(tokens), ort.Value(positions), ort.Value(attentionMask)} + + return inputNames, inputs, nil +} + +func initOutputs(position int64) ([]string, []ort.Value, *ort.Tensor[float32], error) { + outputNames := make([]string, 0, 1+2*nLayers) + outputValues := make([]ort.Value, 0, 1+2*nLayers) + + logits, err := ort.NewEmptyTensor[float32]([]int64{1, 1, int64(vocabSize)}) + + if err != nil { + return nil, nil, nil, err + } + + outputNames = append(outputNames, "logits") + outputValues = append(outputValues, ort.Value(logits)) + + shape := []int64{1, int64(nHeads), position + 1, int64(headDim)} + + for i := range nLayers { + kName := fmt.Sprintf("present.%d.key", i) + vName := fmt.Sprintf("present.%d.value", i) + + kTensor, _ := ort.NewEmptyTensor[float32](shape) + vTensor, _ := ort.NewEmptyTensor[float32](shape) + + outputNames = append(outputNames, kName, vName) + outputValues = append(outputValues, ort.Value(kTensor), ort.Value(vTensor)) + } + + return outputNames, outputValues, logits, nil +} + +func softmax(logits []float32) []float32 { + m := logits[0] + + for _, v := range logits { + if v > m { + m = v + } + } + + s := float32(0.0) + r := make([]float32, len(logits)) + + for i, v := range logits { + e := float32(math.Exp(float64(v - m))) + + r[i] = e + s += e + } + + for i := range r { + r[i] /= s + } + + return r +} + +func topK(p []float32, k int) ([]int, []float32) { + n := len(p) + + if k > n { + k = n + } + + idx := make([]int, n) + + for i := range idx { + idx[i] = i + } + + sort.Slice(idx, func(i, j int) bool { + return p[idx[i]] > p[idx[j]] + }) + + topIdx := idx[:k] + topP := make([]float32, k) + + for i := 0; i < k; i++ { + topP[i] = p[topIdx[i]] + } + + return topIdx, topP +} diff --git a/gpt2/scripts/conv.py b/gpt2/scripts/conv.py new file mode 100644 index 0000000..cf0cbb8 --- /dev/null +++ b/gpt2/scripts/conv.py @@ -0,0 +1,15 @@ +# /// script +# dependencies = [ +# "torch", +# "transformers", +# "optimum", +# "optimum[onnxruntime]", +# ] +# /// + +from transformers import AutoTokenizer, AutoModelForCausalLM +from optimum.onnxruntime import ORTModelForCausalLM + +model_id = "gpt2" +model = ORTModelForCausalLM.from_pretrained(model_id, export=True, use_cache=True) +model.save_pretrained("onnx-gpt2") \ No newline at end of file diff --git a/main.go b/main.go deleted file mode 100644 index 74e7adb..0000000 --- a/main.go +++ /dev/null @@ -1,250 +0,0 @@ -package main - -import ( - "errors" - "fmt" - "log" - "math" - "sort" - - ort "github.com/yalue/onnxruntime_go" -) - -const ( - vocabSize = 50257 - nLayers = 12 - nHeads = 12 - headDim = 64 -) - -func main() { - ort.SetSharedLibraryPath("lib/onnxruntime-osx-arm64-1.22.0/lib/libonnxruntime.1.22.0.dylib") - - if err := ort.InitializeEnvironment(); err != nil { - log.Fatal(err) - } - - defer ort.DestroyEnvironment() - - prompt := []int64{464, 2068, 7586, 21831} - - if out, err := generate("scripts/onnx-gpt2/model.onnx", prompt, 5, nil); err != nil { - log.Fatal(err) - } else { - fmt.Printf("\n%v\n", out) - } -} - -func generate(model string, prompt []int64, steps int64, logits *[][]float32) ([]int64, error) { - if len(prompt) == 0 { - return nil, errors.New("empty prompt") - } - - context := int64(len(prompt)) - - cacheNames, cacheValues := emptyCache() - - token := prompt[0] - - out := make([]int64, 0, steps+1) - - for step := range context + steps { - _, _, outputs, err := forward(model, token, step, cacheNames, cacheValues) - - if err != nil { - return nil, err - } - - l := outputs[0].(*ort.Tensor[float32]).GetData() - - if logits != nil { - *logits = append(*logits, l) - } - - idx, p := topK(softmax(l), 5) - - fmt.Printf("\n%d\n\n", token) - - for i, t := range idx { - fmt.Printf("%.4f %.4f [%d]\n", l[t], p[i], t) - } - - if step < context-1 { - token = prompt[step+1] - } else { - token = int64(idx[0]) // choose best token - out = append(out, token) - } - - cacheValues = outputs[1:] - } - - return out[:steps], nil -} - -func forward(model string, token int64, position int64, cacheNames []string, cacheValues []ort.Value) (*ort.Tensor[float32], []string, []ort.Value, error) { - inputNames, inputs, _ := initInputs(token, position) - outputNames, outputs, logits, _ := initOutputs(position) - - inputNames = append(inputNames, cacheNames...) - inputs = append(inputs, cacheValues...) - - session, err := ort.NewAdvancedSession( - model, - inputNames, - outputNames, - inputs, - outputs, - nil, - ) - - if err != nil { - log.Fatal(err) - } - - defer session.Destroy() - - if err := session.Run(); err != nil { - log.Fatal(err) - } - - return logits, outputNames, outputs, nil -} - -func emptyCache() ([]string, []ort.Value) { - names := make([]string, 0, 2*nLayers) - values := make([]ort.Value, 0, 2*nLayers) - shape := []int64{1, int64(nHeads), 0, int64(headDim)} - - for i := range nLayers { - kName := fmt.Sprintf("past_key_values.%d.key", i) - vName := fmt.Sprintf("past_key_values.%d.value", i) - - kTensor, _ := ort.NewEmptyTensor[float32](shape) - vTensor, _ := ort.NewEmptyTensor[float32](shape) - - names = append(names, kName, vName) - values = append(values, ort.Value(kTensor), ort.Value(vTensor)) - } - - return names, values -} - -func initInputs(token, position int64) ([]string, []ort.Value, error) { - inputNames := []string{"input_ids", "position_ids", "attention_mask"} - - var tokens *ort.Tensor[int64] - var positions *ort.Tensor[int64] - var attentionMask *ort.Tensor[int64] - - if t, err := ort.NewTensor[int64]([]int64{1, 1}, []int64{token}); err != nil { - return nil, nil, err - } else { - tokens = t - } - - if p, err := ort.NewTensor[int64]([]int64{1, 1}, []int64{position}); err != nil { - return nil, nil, err - } else { - positions = p - } - - maskData := make([]int64, position+1) - maskShape := []int64{1, position + 1} - - for i := range maskData { - maskData[i] = 1 - } - - if m, err := ort.NewTensor[int64](maskShape, maskData); err != nil { - return nil, nil, err - } else { - attentionMask = m - } - - inputs := []ort.Value{ort.Value(tokens), ort.Value(positions), ort.Value(attentionMask)} - - return inputNames, inputs, nil -} - -func initOutputs(position int64) ([]string, []ort.Value, *ort.Tensor[float32], error) { - outputNames := make([]string, 0, 1+2*nLayers) - outputValues := make([]ort.Value, 0, 1+2*nLayers) - - logits, err := ort.NewEmptyTensor[float32]([]int64{1, 1, int64(vocabSize)}) - - if err != nil { - return nil, nil, nil, err - } - - outputNames = append(outputNames, "logits") - outputValues = append(outputValues, ort.Value(logits)) - - shape := []int64{1, int64(nHeads), position + 1, int64(headDim)} - - for i := range nLayers { - kName := fmt.Sprintf("present.%d.key", i) - vName := fmt.Sprintf("present.%d.value", i) - - kTensor, _ := ort.NewEmptyTensor[float32](shape) - vTensor, _ := ort.NewEmptyTensor[float32](shape) - - outputNames = append(outputNames, kName, vName) - outputValues = append(outputValues, ort.Value(kTensor), ort.Value(vTensor)) - } - - return outputNames, outputValues, logits, nil -} - -func softmax(logits []float32) []float32 { - m := logits[0] - - for _, v := range logits { - if v > m { - m = v - } - } - - s := float32(0.0) - r := make([]float32, len(logits)) - - for i, v := range logits { - e := float32(math.Exp(float64(v - m))) - - r[i] = e - s += e - } - - for i := range r { - r[i] /= s - } - - return r -} - -func topK(p []float32, k int) ([]int, []float32) { - n := len(p) - - if k > n { - k = n - } - - idx := make([]int, n) - - for i := range idx { - idx[i] = i - } - - sort.Slice(idx, func(i, j int) bool { - return p[idx[i]] > p[idx[j]] - }) - - topIdx := idx[:k] - topP := make([]float32, k) - - for i := 0; i < k; i++ { - topP[i] = p[topIdx[i]] - } - - return topIdx, topP -} diff --git a/scripts/conv.py b/scripts/conv.py deleted file mode 100644 index cf0cbb8..0000000 --- a/scripts/conv.py +++ /dev/null @@ -1,15 +0,0 @@ -# /// script -# dependencies = [ -# "torch", -# "transformers", -# "optimum", -# "optimum[onnxruntime]", -# ] -# /// - -from transformers import AutoTokenizer, AutoModelForCausalLM -from optimum.onnxruntime import ORTModelForCausalLM - -model_id = "gpt2" -model = ORTModelForCausalLM.from_pretrained(model_id, export=True, use_cache=True) -model.save_pretrained("onnx-gpt2") \ No newline at end of file -- cgit v1.3.1