@epfml/discojs 3.0.1-p20250814105822.0 → 3.0.1-p20250924150135.0

package/dist/models/hellaswag.d.ts

@@ -1,6 +1,6 @@
+import type { Tokenizer } from './index.js';
 import { GPT } from './index.js';
-import { PreTrainedTokenizer } from '@xenova/transformers';
-import { ONNXModel } from './onnx.js';
+import type { ONNXModel } from './onnx.js';
 export declare const HELLASWAG_URL = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl";
 /**
  * Represents a single example from the HellaSwag dataset.
@@ -15,7 +15,6 @@ export interface HellaSwagExample {
     label: number;
 }
 export type HellaSwagDataset = HellaSwagExample[];
-type Tokenizer = PreTrainedTokenizer;
 type ModelType = GPT | ONNXModel;
 /**
  * Evaluates the model on a given HellaSwag dataset.

package/dist/models/hellaswag.js

@@ -1,7 +1,6 @@
 import * as tf from '@tensorflow/tfjs';
-import { GPT } from './index.js';
-import { tokenize } from '../processing/text.js';
 import { List } from 'immutable';
+import { GPT } from './index.js';
 export const HELLASWAG_URL = 'https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl';
 // Computes the log likelihood of the input sequence using the tfjs model
 // The input sequence is expected to be a concatenation of the context and the ending
@@ -85,11 +84,11 @@ export async function evaluate(model, tokenizer, dataset, print = true) {
     let correct = 0;
     let total = 0;
     for (const example of dataset) {
-        const endingTokens = example.endings.map(e => tokenize(tokenizer, example.ctx + ' ' + e, {
+        const endingTokens = example.endings.map(e => tokenizer.tokenize(`${example.ctx} ${e}`, {
             truncation: true,
             max_length: 128
         }).toArray());
-        const ctxTokens = tokenize(tokenizer, example.ctx, {
+        const ctxTokens = tokenizer.tokenize(example.ctx, {
             truncation: true,
             max_length: 128
         }).toArray();

package/dist/models/index.d.ts

@@ -1,9 +1,8 @@
 export { Model } from './model.js';
 export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js";
+export { Tokenizer } from "./tokenizer.js";
 export { GPT } from './gpt/index.js';
 export { ONNXModel } from './onnx.js';
 export { GPTConfig } from './gpt/config.js';
-export { evaluate as evaluate_hellaswag } from './hellaswag.js';
+export { evaluate as evaluate_hellaswag, HellaSwagDataset, HellaSwagExample, HELLASWAG_URL, } from "./hellaswag.js";
 export { TFJS } from './tfjs.js';
-export { getTaskTokenizer } from './tokenizer.js';
-export { evaluate, HellaSwagDataset, HellaSwagExample, HELLASWAG_URL } from './hellaswag.js';

package/dist/models/index.js

@@ -1,8 +1,7 @@
 export { Model } from './model.js';
 export { EpochLogs } from "./logs.js";
+export { Tokenizer } from "./tokenizer.js";
 export { GPT } from './gpt/index.js';
 export { ONNXModel } from './onnx.js';
-export { evaluate as evaluate_hellaswag } from './hellaswag.js';
+export { evaluate as evaluate_hellaswag, HELLASWAG_URL, } from "./hellaswag.js";
 export { TFJS } from './tfjs.js';
-export { getTaskTokenizer } from './tokenizer.js';
-export { evaluate, HELLASWAG_URL } from './hellaswag.js';

package/dist/models/tokenizer.d.ts

@@ -1,14 +1,24 @@
-import type { Task } from '../index.js';
-import { PreTrainedTokenizer } from '@xenova/transformers';
-/**
- * A task's tokenizer is initially specified as the tokenizer name, e.g., 'Xenova/gpt2'.
- * The first time the tokenizer is needed, this function initializes the actual tokenizer object
- * and saves it in the task' tokenizer field to be reused in subsequent calls.
- *
- * We are proceeding as such because the task object is sent from the server to the client. Rather than
- * sending complex objects through the network, we simply send the tokenizer name, which is then initialized client-side the
- * first time it is called.
- * @param task the task object specifying which tokenizer to use
- * @returns an initialized tokenizer object
- */
-export declare function getTaskTokenizer(task: Task<'text'>): Promise<PreTrainedTokenizer>;
+import { type PreTrainedTokenizer } from "@xenova/transformers";
+import { List } from "immutable";
+type PaddingSide = "left" | "right";
+type TokenizeConfig = {
+    padding?: undefined;
+    truncation?: false;
+} | (({
+    padding: PaddingSide;
+    truncation?: true;
+} | {
+    truncation: true;
+    padding?: PaddingSide;
+}) & {
+    max_length: number;
+});
+export declare class Tokenizer {
+    #private;
+    readonly name: string;
+    constructor(name: string, to_wrap: PreTrainedTokenizer);
+    static from_pretrained(name: string): Promise<Tokenizer>;
+    tokenize(text: string, config?: TokenizeConfig): List<number>;
+    decode(tokens: number[]): string;
+}
+export {};

