discojs*: rename .unbatch() to .flat()

epfml · Nov 14, 2024 · cb806c0 · cb806c0
1 parent c477bb3
commit cb806c0
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 9 deletions.
diff --git a/cli/src/benchmark_gpt.ts b/cli/src/benchmark_gpt.ts
@@ -79,7 +79,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
     task.trainingInformation.maxSequenceLength = contextLength
     const dataset = loadText('../datasets/wikitext/wiki.train.tokens')
       .map(text => processing.tokenize(tokenizer, text))
-      .unbatch()
+      .flat()
       .batchWithOverlap(config.blockSize)
 
     const preprocessedDataset = dataset

diff --git a/cli/src/train_gpt.ts b/cli/src/train_gpt.ts
@@ -21,7 +21,7 @@ async function main(): Promise<void> {
 
   const tokenDataset = new Dataset([data])
     .map((text: string) => processing.tokenize(tokenizer, text))
-    .unbatch()
+    .flat()
     .batchWithOverlap(config.blockSize)
     .map((tokens) => [tokens.pop(), tokens.last()] as [List<number>, number])
     .repeat()

diff --git a/discojs/src/dataset/dataset.spec.ts b/discojs/src/dataset/dataset.spec.ts
@@ -155,7 +155,7 @@ describe("dataset", () => {
     const blockSize = 4
 
     const parsed = new Dataset([expectedTokens])
-      .unbatch()
+      .flat()
       .batchWithOverlap(blockSize)
 
     // -1 because the last sequence is dropped as there is no next token label

diff --git a/discojs/src/dataset/dataset.ts b/discojs/src/dataset/dataset.ts
@@ -184,8 +184,8 @@ export class Dataset<T> implements AsyncIterable<T> {
     );
   }
 
-  /** Flatten chunks */
-  unbatch<U>(this: Dataset<Batched<U>>): Dataset<U> {
+  /** Flatten batches/arrays of elements */
+  flat<U>(this: Dataset<Batched<U>>): Dataset<U> {
     return new Dataset(
       async function* (this: Dataset<Batched<U>>) {
         for await (const batch of this) yield* batch;

diff --git a/discojs/src/processing/index.ts b/discojs/src/processing/index.ts
@@ -60,7 +60,7 @@ export async function preprocess<D extends DataType>(
 
       const tokenizer = await models.getTaskTokenizer(t);
       return d.map(text => processing.tokenize(tokenizer, text))
-        .unbatch()
+        .flat()
         .batchWithOverlap(blockSize) 
         .map((tokens) => [tokens.pop(), tokens.last()]) as
           Dataset<DataFormat.ModelEncoded[D]>;
@@ -101,7 +101,7 @@ export async function preprocessWithoutLabel<D extends DataType>(
       const tokenizer = await models.getTaskTokenizer(t);
 
       return d.map(text => processing.tokenize(tokenizer, text))
-        .unbatch()
+        .flat()
         .batch(blockSize)
     }
   }

diff --git a/discojs/src/validator.ts b/discojs/src/validator.ts
@@ -22,7 +22,7 @@ export class Validator<D extends DataType> {
           .zip(batch.map(([_, outputs]) => outputs))
           .map(([inferred, truth]) => inferred === truth),
       )
-      .unbatch();
+      .flat();
 
     for await (const e of results) yield e;
   }
@@ -36,7 +36,7 @@ export class Validator<D extends DataType> {
     )
       .batch(this.task.trainingInformation.batchSize)
       .map((batch) => this.#model.predict(batch))
-      .unbatch();
+      .flat();
 
     const predictions = await processing.postprocess(
       this.task,