discojs-core/models: add gpt

Author: peacefulotter

discojs/discojs-core/src/models/gpt/config.ts

@@ -0,0 +1,77 @@
+type ModelType =
+    | 'gpt2'
+    | 'gpt2-medium'
+    | 'gpt2-large'
+    | 'gpt2-xl'
+    | 'gpt-mini'
+    | 'gpt-micro'
+    | 'gpt-nano'
+
+type ModelSize = {
+    nLayer?: number
+    nHead?: number
+    nEmbd?: number
+}
+
+export type GPTConfig = {
+    lr: number
+    batchSize: number
+    blockSize: number
+    vocabSize: number
+    evaluate?: boolean
+    maxEvalBatches?: number
+    evaluateEvery?: number
+    epochs?: number
+    maxIter?: number
+    weightDecay?: number
+    verbose?: boolean
+    bias?: boolean
+    debug?: boolean
+    dropout?: number
+    residDrop?: number
+    embdDrop?: number
+    tokEmb?: boolean
+    lmHead?: boolean
+    modelType: ModelType
+}
+
+export const DEFAULT_CONFIG: Required<GPTConfig> = {
+    lr: 0.001,
+    weightDecay: 0,
+    batchSize: 2,
+    epochs: 9999,
+    maxIter: 10_000,
+    verbose: false,
+    modelType: 'gpt-nano',
+    evaluate: true,
+    maxEvalBatches: 12,
+    evaluateEvery: 100,
+    blockSize: 128,
+    vocabSize: 50258,
+    bias: true,
+    debug: false,
+    dropout: 0.2,
+    residDrop: 0.2,
+    embdDrop: 0.2,
+    tokEmb: true,
+    lmHead: true,
+}
+
+export const getModelSizes = (modelType: ModelType): Required<ModelSize> => {
+    switch (modelType) {
+        case 'gpt2':
+            return { nLayer: 12, nHead: 12, nEmbd: 768 }
+        case 'gpt2-medium':
+            return { nLayer: 24, nHead: 16, nEmbd: 1024 }
+        case 'gpt2-large':
+            return { nLayer: 36, nHead: 20, nEmbd: 1280 }
+        case 'gpt2-xl':
+            return { nLayer: 48, nHead: 25, nEmbd: 1600 }
+        case 'gpt-mini':
+            return { nLayer: 6, nHead: 6, nEmbd: 192 }
+        case 'gpt-micro':
+            return { nLayer: 4, nHead: 4, nEmbd: 128 }
+        case 'gpt-nano':
+            return { nLayer: 3, nHead: 3, nEmbd: 48 }
+    }
+}

discojs/discojs-core/src/models/gpt/evaluate.ts

@@ -0,0 +1,49 @@
+import tf from '@tensorflow/tfjs'
+
+import { data } from '../..'
+import { GPTConfig } from '.'
+
+export default async function evaluate(
+    model: any,
+    dataset: data.Dataset,
+    config: Required<GPTConfig>
+) {
+    console.log('Evaluating..')
+
+    const iter = await dataset.iterator()
+
+    let total_loss = 0
+    const acc: [number, number] = [0, 0]
+
+    let iteration = 0
+    while (iteration < config.maxEvalBatches) {
+        const next = await iter.next()
+        if (!next) break
+        const { xs, ys } = next.value
+        const logits = model.apply(xs)
+
+        // Loss
+        const loss = tf.losses.softmaxCrossEntropy(ys, logits)
+        const lossVal = await loss.array()
+        total_loss += lossVal as number
+
+        // Accuracy
+        const acc_tensor = tf.metrics.categoricalAccuracy(ys, logits)
+        const acc_sum = acc_tensor.sum()
+        acc[0] += (await acc_sum.array()) as number
+        acc[1] += acc_tensor.shape[0] * (acc_tensor.shape[1] as number)
+
+        tf.dispose([acc_tensor, acc_sum, loss, logits, xs, ys])
+
+        iteration++
+    }
+
+    const loss = total_loss / iteration
+    const pp = 2.71828 ** loss
+
+    return {
+        'val/loss': loss,
+        'val/perplexity': pp,
+        'val/acc': acc[0] / acc[1],
+    }
+}

discojs/discojs-core/src/models/gpt/index.ts

@@ -0,0 +1,4 @@
+export * from './train'
+export * from './optimizers'
+export * from './model'
+export * from './config'