@epfml/discojs-web 2.1.2-p20240603114517.0 → 2.1.2-p20240617140649.0

package/dist/data/text_loader.js

@@ -3,7 +3,7 @@ import { data } from '@epfml/discojs';
 export class TextLoader extends data.TextLoader {
     loadDatasetFrom(source) {
         const file = new tf.data.FileDataSource(source);
-        const dataset = new tf.data.TextLineDataset(file).filter(s => s != ' '); // newline creates empty strings
+        const dataset = new tf.data.TextLineDataset(file).filter(s => s !== ' '); // newline creates empty strings
         return Promise.resolve(dataset);
     }
 }

package/dist/memory/memory.d.ts

@@ -2,8 +2,8 @@ import * as tf from '@tensorflow/tfjs';
 import type { Path, Model, ModelInfo, ModelSource } from '@epfml/discojs';
 import { Memory } from '@epfml/discojs';
 export declare class IndexedDB extends Memory {
-    pathFor(source: ModelSource): Path;
-    infoFor(source: ModelSource): ModelInfo;
+    getModelMemoryPath(source: ModelSource): Path;
+    getModelInfo(source: ModelSource): ModelInfo;
     getModelMetadata(source: ModelSource): Promise<tf.io.ModelArtifactsInfo | undefined>;
     contains(source: ModelSource): Promise<boolean>;
     getModel(source: ModelSource): Promise<Model>;

package/dist/memory/memory.js

@@ -8,48 +8,61 @@
  */
 import { Map } from 'immutable';
 import * as tf from '@tensorflow/tfjs';
-import { Memory, ModelType, models } from '@epfml/discojs';
+import { Memory, models } from '@epfml/discojs';
 export class IndexedDB extends Memory {
-    pathFor(source) {
+    getModelMemoryPath(source) {
         if (typeof source === 'string') {
             return source;
         }
-        if (source.type === undefined || source.taskID === undefined || source.name === undefined) {
-            throw new TypeError('source incomplete');
-        }
         const version = source.version ?? 0;
-        return `indexeddb://${source.type}/${source.taskID}/${source.name}@${version}`;
+        return `indexeddb://${source.type}/${source.tensorBackend}/${source.taskID}/${source.name}@${version}`;
     }
-    infoFor(source) {
+    getModelInfo(source) {
         if (typeof source !== 'string') {
             return source;
         }
-        const [stringType, taskID, fullName] = source.split('/').splice(2);
-        const type = stringType === 'working' ? ModelType.WORKING : ModelType.SAVED;
+        const [type, tensorBackend, taskID, fullName] = source.split('/').splice(2);
+        if (type !== 'working' && type !== 'saved') {
+            throw Error("Unknown memory model type");
+        }
         const [name, versionSuffix] = fullName.split('@');
         const version = versionSuffix === undefined ? 0 : Number(versionSuffix);
-        return { type, taskID, name, version };
+        if (tensorBackend !== 'tfjs' && tensorBackend !== 'gpt') {
+            throw Error("Unknown tensor backend");
+        }
+        return { type, taskID, name, version, tensorBackend };
     }
     async getModelMetadata(source) {
         const models = await tf.io.listModels();
-        return models[this.pathFor(source)];
+        return models[this.getModelMemoryPath(source)];
     }
     async contains(source) {
         return await this.getModelMetadata(source) !== undefined;
     }
     async getModel(source) {
-        return new models.TFJS(await tf.loadLayersModel(this.pathFor(source)));
+        const layersModel = await tf.loadLayersModel(this.getModelMemoryPath(source));
+        const tensorBackend = this.getModelInfo(source).tensorBackend;
+        switch (tensorBackend) {
+            case 'tfjs':
+                return new models.TFJS(layersModel);
+            case 'gpt':
+                return new models.GPT(undefined, layersModel);
+            default: {
+                const _ = tensorBackend;
+                throw new Error('should never happen');
+            }
+        }
     }
     async deleteModel(source) {
-        await tf.io.removeModel(this.pathFor(source));
+        await tf.io.removeModel(this.getModelMemoryPath(source));
     }
     async loadModel(source) {
-        const src = this.infoFor(source);
-        if (src.type === ModelType.WORKING) {
+        const src = this.getModelInfo(source);
+        if (src.type === 'working') {
             // Model is already loaded
             return;
         }
-        await tf.io.copyModel(this.pathFor(src), this.pathFor({ ...src, type: ModelType.WORKING, version: 0 }));
+        await tf.io.copyModel(this.getModelMemoryPath(src), this.getModelMemoryPath({ ...src, type: 'working', version: 0 }));
     }
     /**
      * Saves the working model to the source.
@@ -57,67 +70,81 @@ export class IndexedDB extends Memory {
      * @param model the model
      */
     async updateWorkingModel(source, model) {
-        const src = this.infoFor(source);
-        if (src.type !== undefined && src.type !== ModelType.WORKING) {
-            throw new Error('expected working model');
+        const src = this.getModelInfo(source);
+        if (src.type !== 'working') {
+            throw new Error('expected working type model');
         }
+        // Enforce version 0 to always keep a single working model at a time
+        const modelInfo = { ...src, type: 'working', version: 0 };
+        let includeOptimizer;
         if (model instanceof models.TFJS) {
-            await model.extract().save(this.pathFor({ ...src, type: ModelType.WORKING, version: 0 }), { includeOptimizer: true });
+            modelInfo['tensorBackend'] = 'tfjs';
+            includeOptimizer = true;
+        }
+        else if (model instanceof models.GPT) {
+            modelInfo['tensorBackend'] = 'gpt';
+            includeOptimizer = false; // true raises an error
         }
         else {
             throw new Error('unknown model type');
         }
-        // Enforce version 0 to always keep a single working model at a time
+        const indexedDBURL = this.getModelMemoryPath(modelInfo);
+        await model.extract().save(indexedDBURL, { includeOptimizer });
     }
     /**
    * Creates a saved copy of the working model corresponding to the source.
    * @param source the source
    */
     async saveWorkingModel(source) {
-        const src = this.infoFor(source);
-        if (src.type !== undefined && src.type !== ModelType.WORKING) {
-            throw new Error('expected working model');
+        const src = this.getModelInfo(source);
+        if (src.type !== 'working') {
+            throw new Error('expected working type model');
         }
-        const dst = this.pathFor(await this.duplicateSource({ ...src, type: ModelType.SAVED }));
-        await tf.io.copyModel(this.pathFor({ ...src, type: ModelType.WORKING }), dst);
+        const dst = this.getModelMemoryPath(await this.duplicateSource({ ...src, type: 'saved' }));
+        await tf.io.copyModel(this.getModelMemoryPath({ ...src, type: 'working' }), dst);
         return dst;
     }
     async saveModel(source, model) {
-        const src = this.infoFor(source);
-        if (src.type !== undefined && src.type !== ModelType.SAVED) {
-            throw new Error('expected saved model');
+        const src = this.getModelInfo(source);
+        if (src.type !== 'saved') {
+            throw new Error('expected saved type model');
         }
-        const dst = this.pathFor(await this.duplicateSource({ ...src, type: ModelType.SAVED }));
+        const modelInfo = await this.duplicateSource({ ...src, type: 'saved' });
+        let includeOptimizer;
         if (model instanceof models.TFJS) {
-            await model.extract().save(dst, { includeOptimizer: true });
+            modelInfo['tensorBackend'] = 'tfjs';
+            includeOptimizer = true;
+        }
+        else if (model instanceof models.GPT) {
+            modelInfo['tensorBackend'] = 'gpt';
+            includeOptimizer = false; // true raises an error
         }
         else {
             throw new Error('unknown model type');
         }
-        return dst;
+        const indexedDBURL = this.getModelMemoryPath(modelInfo);
+        await model.extract().save(indexedDBURL, { includeOptimizer });
+        return indexedDBURL;
     }
     /**
    * Downloads the model corresponding to the source.
    * @param source the source
    */
     async downloadModel(source) {
-        const src = this.infoFor(source);
-        await tf.io.copyModel(this.pathFor(source), `downloads://${src.taskID}_${src.name}`);
+        const src = this.getModelInfo(source);
+        await tf.io.copyModel(this.getModelMemoryPath(source), `downloads://${src.taskID}_${src.name}`);
     }
     async latestDuplicate(source) {
         if (typeof source !== 'string') {
-            source = this.pathFor({ ...source, version: 0 });
+            source = this.getModelMemoryPath({ ...source, version: 0 });
         }
         // perform a single memory read
         const paths = Map(await tf.io.listModels());
         if (!paths.has(source)) {
             return undefined;
         }
-        const latest = Map(paths)
-            .keySeq()
-            .toList()
-            .map((p) => this.infoFor(p).version)
-            .max();
+        const latest = Map(paths).keySeq().toList()
+            .map((p) => this.getModelInfo(p).version).max();
         if (latest === undefined) {
             return 0;
         }
@@ -125,7 +152,7 @@ export class IndexedDB extends Memory {
     }
     async duplicateSource(source) {
         const latestDuplicate = await this.latestDuplicate(source);
-        source = this.infoFor(source);
+        source = this.getModelInfo(source);
         if (latestDuplicate === undefined) {
             return source;
         }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epfml/discojs-web",
-  "version": "2.1.2-p20240603114517.0",
+  "version": "2.1.2-p20240617140649.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",