From 40829fc22ffc8e20ec1e02e9870edee860cd3f3c Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Tue, 11 Nov 2025 01:35:22 +0100 Subject: Return generated tokens --- main.go | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/main.go b/main.go index 2c2e567..82461da 100644 --- a/main.go +++ b/main.go @@ -28,14 +28,16 @@ func main() { prompt := []int64{464, 257, 5025} - if err := generate("scripts/onnx-gpt2/model.onnx", prompt, 5); err != nil { + if out, err := generate("scripts/onnx-gpt2/model.onnx", prompt, 5); err != nil { log.Fatal(err) + } else { + fmt.Printf("\n%v\n", out) } } -func generate(model string, prompt []int64, steps int64) error { +func generate(model string, prompt []int64, steps int64) ([]int64, error) { if len(prompt) == 0 { - return errors.New("empty prompt") + return nil, errors.New("empty prompt") } cacheNames, cacheValues := emptyCache() @@ -52,7 +54,7 @@ func generate(model string, prompt []int64, steps int64) error { _, _, outputs, err := forward(model, token, position, cacheNames, cacheValues) if err != nil { - return err + return nil, err } cacheValues = outputs[1:] @@ -62,11 +64,13 @@ func generate(model string, prompt []int64, steps int64) error { token := prompt[len(prompt)-1] offset := int64(len(prompt)) - 1 + out := make([]int64, steps) + for step := range steps { logits, _, outputs, err := forward(model, token, offset+step, cacheNames, cacheValues) if err != nil { - return err + return nil, err } probs := softmax(logits.GetData()) @@ -81,10 +85,12 @@ func generate(model string, prompt []int64, steps int64) error { token = int64(idx[0]) // choose best token + out[step] = token + cacheValues = outputs[1:] } - return nil + return out, nil } func forward(model string, token int64, position int64, cacheNames []string, cacheValues []ort.Value) (*ort.Tensor[float32], []string, []ort.Value, error) { -- cgit v1.3.1