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');