package/dist/models/tokenizer.js

@@ -1,26 +1,47 @@
-import { AutoTokenizer, env } from '@xenova/transformers';
-/**
- * A task's tokenizer is initially specified as the tokenizer name, e.g., 'Xenova/gpt2'.
- * The first time the tokenizer is needed, this function initializes the actual tokenizer object
- * and saves it in the task' tokenizer field to be reused in subsequent calls.
- *
- * We are proceeding as such because the task object is sent from the server to the client. Rather than
- * sending complex objects through the network, we simply send the tokenizer name, which is then initialized client-side the
- * first time it is called.
- * @param task the task object specifying which tokenizer to use
- * @returns an initialized tokenizer object
- */
-export async function getTaskTokenizer(task) {
-    let tokenizer = task.trainingInformation.tokenizer;
-    if (tokenizer === undefined)
-        throw Error('No tokenizer specified in the task training information');
-    if (typeof tokenizer == 'string') {
-        // Needs to be false in order to prevent transformers.js from reading the local cache 
+import { AutoTokenizer, env, } from "@xenova/transformers";
+import { List } from "immutable";
+// serializable wrapper around PreTrainedTokenizer
+export class Tokenizer {
+    name;
+    #wrapped;
+    constructor(name, to_wrap) {
+        this.name = name;
+        this.#wrapped = to_wrap;
+    }
+    static async from_pretrained(name) {
+        // Needs to be false in order to prevent transformers.js from reading the local cache
         // and triggering an error when running in the browser
         // Reference: https://medium.com/@GenerationAI/transformers-js-onnx-runtime-webgpu-46c3e58d547c
         env.allowLocalModels = false;
-        tokenizer = await AutoTokenizer.from_pretrained(tokenizer);
-        task.trainingInformation.tokenizer = tokenizer;
+        const to_wrap = await AutoTokenizer.from_pretrained(name);
+        return new Tokenizer(name, to_wrap);
+    }
+    tokenize(text, config = {}) {
+        if (config.padding || config.truncation)
+            if (!Number.isInteger(config.max_length))
+                throw new Error("max_length should be an integer");
+        if (config.padding) {
+            // The padding side is set as an attribute, not in the config
+            this.#wrapped.padding_side = config.padding ?? "left";
+            config.truncation = true;
+        }
+        const tokenizerResult = this.#wrapped(text, {
+            padding: config.padding !== undefined,
+            truncation: config.truncation,
+            max_length: "max_length" in config ? config.max_length : undefined,
+            return_tensor: false,
+        });
+        if (typeof tokenizerResult !== "object" ||
+            tokenizerResult === null ||
+            !("input_ids" in tokenizerResult) ||
+            !isArrayOfNumber(tokenizerResult.input_ids))
+            throw new Error("tokenizer returned unexpected type");
+        return List(tokenizerResult.input_ids);
     }
-    return tokenizer;
+    decode(tokens) {
+        return this.#wrapped.decode(tokens);
+    }
+}
+function isArrayOfNumber(raw) {
+    return Array.isArray(raw) && raw.every((e) => typeof e === "number");
 }

