discojs/model/tfjs: drop text support

epfml · Oct 4, 2024 · 0ac6c29 · 0ac6c29
1 parent 9f10e3e
commit 0ac6c29
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 34 deletions.
diff --git a/discojs/src/models/tfjs.ts b/discojs/src/models/tfjs.ts
@@ -16,7 +16,7 @@ import { EpochLogs } from './logs.js'
 type Serialized<D extends DataType> = [D, tf.io.ModelArtifacts];
 
 /** TensorFlow JavaScript model with standard training */
-export class TFJS<D extends DataType> extends Model<D> {
+export class TFJS<D extends "image" | "tabular"> extends Model<D> {
   /** Wrap the given trainable model */
   constructor (
     public readonly datatype: D,
@@ -168,7 +168,7 @@ export class TFJS<D extends DataType> extends Model<D> {
     return ret
   }
 
-  static async deserialize<D extends DataType>([
+  static async deserialize<D extends "image" | "tabular">([
     datatype,
     artifacts,
   ]: Serialized<D>): Promise<TFJS<D>> {
@@ -259,23 +259,6 @@ export class TFJS<D extends DataType> extends Model<D> {
           ys: tf.stack(b.map(([_, output]) => tf.tensor1d([output])).toArray()),
         }));
       }
-      case "text": {
-        // cast as typescript doesn't reduce generic type
-        const b = batch as Batched<ModelEncoded["text"]>;
-
-        return {
-          xs: tf.stack(
-            b.map(([line]) => tf.tensor1d(line.toArray())).toArray(),
-          ),
-          ys: tf.stack(
-            b
-              .map(([line, next]) =>
-                tf.oneHot(line.shift().push(next).toArray(), outputSize),
-              )
-              .toArray(),
-          ),
-        };
-      }
     }
 
     const _: never = this.datatype;
@@ -311,14 +294,6 @@ export class TFJS<D extends DataType> extends Model<D> {
           ),
         );
       }
-      case "text": {
-        // cast as typescript doesn't reduce generic type
-        const b = batch as Batched<ModelEncoded["text"][0]>;
-
-        return tf.stack(
-            b.map((line) => tf.tensor1d(line.toArray())).toArray(),
-          )
-      }
     }
 
     const _: never = this.datatype;

diff --git a/discojs/src/serialization/model.spec.ts b/discojs/src/serialization/model.spec.ts
@@ -35,7 +35,9 @@ describe('serialization', () => {
     const decoded = await serialization.model.decode(encoded)
 
     expect(decoded).to.be.an.instanceof(models.TFJS);
-    expect((decoded as models.TFJS<DataType>).datatype).to.equal("image")
+    expect((decoded as models.TFJS<"image" | "tabular">).datatype).to.equal(
+      "image",
+    );
     assert.sameDeepOrderedMembers(
       await getRawWeights(model),
       await getRawWeights(decoded)

diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts
@@ -61,11 +61,10 @@ export async function decode(encoded: unknown): Promise<Model<DataType>> {
         );
       const [rawDatatype, rawModel] = raw.slice(1) as unknown[];
 
-      let datatype: DataType;
+      let datatype;
       switch (rawDatatype) {
         case "image":
         case "tabular":
-        case "text":
           datatype = rawDatatype;
           break;
         default:

diff --git a/webapp/src/components/task_creation_form/TaskForm.vue b/webapp/src/components/task_creation_form/TaskForm.vue
@@ -279,10 +279,20 @@ const onSubmit = async (rawTask: any): Promise<void> => {
 
   let model
   try {
-    model = new models.TFJS(
-      task.trainingInformation.dataType,
-      await tf.loadLayersModel(tf.io.browserFiles(modelFiles.value.toArray())),
-    );
+    switch (task.trainingInformation.dataType) {
+      case "image":
+      case "tabular":
+        model = new models.TFJS(
+          task.trainingInformation.dataType,
+          await tf.loadLayersModel(
+            tf.io.browserFiles(modelFiles.value.toArray()),
+          ),
+        );
+        break;
+      case "text":
+        toaster.error("Currently no support of TFJS text model");
+        return;
+    }
   } catch (e) {
     debug("while loading model:%o", e);
     toaster.error('Model loading failed');