Verify sequence length

1 files changed, 9 insertions, 4 deletions

diff --git a/gpt2/model.go b/gpt2/model.go
index 3655b8a..79a4407 100644
--- a/gpt2/model.go
+++ b/gpt2/model.go
@@ -14,10 +14,11 @@ import (
 )
 
 const (
-	vocabSize = 50257
-	nLayers   = 12
-	nHeads    = 12
-	headDim   = 64
+	vocabSize  = 50257
+	nLayers    = 12
+	nHeads     = 12
+	headDim    = 64
+	nPositions = 1024
 )
 
 type Model struct {
@@ -122,6 +123,10 @@ func (m *Model) Generate(prompt []int64, steps int64, logits *[][]float32) ([]in
 		return nil, errors.New("empty prompt")
 	}
 
+	if int64(len(prompt))+steps > nPositions {
+		return nil, errors.New("sequence length exceeds context limit")
+	}
+
 	n := int64(len(prompt))
 	r := make([]int64, steps)