feat: add support for text data & tokenization

discojs/discojs-core/src/dataset/data/data.ts

@@ -1,4 +1,4 @@
-import { Task } from '../..'
+import { tf, Task } from '../..'
 import { Dataset } from '../dataset'
 
 export abstract class Data {
@@ -17,5 +17,22 @@ export abstract class Data {
 
   abstract batch (): Data
 
-  abstract preprocess (): Data
+  preprocess (): Data {
+    return this
+  }
+
+  applyPreprocessing (samples: tf.TensorContainer): tf.TensorContainer {
+    return samples
+  }
+
+  get batchedDataset (): Dataset {
+    const batchSize = this.task.trainingInformation.batchSize
+    return batchSize === undefined
+      ? this.dataset
+      : this.dataset.batch(batchSize)
+  }
+
+  get preprocessedDataset (): Dataset {
+    return this.dataset.map(this.applyPreprocessing)
+  }
 }

discojs/discojs-core/src/dataset/data/image_data.ts

@@ -1,5 +1,5 @@
 import { tf, Task } from '../..'
-import { getPreprocessImage, ImagePreprocessing } from './preprocessing'
+import { ImagePreprocessing } from './preprocessing'
 import { Dataset } from '../dataset'
 import { Data } from './data'
 
@@ -43,16 +43,33 @@ export class ImageData extends Data {
   }
 
   batch (): Data {
-    const batchSize = this.task.trainingInformation.batchSize
-    const newDataset = batchSize === undefined ? this.dataset : this.dataset.batch(batchSize)
-
-    return new ImageData(newDataset, this.task, this.size)
+    return new ImageData(this.batchedDataset, this.task, this.size)
   }
 
   preprocess (): Data {
-    let newDataset = this.dataset
-    const preprocessImage = getPreprocessImage(this.task)
-    newDataset = newDataset.map((x: tf.TensorContainer) => preprocessImage(x))
-    return new ImageData(newDataset, this.task, this.size)
+    return new ImageData(this.preprocessedDataset, this.task, this.size)
+  }
+
+  applyPreprocessing (tensorContainer: tf.TensorContainer): tf.TensorContainer {
+    const params = this.task.trainingInformation
+    // TODO unsafe cast, tfjs does not provide the right interface
+    let { xs, ys } = tensorContainer as {
+      xs: tf.Tensor3D | tf.Tensor4D
+      ys: tf.Tensor1D | number | undefined
+    }
+    if (params.preprocessingFunctions?.includes(ImagePreprocessing.Normalize)) {
+      xs = xs.div(tf.scalar(255))
+    }
+    if (params.preprocessingFunctions?.includes(ImagePreprocessing.Resize) &&
+      params.IMAGE_H !== undefined &&
+      params.IMAGE_W !== undefined) {
+      xs = tf.image.resizeBilinear(xs, [
+        params.IMAGE_H, params.IMAGE_W
+      ])
+    }
+    return {
+      xs,
+      ys
+    }
   }
 }

discojs/discojs-core/src/dataset/data/preprocessing.ts

@@ -1,38 +1,10 @@
-import { tf, Task } from '../..'
-
-type PreprocessImage = (image: tf.TensorContainer) => tf.TensorContainer
-
-export type Preprocessing = ImagePreprocessing
-
-export interface ImageTensorContainer extends tf.TensorContainerObject {
-  xs: tf.Tensor3D | tf.Tensor4D
-  ys: tf.Tensor1D | number | undefined
-}
-
 export enum ImagePreprocessing {
   Normalize = 'normalize',
-  Resize = 'resize'
+  Resize = 'resize',
 }
 
-export function getPreprocessImage (task: Task): PreprocessImage {
-  const preprocessImage: PreprocessImage = (tensorContainer: tf.TensorContainer): tf.TensorContainer => {
-    // TODO unsafe cast, tfjs does not provide the right interface
-    const info = task.trainingInformation
-    let { xs, ys } = tensorContainer as ImageTensorContainer
-    if (info.preprocessingFunctions?.includes(ImagePreprocessing.Normalize)) {
-      xs = xs.div(tf.scalar(255))
-    }
-    if (info.preprocessingFunctions?.includes(ImagePreprocessing.Resize) &&
-      info.IMAGE_H !== undefined &&
-      info.IMAGE_W !== undefined) {
-      xs = tf.image.resizeBilinear(xs, [
-        info.IMAGE_H, info.IMAGE_W
-      ])
-    }
-    return {
-      xs,
-      ys
-    }
-  }
-  return preprocessImage
-}
+export enum TextPreprocessing {}
+
+export enum TabularPreprocessing {}
+
+export type Preprocessing = ImagePreprocessing | TextPreprocessing | TabularPreprocessing

discojs/discojs-core/src/dataset/data/tabular_data.ts

@@ -22,13 +22,6 @@ export class TabularData extends Data {
   }
 
   batch (): Data {
-    const batchSize = this.task.trainingInformation.batchSize
-    const newDataset = batchSize === undefined ? this.dataset : this.dataset.batch(batchSize)
-
-    return new TabularData(newDataset, this.task, this.size)
-  }
-
-  preprocess (): Data {
-    return this
+    return new TabularData(this.batchedDataset, this.task, this.size)
   }
 }

