diff options
| -rw-r--r-- | gpt2/model.go | 13 |
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) |