summaryrefslogtreecommitdiff
path: root/gpt2/config.go
blob: 9f1a65f5a1866038438a24e40c9c44de20789e21 (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
26
27
28
29
30
31
32
33
34
35
package gpt2

type Config struct {
	VocabSize    int
	NumLayers    int
	NumHeads     int
	HeadDim      int
	NumPositions int
}

type Options struct {
	WithCache    bool
	WithLogits   bool
	WithLogProbs bool
}

func ConfigDefault() Config {
	return Config{
		VocabSize:    50257,
		NumLayers:    12,
		NumHeads:     12,
		HeadDim:      64,
		NumPositions: 1024,
	}
}

func ConfigLarge() Config {
	return Config{
		VocabSize:    50257,
		NumLayers:    36,
		NumHeads:     20,
		HeadDim:      64,
		NumPositions: 1024,
	}
}