feat: add makeChunkedDataset helper

Author: dmitryusenko

packages/tfjs-node-helpers/src/index.ts

@@ -6,3 +6,4 @@ export * from './feature-engineering/feature-extractor';
 export * from './feature-engineering/feature';
 export * from './feature-engineering/prepare-datasets-for-binary-classification';
 export * from './training/sample';
+export * from './utils/make-chunked-dataset';

packages/tfjs-node-helpers/src/utils/make-chunked-dataset-generator.ts

@@ -0,0 +1,30 @@
+import { Tensor, tensor } from '@tensorflow/tfjs-node';
+import { Sample } from '../training/sample';
+
+export const makeChunkedDatasetGenerator = async function* ({
+  loadChunk,
+  chunkSize
+}: {
+  loadChunk: (skip: number, take: number) => Promise<Array<Sample>>,
+  chunkSize: number
+}): AsyncGenerator<{ xs: Tensor, ys: Tensor }> {
+  let skip = 0;
+  let take = chunkSize;
+
+  while (true) {
+    const samples = await loadChunk(skip, take);
+
+    for (const sample of samples) {
+      yield {
+        xs: tensor(sample.input),
+        ys: tensor(sample.output)
+      };
+    }
+
+    if (samples.length < take) {
+      break;
+    }
+
+    skip += take;
+  }
+};

packages/tfjs-node-helpers/src/utils/make-chunked-dataset.ts

@@ -0,0 +1,22 @@
+import { data, TensorContainer } from '@tensorflow/tfjs-node';
+import { Sample } from '../training/sample';
+import { makeChunkedDatasetGenerator } from '../utils/make-chunked-dataset-generator';
+
+export const makeChunkedDataset = ({
+  loadChunk,
+  chunkSize,
+  batchSize
+}: {
+  loadChunk: (skip: number, take: number) => Promise<Array<Sample>>,
+  chunkSize: number,
+  batchSize: number
+}): data.Dataset<TensorContainer> => {
+  return data
+    .generator(
+      () => makeChunkedDatasetGenerator({
+        loadChunk,
+        chunkSize
+      }) as any
+    )
+    .batch(batchSize);
+};

packages/tfjs-node-helpers/tsconfig.lib.json

@@ -3,7 +3,8 @@
   "compilerOptions": {
     "outDir": "../../dist/out-tsc",
     "declaration": true,
-    "types": []
+    "types": [],
+    "lib": ["es2018"]
   },
   "include": ["**/*.ts"],
   "exclude": ["jest.config.ts", "**/*.spec.ts", "**/*.test.ts"]