Add config struct

1 files changed, 25 insertions, 0 deletions

diff --git a/gpt2/config.go b/gpt2/config.go
new file mode 100644
index 0000000..b7cede5
--- /dev/null
+++ b/gpt2/config.go
@@ -0,0 +1,25 @@
+package gpt2
+
+type Config struct {
+	vocabSize  int
+	nLayers    int
+	nHeads     int
+	headDim    int
+	nPositions int
+}
+
+func NewDefaultConfig() 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
+}