package/dist/processing/index.d.ts

@@ -1,8 +1,7 @@
 /** Dataset shapers, convenient to map with */
-import type { Dataset, DataFormat, DataType, Task } from "../index.js";
+import type { Dataset, DataFormat, DataType, Task, Network } from "../index.js";
 export * from "./image.js";
 export * from "./tabular.js";
-export * from "./text.js";
-export declare function preprocess<D extends DataType>(task: Task<D>, dataset: Dataset<DataFormat.Raw[D]>): Promise<Dataset<DataFormat.ModelEncoded[D]>>;
-export declare function preprocessWithoutLabel<D extends DataType>(task: Task<D>, dataset: Dataset<DataFormat.RawWithoutLabel[D]>): Promise<Dataset<DataFormat.ModelEncoded[D][0]>>;
-export declare function postprocess<D extends DataType>(task: Task<D>, encoded: DataFormat.ModelEncoded[D][1]): Promise<DataFormat.Inferred[D]>;
+export declare function preprocess<D extends DataType, N extends Network>(task: Task<D, N>, dataset: Dataset<DataFormat.Raw[D]>): Dataset<DataFormat.ModelEncoded[D]>;
+export declare function preprocessWithoutLabel<D extends DataType>(task: Task<D, Network>, dataset: Dataset<DataFormat.RawWithoutLabel[D]>): Dataset<DataFormat.ModelEncoded[D][0]>;
+export declare function postprocess<D extends DataType>(task: Task<D, Network>, encoded: DataFormat.ModelEncoded[D][1]): DataFormat.Inferred[D];

package/dist/processing/index.js

@@ -1,12 +1,10 @@
 /** Dataset shapers, convenient to map with */
 import { List } from "immutable";
-import { models } from "../index.js";
 import * as processing from "./index.js";
 export * from "./image.js";
 export * from "./tabular.js";
