@epfml/discojs-web 2.1.2-p20240515132210.0 → 2.1.2-p20240515133413.0

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 {
-    getModelMemoryPath(source: ModelSource): Path;
-    getModelInfo(source: ModelSource): ModelInfo;
+    pathFor(source: ModelSource): Path;
+    infoFor(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,9 +8,9 @@
  */
 import { Map } from 'immutable';
 import * as tf from '@tensorflow/tfjs';
-import { Memory, StoredModelType, models } from '@epfml/discojs';
+import { Memory, ModelType, models } from '@epfml/discojs';
 export class IndexedDB extends Memory {
-    getModelMemoryPath(source) {
+    pathFor(source) {
         if (typeof source === 'string') {
             return source;
         }
@@ -20,41 +20,36 @@ export class IndexedDB extends Memory {
         const version = source.version ?? 0;
         return `indexeddb://${source.type}/${source.taskID}/${source.name}@${version}`;
     }
-    getModelInfo(source) {
+    infoFor(source) {
         if (typeof source !== 'string') {
             return source;
         }
         const [stringType, taskID, fullName] = source.split('/').splice(2);
-        const type = stringType === 'working' ? StoredModelType.WORKING : StoredModelType.SAVED;
+        const type = stringType === 'working' ? ModelType.WORKING : ModelType.SAVED;
         const [name, versionSuffix] = fullName.split('@');
         const version = versionSuffix === undefined ? 0 : Number(versionSuffix);
         return { type, taskID, name, version };
     }
     async getModelMetadata(source) {
         const models = await tf.io.listModels();
-        return models[this.getModelMemoryPath(source)];
+        return models[this.pathFor(source)];
     }
     async contains(source) {
         return await this.getModelMetadata(source) !== undefined;
     }
     async getModel(source) {
-        console.log("source", source);
-        console.log("memory path", this.getModelMemoryPath(source));
-        const layersModel = await tf.loadLayersModel(this.getModelMemoryPath(source));
-        console.log("layers model", layersModel);
-        return new models.TFJS(layersModel);
+        return new models.TFJS(await tf.loadLayersModel(this.pathFor(source)));
     }
     async deleteModel(source) {
-        await tf.io.removeModel(this.getModelMemoryPath(source));
+        await tf.io.removeModel(this.pathFor(source));
     }
     async loadModel(source) {
-        console.log("Loading model");
-        const src = this.getModelInfo(source);
-        if (src.type === StoredModelType.WORKING) {
+        const src = this.infoFor(source);
+        if (src.type === ModelType.WORKING) {
             // Model is already loaded
             return;
         }
-        await tf.io.copyModel(this.getModelMemoryPath(src), this.getModelMemoryPath({ ...src, type: StoredModelType.WORKING, version: 0 }));
+        await tf.io.copyModel(this.pathFor(src), this.pathFor({ ...src, type: ModelType.WORKING, version: 0 }));
     }
     /**
      * Saves the working model to the source.
@@ -62,16 +57,12 @@ export class IndexedDB extends Memory {
      * @param model the model
      */
     async updateWorkingModel(source, model) {
-        const src = this.getModelInfo(source);
-        if (src.type !== undefined && src.type !== StoredModelType.WORKING) {
+        const src = this.infoFor(source);
+        if (src.type !== undefined && src.type !== ModelType.WORKING) {
             throw new Error('expected working model');
         }
-        const indexedDBURL = this.getModelMemoryPath({ ...src, type: StoredModelType.WORKING, version: 0 });
         if (model instanceof models.TFJS) {
-            await model.extract().save(indexedDBURL, { includeOptimizer: true });
-        }
-        else if (model instanceof models.GPT) {
-            await model.extract().save(indexedDBURL, { includeOptimizer: false }); // true raises an error
+            await model.extract().save(this.pathFor({ ...src, type: ModelType.WORKING, version: 0 }), { includeOptimizer: true });
         }
         else {
             throw new Error('unknown model type');
@@ -83,42 +74,39 @@ export class IndexedDB extends Memory {
    * @param source the source
    */
     async saveWorkingModel(source) {
-        const src = this.getModelInfo(source);
-        if (src.type !== undefined && src.type !== StoredModelType.WORKING) {
+        const src = this.infoFor(source);
+        if (src.type !== undefined && src.type !== ModelType.WORKING) {
             throw new Error('expected working model');
         }
-        const dst = this.getModelMemoryPath(await this.duplicateSource({ ...src, type: StoredModelType.SAVED }));
-        await tf.io.copyModel(this.getModelMemoryPath({ ...src, type: StoredModelType.WORKING }), dst);
+        const dst = this.pathFor(await this.duplicateSource({ ...src, type: ModelType.SAVED }));
+        await tf.io.copyModel(this.pathFor({ ...src, type: ModelType.WORKING }), dst);
         return dst;
     }
     async saveModel(source, model) {
-        const src = this.getModelInfo(source);
-        if (src.type !== undefined && src.type !== StoredModelType.SAVED) {
+        const src = this.infoFor(source);
+        if (src.type !== undefined && src.type !== ModelType.SAVED) {
             throw new Error('expected saved model');
         }
-        const indexedDBURL = this.getModelMemoryPath(await this.duplicateSource({ ...src, type: StoredModelType.SAVED }));
+        const dst = this.pathFor(await this.duplicateSource({ ...src, type: ModelType.SAVED }));
         if (model instanceof models.TFJS) {
-            await model.extract().save(indexedDBURL, { includeOptimizer: true });
-        }
-        else if (model instanceof models.GPT) {
-            await model.extract().save(indexedDBURL, { includeOptimizer: false }); // true raises an error
+            await model.extract().save(dst, { includeOptimizer: true });
         }
         else {
             throw new Error('unknown model type');
         }
-        return indexedDBURL;
+        return dst;
     }
     /**
    * Downloads the model corresponding to the source.
    * @param source the source
    */
     async downloadModel(source) {
-        const src = this.getModelInfo(source);
-        await tf.io.copyModel(this.getModelMemoryPath(source), `downloads://${src.taskID}_${src.name}`);
+        const src = this.infoFor(source);
+        await tf.io.copyModel(this.pathFor(source), `downloads://${src.taskID}_${src.name}`);
     }
     async latestDuplicate(source) {
         if (typeof source !== 'string') {
-            source = this.getModelMemoryPath({ ...source, version: 0 });
+            source = this.pathFor({ ...source, version: 0 });
         }
         // perform a single memory read
         const paths = Map(await tf.io.listModels());
@@ -128,7 +116,7 @@ export class IndexedDB extends Memory {
         const latest = Map(paths)
             .keySeq()
             .toList()
-            .map((p) => this.getModelInfo(p).version)
+            .map((p) => this.infoFor(p).version)
             .max();
         if (latest === undefined) {
             return 0;
@@ -137,7 +125,7 @@ export class IndexedDB extends Memory {
     }
     async duplicateSource(source) {
         const latestDuplicate = await this.latestDuplicate(source);
-        source = this.getModelInfo(source);
+        source = this.infoFor(source);
         if (latestDuplicate === undefined) {
             return source;
         }

package/package.json

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