discojs/discojs-core/src/dataset/data/text_data.ts

@@ -0,0 +1,7 @@
+import { Data } from './data'
+
+export class TextData extends Data {
+  batch (): TextData {
+    return new TextData(this.batchedDataset, this.task, this.size)
+  }
+}

discojs/discojs-core/src/dataset/data_loader/tabular_loader.ts

@@ -9,11 +9,8 @@ import { DataLoader, DataConfig } from '../data_loader'
 const BUFFER_SIZE = 1000
 
 export abstract class TabularLoader<Source> extends DataLoader<Source> {
-  private readonly delimiter: string
-
-  constructor (task: Task, delimiter: string) {
+  constructor (task: Task, public readonly delimiter = ',') {
     super(task)
-    this.delimiter = delimiter
   }
 
   /**

discojs/discojs-core/src/dataset/data_loader/text_loader.ts

@@ -0,0 +1,57 @@
+import GPT3Tokenizer from 'gpt3-tokenizer'
+import { List } from 'immutable'
+
+import { Task } from '../..'
+import { DataLoader, DataConfig } from './data_loader'
+import { Dataset } from '../dataset'
+import { DataSplit } from '../data/data_split'
+import { TextData } from '../data/text_data'
+
+const BUFFER_SIZE = 50
+
+type Tokenizer = 'gpt3'
+
+export abstract class TextLoader<Source> extends DataLoader<Source> {
+  public readonly tokenize: (token: string) => string
+
+  constructor (
+    task: Task,
+    tokenizerType?: Tokenizer,
+    public readonly delimiter = '\n'
+  ) {
+    super(task)
+
+    if (tokenizerType === undefined) {
+      tokenizerType = 'gpt3'
+    }
+
+    switch (tokenizerType) {
+      default:
+        this.tokenize = new GPT3Tokenizer({ type: tokenizerType }).bpe
+        break
+    }
+  }
+
+  abstract loadTextDatasetFrom (source: Source): Promise<Dataset>
+
+  async load (source: Source, config?: DataConfig): Promise<Dataset> {
+    const dataset = await this.loadTextDatasetFrom(source)
+    return config?.shuffle ? dataset.shuffle(BUFFER_SIZE) : dataset
+  }
+
+  async loadAll (sources: Source[], config: DataConfig): Promise<DataSplit> {
+    const datasets = await Promise.all(sources.map(async (source) =>
+      await this.load(source, { ...config, shuffle: false })))
+    let dataset = List(datasets).reduce((acc: Dataset, dataset) =>
+      acc.concatenate(dataset))
+    dataset = config?.shuffle ? dataset.shuffle(BUFFER_SIZE) : dataset
+    const data = await TextData.init(
+      dataset,
+      this.task,
+      undefined
+    )
+    return {
+      train: data
+    }
+  }
+}

discojs/discojs-node/src/dataset/data_loader/image_loader.ts

@@ -1,6 +1,6 @@
 import fs from 'fs'
 
-import { tf, data } from '../..'
+import { tf, data } from 'core'
 
 export class NodeImageLoader extends data.ImageLoader<string> {
   async readImageFrom (source: string): Promise<tf.Tensor3D> {

discojs/discojs-node/src/dataset/data_loader/tabular_loader.ts

@@ -1,4 +1,4 @@
-import { tf, data } from '../..'
+import { tf, data } from 'core'
 
 export class NodeTabularLoader extends data.TabularLoader<string> {
   loadTabularDatasetFrom (source: string, csvConfig: Record<string, unknown>): tf.data.CSVDataset {

discojs/discojs-node/src/dataset/data_loader/text_loader.ts

@@ -0,0 +1,30 @@
+import fs from 'node:fs'
+
+import split2 from 'split2'
+
+import { tf } from '../..'
+import { TextLoader } from 'core/dataset/data_loader/text_loader'
+import { Dataset } from 'core/dataset'
+import { DataConfig } from 'core/dataset/data_loader'
+
+export class NodeTextLoader extends TextLoader<string> {
+  async loadTextDatasetFrom (source: string, config?: DataConfig): Promise<Dataset> {
+    const prefix = 'file://'
+    if (source.slice(0, 7) !== prefix) {
+      source = prefix + source
+    }
+    // create stream being read by generator
+    const stream = fs.createReadStream(source, { encoding: 'utf-8' })
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const self = this
+
+    async function * dataGenerator (): AsyncGenerator<tf.TensorContainer> {
+      // TODO @s314cy
+      const withLabels = config?.labels !== undefined
+      stream.pipe(split2())
+      stream.on('data', (data) => yield self.tokenize(data))
+    }
+
+    return tf.data.generator(dataGenerator)
+  }
+}

discojs/discojs-web/src/dataset/data_loader/text_loader.ts

@@ -0,0 +1,9 @@
+import { tf } from '../..'
+import { Dataset } from 'core/dataset'
+import { TextLoader } from 'core/dataset/data_loader/text_loader'
+
+export class WebTextLoader extends TextLoader<File> {
+  async loadTextDatasetFrom (source: File): Promise<Dataset> {
+    return new tf.data.TextLineDataset(new tf.data.FileDataSource(source))
+  }
+}