fixup! discojs-core/models: add gpt

Author: tharvik

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

@@ -58,7 +58,7 @@ export class GPT extends Model {
     this.model.setWeights(ws.weights)
   }
 
-  private convertCharDataset (dataset: Dataset): Dataset {
+  private convertCharDataset (dataset: Dataset): tf.data.Dataset<{ xs: tf.Tensor2D, ys: tf.Tensor3D }> {
     const batchSize = 4
     const sampleSize = GPT.blockSize + 1
     const chunkSize = sampleSize * batchSize * 2
@@ -80,8 +80,8 @@ export class GPT extends Model {
 
       const buffer = await chunk.buffer()
 
-      const xs = tf.buffer([batchSize, GPT.blockSize], 'int32')
-      const ys = tf.buffer([batchSize, GPT.blockSize, GPT.vocabSize], 'int32')
+      const xs = tf.buffer<tf.Rank.R2, 'int32'>([batchSize, GPT.blockSize], 'int32')
+      const ys = tf.buffer<tf.Rank.R3, 'int32'>([batchSize, GPT.blockSize, GPT.vocabSize], 'int32')
 
       for (let i = 0; i < batchSize; i++) {
         for (let j = 0; j < sampleSize; j++) {