Add logit tracing

1 files changed, 9 insertions, 5 deletions

diff --git a/main.go b/main.go
index e1f2cdb..74e7adb 100644
--- a/main.go
+++ b/main.go
@@ -28,14 +28,14 @@ func main() {
 
 	prompt := []int64{464, 2068, 7586, 21831}
 
-	if out, err := generate("scripts/onnx-gpt2/model.onnx", prompt, 5); err != nil {
+	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) ([]int64, error) {
+func generate(model string, prompt []int64, steps int64, logits *[][]float32) ([]int64, error) {
 	if len(prompt) == 0 {
 		return nil, errors.New("empty prompt")
 	}
@@ -55,14 +55,18 @@ func generate(model string, prompt []int64, steps int64) ([]int64, error) {
 			return nil, err
 		}
 
-		logits := outputs[0].(*ort.Tensor[float32]).GetData()
+		l := outputs[0].(*ort.Tensor[float32]).GetData()
 
-		idx, p := topK(softmax(logits), 5)
+		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", logits[t], p[i], t)
+			fmt.Printf("%.4f %.4f [%d]\n", l[t], p[i], t)
 		}
 
 		if step < context-1 {