-export * from "./text.js";
-export async function preprocess(task, dataset) {
-    switch (task.trainingInformation.dataType) {
+export function preprocess(task, dataset) {
+    switch (task.dataType) {
         case "image": {
             // cast as typescript doesn't reduce generic type
             const d = dataset;
@@ -32,18 +30,17 @@ export async function preprocess(task, dataset) {
         case "text": {
             // cast as typescript doesn't reduce generic type
             const d = dataset;
-            const t = task;
-            const contextLength = task.trainingInformation.contextLength;
-            const tokenizer = await models.getTaskTokenizer(t);
-            return d.map(text => processing.tokenize(tokenizer, text))
+            const { contextLength, tokenizer } = task.trainingInformation;
+            return d
+                .map((text) => tokenizer.tokenize(text))
                 .flatten()
                 .batch(contextLength + 1, 1)
                 .map((tokens) => [tokens.pop(), tokens.last()]);
         }
     }
 }
-export async function preprocessWithoutLabel(task, dataset) {
-    switch (task.trainingInformation.dataType) {
+export function preprocessWithoutLabel(task, dataset) {
+    switch (task.dataType) {
         case "image": {
             // cast as typescript doesn't reduce generic type
             const d = dataset;
@@ -59,22 +56,20 @@ export async function preprocessWithoutLabel(task, dataset) {
         case "text": {
             // cast as typescript doesn't reduce generic type
             const d = dataset;
-            const t = task;
-            const contextLength = task.trainingInformation.contextLength;
-            const tokenizer = await models.getTaskTokenizer(t);
-            return d.map(text => processing.tokenize(tokenizer, text))
+            const { contextLength, tokenizer } = task.trainingInformation;
+            return d
+                .map((text) => tokenizer.tokenize(text))
                 .flatten()
                 .batch(contextLength);
         }
     }
 }
-export async function postprocess(task, encoded) {
-    switch (task.trainingInformation.dataType) {
+export function postprocess(task, encoded) {
+    switch (task.dataType) {
         case "image": {
             // cast as typescript doesn't reduce generic type
             const index = encoded;
-            const { LABEL_LIST } = task.trainingInformation;
-            const labels = List(LABEL_LIST);
+            const labels = List(task.trainingInformation.LABEL_LIST);
             const v = labels.get(index);
             if (v === undefined)
                 throw new Error("index not found in labels");
@@ -88,9 +83,9 @@ export async function postprocess(task, encoded) {
         case "text": {
             // cast as typescript doesn't reduce generic type
             const token = encoded;
-            const t = task;
-            const tokenizer = await models.getTaskTokenizer(t);
-            return tokenizer.decode([token]);
+            return task.trainingInformation.tokenizer.decode([
+                token,
+            ]);
         }
     }
 }

package/dist/serialization/coder.d.ts

@@ -1,4 +1,8 @@
 export type Encoded = Uint8Array;
 export declare function isEncoded(raw: unknown): raw is Encoded;
-export declare function encode(serialized: unknown): Encoded;
+type Encodable = null | undefined | boolean | number | string | Float32Array | ArrayBuffer | Encodable[] | {
+    [_: string]: Encodable;
+};
+export declare function encode(serialized: Encodable): Encoded;
 export declare function decode(encoded: Encoded): unknown;
+export {};

package/dist/serialization/coder.js

@@ -44,7 +44,10 @@ CODEC.register({
     copy(raw),
 });
 export function encode(serialized) {
-    return msgpack.encode(serialized, { extensionCodec: CODEC });
+    return msgpack.encode(serialized, {
+        extensionCodec: CODEC,
+        ignoreUndefined: true,
+    });
 }
 export function decode(encoded) {
     return msgpack.decode(encoded, { extensionCodec: CODEC });

package/dist/serialization/index.d.ts

@@ -1,4 +1,8 @@
 export * as model from './model.js';
+export * as task from "./task.js";
 export * as weights from './weights.js';
 export type { Encoded } from "./coder.js";
 export { isEncoded } from "./coder.js";
+export type JSON = null | undefined | boolean | number | string | JSON[] | {
+    [_: string]: JSON;
+};

package/dist/serialization/index.js

@@ -1,3 +1,4 @@
 export * as model from './model.js';
+export * as task from "./task.js";
 export * as weights from './weights.js';
 export { isEncoded } from "./coder.js";

package/dist/serialization/task.d.ts

@@ -0,0 +1,5 @@
+import type { DataType, Network } from "../index.js";
+import { Task } from "../index.js";
+import type { JSON } from "./index.js";
+export declare function serializeToJSON(task: Task<DataType, Network>): JSON;
+export declare function deserializeFromJSON(serialized: JSON): Promise<Task<DataType, Network>>;

package/dist/serialization/task.js

@@ -0,0 +1,34 @@
+import { z } from "zod";
+import { Task, Tokenizer } from "../index.js";
+export function serializeToJSON(task) {
+    switch (task.dataType) {
+        case "image":
+        case "tabular":
+            return task;
+        case "text": {
+            return {
+                ...task,
+                trainingInformation: {
+                    ...task.trainingInformation,
+                    tokenizer: task.trainingInformation.tokenizer.name,
+                },
+            };
+        }
+    }
+}
+export async function deserializeFromJSON(serialized) {
+    return await z
+        .object({
+        trainingInformation: z
+            .object({
+            tokenizer: z
+                .string()
+                .transform((name) => Tokenizer.from_pretrained(name))
+                .optional(),
+        })
+            .passthrough(),
+    })
+        .passthrough()
+        .pipe(Task.schema)
+        .parseAsync(serialized);
+}

package/dist/task/display_information.d.ts

@@ -1,15 +1,92 @@
-import { type Summary } from './summary.js';
-import { type DataExample } from './data_example.js';
-export interface DisplayInformation {
-    taskTitle: string;
-    summary: Summary;
-    dataFormatInformation?: string;
-    dataExampleText?: string;
-    model?: string;
-    dataExample?: DataExample[];
-    headers?: string[];
-    dataExampleImage?: string;
-    sampleDatasetLink?: string;
-    sampleDatasetInstructions?: string;
+import { z } from "zod";
+import type { DataType } from "../types/index.js";
+export declare namespace DisplayInformation {
+    const baseSchema: z.ZodObject<{
+        title: z.ZodString;
+        summary: z.ZodObject<{
+            preview: z.ZodString;
+            overview: z.ZodString;
+        }, "strip", z.ZodTypeAny, {
+            preview: string;
+            overview: string;
+        }, {
+            preview: string;
+            overview: string;
+        }>;
+        dataFormatInformation: z.ZodOptional<z.ZodString>;
+        model: z.ZodOptional<z.ZodString>;
+        sampleDataset: z.ZodOptional<z.ZodObject<{
+            link: z.ZodString;
+            instructions: z.ZodString;
+        }, "strip", z.ZodTypeAny, {
+            link: string;
+            instructions: string;
+        }, {
+            link: string;
+            instructions: string;
+        }>>;
+    }, "strip", z.ZodTypeAny, {
+        title: string;
+        summary: {
+            preview: string;
+            overview: string;
+        };
+        dataFormatInformation?: string | undefined;
+        model?: string | undefined;
+        sampleDataset?: {
+            link: string;
+            instructions: string;
+        } | undefined;
+    }, {
+        title: string;
+        summary: {
+            preview: string;
+            overview: string;
+        };
+        dataFormatInformation?: string | undefined;
+        model?: string | undefined;
+        sampleDataset?: {
+            link: string;
+            instructions: string;
+        } | undefined;
+    }>;
+    const dataTypeToSchema: {
+        image: z.ZodObject<{
+            dataExample: z.ZodOptional<z.ZodString>;
+        }, "strip", z.ZodTypeAny, {
+            dataExample?: string | undefined;
+        }, {
+            dataExample?: string | undefined;
+        }>;
+        tabular: z.ZodObject<{
+            dataExample: z.ZodOptional<z.ZodArray<z.ZodObject<{
+                name: z.ZodString;
+                data: z.ZodString;
+            }, "strip", z.ZodTypeAny, {
+                name: string;
+                data: string;
+            }, {
+                name: string;
+                data: string;
+            }>, "many">>;
+        }, "strip", z.ZodTypeAny, {
+            dataExample?: {
+                name: string;
+                data: string;
+            }[] | undefined;
+        }, {
+            dataExample?: {
+                name: string;
+                data: string;
+            }[] | undefined;
+        }>;
+        text: z.ZodObject<{
+            dataExample: z.ZodOptional<z.ZodString>;
+        }, "strip", z.ZodTypeAny, {
+            dataExample?: string | undefined;
+        }, {
+            dataExample?: string | undefined;
+        }>;
+    };
 }
-export declare function isDisplayInformation(raw: unknown): raw is DisplayInformation;
+export type DisplayInformation<D extends DataType> = (typeof DisplayInformation.dataTypeToSchema)[D];

package/dist/task/display_information.js

@@ -1,59 +1,35 @@
-import { isSummary } from './summary.js';
-import { isDataExample } from './data_example.js';
-export function isDisplayInformation(raw) {
-    if (typeof raw !== 'object' || raw === null) {
-        return false;
-    }
-    const { dataExample, dataExampleImage, dataExampleText, dataFormatInformation, sampleDatasetLink, sampleDatasetInstructions, headers, model, summary, taskTitle, } = raw;
-    if (typeof taskTitle !== 'string' ||
-        (dataExampleText !== undefined && typeof dataExampleText !== 'string') ||
-        (sampleDatasetLink !== undefined && typeof sampleDatasetLink !== 'string') ||
-        (dataFormatInformation !== undefined && typeof dataFormatInformation !== 'string') ||
-        (model !== undefined && typeof model !== 'string') ||
-        (dataExampleImage !== undefined && typeof dataExampleImage !== 'string') ||
-        (sampleDatasetInstructions !== undefined && typeof sampleDatasetInstructions !== 'string')) {
-        return false;
-    }
-    if (!isSummary(summary)) {
-        return false;
-    }
-    if (sampleDatasetLink !== undefined) {
-        try {
-            new URL(sampleDatasetLink);
-        }
-        catch {
-            return false;
-        }
-    }
-    if (dataExampleImage !== undefined) {
-        try {
-            new URL(dataExampleImage);
-        }
-        catch {
-            return false;
-        }
-    }
-    if (dataExample !== undefined && !(Array.isArray(dataExample) &&
-        dataExample.every(isDataExample))) {
-        return false;
-    }
-    if (headers !== undefined && !(Array.isArray(headers) &&
-        headers.every((e) => typeof e === 'string'))) {
-        return false;
-    }
-    const repack = {
-        dataExample,
-        dataExampleImage,
-        dataExampleText,
-        dataFormatInformation,
-        sampleDatasetLink,
-        sampleDatasetInstructions,
-        headers,
-        model,
-        summary,
-        taskTitle,
+import { z } from "zod";
+export var DisplayInformation;
+(function (DisplayInformation) {
+    DisplayInformation.baseSchema = z.object({
+        title: z.string(),
+        summary: z.object({
+            preview: z.string(),
+            overview: z.string(),
+        }),
+        dataFormatInformation: z.string().optional(),
+        model: z.string().optional(),
+        sampleDataset: z
+            .object({
+            // URL to download a dataset for the task, is displayed in the UI when asking to connect data
+            link: z.string(),
+            // Instructions to download, unzip, and connect the right file of the sample dataset
+            instructions: z.string(),
+        })
+            .optional(),
+    });
+    DisplayInformation.dataTypeToSchema = {
+        image: z.object({
+            // url to an image
+            dataExample: z.string().optional(),
+        }),
+        tabular: z.object({
+            dataExample: z
+                .array(z.object({ name: z.string(), data: z.string() }))
+                .optional(),
+        }),
+        text: z.object({
+            dataExample: z.string().optional(),
+        }),
     };
-    const _correct = repack;
-    const _total = repack;
-    return true;
-}
+})(DisplayInformation || (DisplayInformation = {}));

package/dist/task/index.d.ts

@@ -1,5 +1,5 @@
-export { isTask, type Task, isTaskID, type TaskID } from './task.js';
-export { type TaskProvider } from './task_provider.js';
-export { isDisplayInformation, type DisplayInformation } from './display_information.js';
-export type { TrainingInformation } from './training_information.js';
-export { pushTask, fetchTasks } from './task_handler.js';
+export { Task } from "./task.js";
+export type { TaskProvider } from "./task_provider.js";
+export { DisplayInformation } from "./display_information.js";
+export { TrainingInformation } from "./training_information.js";
+export { pushTask, fetchTasks } from "./task_handler.js";

package/dist/task/index.js

@@ -1,3 +1,4 @@
-export { isTask, isTaskID } from './task.js';
-export { isDisplayInformation } from './display_information.js';
-export { pushTask, fetchTasks } from './task_handler.js';
+export { Task } from "./task.js";
+export { DisplayInformation } from "./display_information.js";
+export { TrainingInformation } from "./training_information.js";
+export { pushTask, fetchTasks } from "./task_handler.js";