summaryrefslogtreecommitdiff
path: root/gpt2/config.go
blob: be7db3027a415e666452e9919d6b3e70491ab6b2 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
package gpt2

type Config struct {
	vocabSize  int
	nLayers    int
	nHeads     int
	headDim    int
	nPositions int
}

func DefaultConfig() Config {
	return Config{
		vocabSize:  50257,
		nLayers:    12,
		nHeads:     12,
		headDim:    64,
		nPositions: 1024,
	}
}

func (c Config) WithVocabSize(vocabSize int) Config {
	c.vocabSize = vocabSize

	return c
}