From 58f916cff11b7a377b99828bcbc06c2eea3e0e3a Mon Sep 17 00:00:00 2001
From: tharvik <tharvik@users.noreply.github.com>
Date: Fri, 16 Aug 2024 23:49:25 +0200
Subject: [PATCH] discojs/validation: simplify

---
 discojs/src/dataset/data/image_data.ts        |   4 +-
 discojs/src/validation/validator.ts           | 254 +++++-----
 server/tests/validator.spec.ts                | 100 ++--
 .../src/components/containers/TableLayout.vue |  30 +-
 .../src/components/testing/PredictSteps.vue   | 327 +++++++------
 webapp/src/components/testing/TestSteps.vue   | 440 ++++++++++--------
 webapp/src/components/testing/Testing.vue     |  47 +-
 7 files changed, 601 insertions(+), 601 deletions(-)

diff --git a/discojs/src/dataset/data/image_data.ts b/discojs/src/dataset/data/image_data.ts
index d6d7a4849..e4c771a02 100644
--- a/discojs/src/dataset/data/image_data.ts
+++ b/discojs/src/dataset/data/image_data.ts
@@ -32,8 +32,8 @@ export class ImageData extends Data {
       }
 
       let shape
-      if ('xs' in sample && 'ys' in sample) {
-        shape = (sample as { xs: tf.Tensor, ys: number[] }).xs.shape
+      if ('xs' in sample) {
+        shape = (sample as { xs: tf.Tensor }).xs.shape
       } else {
         shape = (sample as tf.Tensor3D).shape
       }
diff --git a/discojs/src/validation/validator.ts b/discojs/src/validation/validator.ts
index fd373e085..437cbf73a 100644
--- a/discojs/src/validation/validator.ts
+++ b/discojs/src/validation/validator.ts
@@ -1,149 +1,165 @@
-import { List } from 'immutable'
-import * as tf from '@tensorflow/tfjs'
+import * as tf from "@tensorflow/tfjs";
 
 import type {
   Model,
   Task,
-  Logger,
-  client as clients,
-  data,
-  Memory,
-  ModelSource,
   TypedDataset,
   TypedLabeledDataset,
 } from "../index.js";
-import { datasetToData, labeledDatasetToData } from "../dataset/data/helpers.js";
+import {
+  datasetToData,
+  labeledDatasetToData,
+} from "../dataset/data/helpers.js";
+
+function intoTFDataset<T extends tf.TensorContainer>(
+  iter: AsyncIterable<T>,
+): tf.data.Dataset<T> {
+  // @ts-expect-error generator
+  return tf.data.generator(async function* () {
+    yield* iter;
+  });
+}
 
 export class Validator {
-  private size = 0
-  private _confusionMatrix: number[][] | undefined
-  private rollingAccuracy = 0
+  readonly #model: Model;
 
-  constructor (
+  constructor(
     public readonly task: Task,
-    public readonly logger: Logger,
-    private readonly memory: Memory,
-    private readonly source?: ModelSource,
-    private readonly client?: clients.Client
+    model: Model,
   ) {
-    if (source === undefined && client === undefined) {
-      throw new Error('To initialize a Validator, either or both a source and client need to be specified')
-    }
+    this.#model = model;
   }
 
-  private async getLabel(ys: tf.Tensor): Promise<Float32Array | Int32Array | Uint8Array> {
-    // Binary classification
-    if (ys.shape[1] == 1) {
-      const threshold = tf.scalar(0.5)
-      const binaryTensor = ys.greaterEqual(threshold)
-      const binaryArray = await binaryTensor.data()
-      tf.dispose([binaryTensor, threshold])
-      return binaryArray
-      // Multi-class classification
-      } else {
-      const yIdxTensor = ys.argMax(-1)
-      const yIdx = await yIdxTensor.data()
-      tf.dispose([yIdxTensor])
-      return yIdx
+  /** infer every line of the dataset and check that it is as labeled */
+  async *test(dataset: TypedLabeledDataset): AsyncGenerator<boolean> {
+    const preprocessed = (
+      await labeledDatasetToData(this.task, dataset)
+    ).preprocess();
+    const batched = preprocessed.batch().dataset;
+
+    const iterator = await tf.data
+      .zip<[tf.Tensor1D | tf.Tensor2D, number]>([
+        preprocessed.dataset.map((t) => {
+          if (
+            typeof t !== "object" ||
+            !("ys" in t) ||
+            !(t.ys instanceof tf.Tensor) ||
+            !(t.ys.rank === 1 || t.ys.rank === 2)
+          )
+            throw new Error("unexpected preprocessed dataset");
+          if ("xs" in t) tf.dispose(t.xs);
+          return t.ys;
+        }),
+        intoTFDataset(this.#inferOnBatchedData(batched)),
+      ])
+      .iterator();
+    for (
+      let iter = await iterator.next();
+      iter.done !== true;
+      iter = await iterator.next()
+    ) {
+      const zipped = iter.value;
+
+      const label = await getLabel(zipped[0]);
+      tf.dispose(zipped[0]);
+      const infered = zipped[1];
+
+      yield label === infered;
     }
-    // Multi-label classification is not supported
   }
 
-  // test assumes data comes with labels while predict doesn't
-  async *test(dataset: data.Data | TypedLabeledDataset):
-    AsyncGenerator<Array<{ groundTruth: number, pred: number, features: number[] }>, void> {
-    if (Array.isArray(dataset))
-      dataset = await labeledDatasetToData(this.task, dataset)
-    const batched = dataset
-      .preprocess()
-      .dataset.batch(this.task.trainingInformation.batchSize);
-
-    const model = await this.getModel()
-    let hits = 0
-    const iterator = await batched.iterator()
-    let next = await iterator.next()
-    while (next.done !== true) {
-      const { xs, ys } = next.value as { xs: tf.Tensor2D, ys: tf.Tensor3D }
-      const ysLabel = await this.getLabel(ys)
-      const yPredTensor = await model.predict(xs)
-      const pred = await this.getLabel(yPredTensor)
-      const currentFeatures = await xs.array()
-      this.size += ysLabel.length
-      hits += List(pred).zip(List(ysLabel)).filter(([p, y]) => p === y).size
-      this.rollingAccuracy = hits / this.size
-      tf.dispose([xs, ys, yPredTensor])
-
-      yield (List(ysLabel).zip(List(pred), List(currentFeatures)))
-        .map(([gt, p, f]) => ({ groundTruth: gt, pred: p, features: f }))
-        .toArray()
-      
-      next = await iterator.next()
-    }
-    
-    this.logger.success(`Obtained validation accuracy of ${this.accuracy}`)
-    this.logger.success(`Visited ${this.visitedSamples} samples`)
+  /** use the model to predict every line of the dataset */
+  async *infer(dataset: TypedDataset): AsyncGenerator<number, void> {
+    const data = await datasetToData(this.task, dataset);
+
+    const batched = data.preprocess().batch().dataset;
+
+    yield* this.#inferOnBatchedData(batched);
   }
 
-  async *inference (dataset: data.Data | TypedDataset): AsyncGenerator<Array<{ features: number[], pred: number }>, void> {
-    if (Array.isArray(dataset))
-      dataset = await datasetToData(this.task, dataset)
-    const batched = dataset
-      .preprocess()
-      .dataset.batch(this.task.trainingInformation.batchSize);
-
-    const model = await this.getModel()
-    const iterator = await batched.iterator()
-    let next = await iterator.next()
-
-    while (next.done !== true) {
-      let xs: tf.Tensor2D
-      if (next.value instanceof tf.Tensor) {
-        xs = next.value as tf.Tensor2D
-      } else {
-        const tensors = (next.value as { xs: tf.Tensor2D, ys: tf.Tensor2D })
-        xs = tensors['xs']
-        tf.dispose([tensors['ys']])
+  async *#inferOnBatchedData(
+    batched: tf.data.Dataset<tf.TensorContainer>,
+  ): AsyncGenerator<number, void> {
+    const iterator = await batched.iterator();
+    for (
+      let iter = await iterator.next();
+      iter.done !== true;
+      iter = await iterator.next()
+    ) {
+      const row = iter.value;
+      if (
+        typeof row !== "object" ||
+        !("xs" in row) ||
+        !(row.xs instanceof tf.Tensor)
+      )
+        throw new Error("unexpected shape of dataset");
+
+      const prediction = await this.#model.predict(row.xs);
+      tf.dispose(row);
+      let predictions: number[];
+      switch (prediction.rank) {
+        case 2:
+        case 3:
+          predictions = await getLabels(
+            // cast as rank was just checked
+            prediction as tf.Tensor2D | tf.Tensor3D,
+          );
+          prediction.dispose();
+          break;
+        default:
+          throw new Error("unexpected batched prediction shape");
       }
-      const currentFeatures = await xs.array()
-      const yPredTensor = await model.predict(xs)
-      const pred = await this.getLabel(yPredTensor)
-      this.size += pred.length
-      if (!Array.isArray(currentFeatures)) {
-        throw new TypeError('Data format is incorrect')
-      }
-      tf.dispose([xs, yPredTensor])
-      
-      yield List(currentFeatures).zip(List(pred))
-        .map(([f, p]) => ({ features: f, pred: p }))
-        .toArray()
-      
-      next = await iterator.next()
+      prediction.dispose();
+
+      for (const prediction of predictions) yield prediction;
     }
-    
-    this.logger.success(`Visited ${this.visitedSamples} samples`)
   }
+}
 
-  async getModel (): Promise<Model> {
-    if (this.source !== undefined && await this.memory.contains(this.source)) {
-      return await this.memory.getModel(this.source)
-    }
+async function getLabels(ys: tf.Tensor2D | tf.Tensor3D): Promise<number[]> {
+  // cast as unstack drop a dimension and tfjs doesn't type correctly
+  return Promise.all(
+    tf.unstack(ys).map((y) => {
+      const ret = getLabel(y as tf.Tensor1D | tf.Tensor2D);
+      y.dispose();
+      return ret;
+    }),
+  );
+}
 
-    if (this.client !== undefined) {
-      return await this.client.getLatestModel()
-    }
+async function getLabel(ys: tf.Tensor1D | tf.Tensor2D): Promise<number> {
+  switch (ys.rank) {
+    case 1: {
+      if (ys.shape[0] == 1) {
+        // Binary classification
+        const threshold = tf.scalar(0.5);
+        const binaryTensor = ys.greaterEqual(threshold);
 
-    throw new Error('Could not load the model')
-  }
+        const binaryArray = await binaryTensor.data();
+        tf.dispose([binaryTensor, threshold]);
 
-  get accuracy (): number {
-    return this.rollingAccuracy
-  }
+        return binaryArray[0];
+      }
 
-  get visitedSamples (): number {
-    return this.size
-  }
+      // Multi-class classification
+      const indexTensor = ys.argMax();
 
-  get confusionMatrix (): number[][] | undefined {
-    return this._confusionMatrix
+      const indexArray = await indexTensor.data();
+      tf.dispose([indexTensor]);
+
+      return indexArray[0];
+
+      // Multi-label classification is not supported
+    }
+    case 2: {
+      // it's LLM, we only extract the next token
+      const firstToken = tf.tidy(() => ys.gather([0]).squeeze().argMax());
+      const raw = await firstToken.data();
+      firstToken.dispose();
+
+      return raw[0];
+    }
+    default:
+      throw new Error("unexpected tensor rank");
   }
 }
diff --git a/server/tests/validator.spec.ts b/server/tests/validator.spec.ts
index 89366416f..8dcd862cd 100644
--- a/server/tests/validator.spec.ts
+++ b/server/tests/validator.spec.ts
@@ -1,35 +1,12 @@
 import { expect } from "chai";
 import { Repeat } from "immutable";
-import type * as http from "node:http";
-
-import {
-  Validator,
-  ConsoleLogger,
-  EmptyMemory,
-  client as clients,
-  aggregator as aggregators,
-  defaultTasks,
-} from "@epfml/discojs";
-import { loadCSV, loadImagesInDir } from "@epfml/discojs-node";
 
-import { Server } from "../src/index.js";
+import { Validator, defaultTasks } from "@epfml/discojs";
+import { loadCSV, loadImagesInDir } from "@epfml/discojs-node";
 
 describe("validator", function () {
-  this.timeout("10s");
-
-  let server: http.Server;
-  let url: URL;
-  beforeEach(async () => {
-    [server, url] = await Server.of(
-      defaultTasks.simpleFace,
-      defaultTasks.lusCovid,
-      defaultTasks.titanic,
-    ).then((s) => s.serve());
-  });
-  afterEach(() => server?.close());
-
   it("can read and predict randomly on simple_face", async () => {
-    const task = defaultTasks.simpleFace.getTask();
+    const provider = defaultTasks.simpleFace;
 
     const [adult, child] = [
       (await loadImagesInDir("../datasets/simple_face/adult")).zip(
@@ -41,48 +18,44 @@ describe("validator", function () {
     ];
     const dataset = adult.chain(child);
 
-    // Init a validator instance
-    const meanAggregator = aggregators.getAggregator(task, { scheme: "local" });
-    const client = new clients.Local(url, task, meanAggregator);
     const validator = new Validator(
-      task,
-      new ConsoleLogger(),
-      new EmptyMemory(),
-      undefined,
-      client,
+      provider.getTask(),
+      await provider.getModel(),
     );
 
-    for await (const _ of validator.test(["image", dataset]));
+    let hits = 0;
+    let size = 0;
+    for await (const correct of validator.test(["image", dataset])) {
+      if (correct) hits++;
+      size++;
+    }
 
-    expect(validator.visitedSamples).to.equal(await dataset.size());
-    expect(validator.accuracy).to.be.greaterThan(0.3);
-  }).timeout("5s");
+    expect(hits / size).to.be.greaterThan(0.3);
+  });
 
   it("can read and predict randomly on titanic", async () => {
-    const task = defaultTasks.titanic.getTask();
+    const provider = defaultTasks.titanic;
 
     const dataset = loadCSV("../datasets/titanic_train.csv");
 
-    const meanAggregator = aggregators.getAggregator(task, { scheme: "local" });
-    const client = new clients.Local(url, task, meanAggregator);
     const validator = new Validator(
-      task,
-      new ConsoleLogger(),
-      new EmptyMemory(),
-      undefined,
-      client,
+      provider.getTask(),
+      await provider.getModel(),
     );
 
-    for await (const _ of validator.test(["tabular", dataset]));
+    let hits = 0;
+    let size = 0;
+    for await (const correct of validator.test(["tabular", dataset])) {
+      if (correct) hits++;
+      size++;
+    }
 
-    expect(validator.visitedSamples).to.equal(await dataset.size());
-    expect(validator.accuracy).to.be.greaterThan(0.3);
-  }).timeout("1s");
+    expect(hits / size).to.be.greaterThan(0.3);
+  });
 
   it("can read and predict randomly on lus_covid", async () => {
-    const task = defaultTasks.lusCovid.getTask();
+    const provider = defaultTasks.lusCovid;
 
-    // Load the data
     const [positive, negative] = [
       (await loadImagesInDir("../datasets/lus_covid/COVID+")).zip(
         Repeat("COVID-Positive"),
@@ -93,21 +66,18 @@ describe("validator", function () {
     ];
     const dataset = positive.chain(negative);
 
-    // Initialize a validator instance
-    const meanAggregator = aggregators.getAggregator(task, { scheme: "local" });
-    const client = new clients.Local(url, task, meanAggregator);
     const validator = new Validator(
-      task,
-      new ConsoleLogger(),
-      new EmptyMemory(),
-      undefined,
-      client,
+      provider.getTask(),
+      await provider.getModel(),
     );
 
-    // Assert random initialization metrics
-    for await (const _ of validator.test(["image", dataset]));
+    let hits = 0;
+    let size = 0;
+    for await (const correct of validator.test(["image", dataset])) {
+      if (correct) hits++;
+      size++;
+    }
 
-    expect(validator.visitedSamples).to.equal(await dataset.size());
-    expect(validator.accuracy).to.be.greaterThan(0.3);
-  }).timeout("1s");
+    expect(hits / size).to.be.greaterThan(0.3);
+  });
 });
diff --git a/webapp/src/components/containers/TableLayout.vue b/webapp/src/components/containers/TableLayout.vue
index b6dcb700e..19001de18 100644
--- a/webapp/src/components/containers/TableLayout.vue
+++ b/webapp/src/components/containers/TableLayout.vue
@@ -12,18 +12,9 @@
       </tr>
     </thead>
 
-    <tbody
-      class="flex flex-col overflow-y-scroll w-full h-128 px-2"
-    >
-      <tr
-        v-for="(row, i) in rows"
-        :key="i"
-      >
-        <td
-          v-for="(value, j) in row"
-          :key="j"
-          class="min-w-[200px]"
-        >
+    <tbody class="flex flex-col overflow-y-scroll w-full h-128 px-2">
+      <tr v-for="(row, i) in rows" :key="i">
+        <td v-for="(value, j) in row" :key="j" class="min-w-[200px]">
           {{ value }}
         </td>
       </tr>
@@ -31,14 +22,11 @@
   </table>
 </template>
 
-<script lang="ts">
-import { defineComponent } from 'vue'
+<script lang="ts" setup>
+import { List } from "immutable";
 
-export default defineComponent({
-  name: 'TableLayout',
-  props: {
-    columns: { default: () => {}, type: Array },
-    rows: { default: () => {}, type: Array }
-  }
-})
+defineProps<{
+  columns: List<string>;
+  rows: List<List<string>>;
+}>();
 </script>
diff --git a/webapp/src/components/testing/PredictSteps.vue b/webapp/src/components/testing/PredictSteps.vue
index aa3ff2c32..a4e1aeb7f 100644
--- a/webapp/src/components/testing/PredictSteps.vue
+++ b/webapp/src/components/testing/PredictSteps.vue
@@ -39,12 +39,12 @@
     <IconCard class="mx-auto mt-10 lg:w-1/2" title-placement="left">
       <template #title> Run model inference </template>
 
-      <div v-if="inferenceGenerator === undefined">
+      <div v-if="generator === undefined">
         By clicking the button below, you will be able to predict using the
         selected model with chosen dataset of yours.
 
         <div class="flex justify-center mt-4">
-          <CustomButton @click="modelInference()"> predict </CustomButton>
+          <CustomButton @click="startInference()"> predict </CustomButton>
         </div>
       </div>
       <div v-else>
@@ -62,69 +62,56 @@
       <!-- stats -->
       <div class="flex justify-center p-4 font-medium text-slate-500">
         <div class="text-center">
-          <span class="text-2xl">{{ visitedSamples }}</span>
+          <span class="text-2xl">
+            {{ predictions?.results.size ?? 0 }}
+          </span>
           <span class="text-sm">&nbsp;samples visited</span>
         </div>
       </div>
     </div>
 
-    <div v-if="dataWithPred !== undefined">
+    <div v-if="predictions !== undefined">
       <div class="mx-auto lg:w-1/2 text-center pb-8">
         <CustomButton @click="saveCsv()"> download as csv </CustomButton>
-        <a ref="downloadLink" class="hidden" />
       </div>
 
-      <!-- Image prediction gallery -->
-      <div v-if="dataset?.[0] === 'image'" class="grid grid-cols-6 gap-6">
+      <!-- Image gallery -->
+      <div v-if="predictions.type === 'image'" class="grid grid-cols-6 gap-6">
         <ImageCard
-          v-for="(value, key) in dataWithPred"
-          :key="key"
-          :image="(value.data as ImageWithUrl).image"
+          v-for="(result, index) in predictions.results"
+          :key="index"
+          :image="result.input.image"
         >
           <template #title>
-            <span class="font-bold">{{ getLabelName(value.prediction) }}</span>
+            <span class="font-bold">{{ result.output }}</span>
           </template>
         </ImageCard>
       </div>
 
       <div
-        v-else-if="props.task.trainingInformation.dataType === 'tabular'"
+        v-else-if="predictions.type === 'tabular'"
         class="mx-auto lg:w-3/4 h-full bg-white rounded-md max-h-128 overflow-x-scroll overflow-y-hidden"
       >
         <TableLayout
-          :columns="featuresNames"
-          :rows="dataWithPred.map((pred) => pred.data)"
+          :columns="predictions.labels.input.concat(predictions.labels.output)"
+          :rows="
+            predictions.results.map(({ input, output }) => input.push(output))
+          "
         />
       </div>
-
-      <div
-        v-else-if="props.task.trainingInformation.dataType === 'text'"
-        class="mx-auto lg:w-3/4 h-full bg-white rounded-md max-h-128 overflow-x-scroll overflow-y-hidden"
-      >
-        <!-- Display nothing for now -->
-      </div>
     </div>
   </div>
 </template>
 <script lang="ts" setup>
 import * as d3 from "d3";
 import createDebug from "debug";
-import { Set } from "immutable";
-import { storeToRefs } from "pinia";
-import { computed, ref, toRaw } from "vue";
+import { List, Range, Set } from "immutable";
+import { computed, ref, toRaw, watch } from "vue";
 
-import type {
-  Dataset,
-  Tabular,
-  Task,
-  Text,
-  TypedDataset,
-} from "@epfml/discojs";
-import { ConsoleLogger, EmptyMemory, Validator } from "@epfml/discojs";
-import { IndexedDB } from "@epfml/discojs-web";
+import type { Dataset, Model, Tabular, Task, Text } from "@epfml/discojs";
+import { Validator } from "@epfml/discojs";
 
 import InfoIcon from "@/assets/svg/InfoIcon.vue";
-import { useMemoryStore } from "@/store/memory";
 import { useToaster } from "@/composables/toaster";
 import { useValidationStore } from "@/store/validation";
 
@@ -142,25 +129,31 @@ import type {
   TypedNamedDataset,
 } from "@/components/dataset_input/types.js";
 
-const debug = createDebug("webapp:Tester");
-const { useIndexedDB } = storeToRefs(useMemoryStore());
+const debug = createDebug("webapp:testing:PredictSteps");
 const toaster = useToaster();
 const validationStore = useValidationStore();
 
 const props = defineProps<{
   task: Task;
-  modelID: string;
+  model: Model;
 }>();
 
-interface ImageWithUrl {
-  filename: string;
-  image: ImageData;
-}
-
-interface DataWithPrediction {
-  data: ImageWithUrl | number[];
-  prediction?: number;
-}
+type Results =
+  | {
+      type: "image";
+      results: List<{
+        input: {
+          filename: string;
+          image: ImageData;
+        };
+        output: string;
+      }>;
+    }
+  | {
+      type: "tabular";
+      labels: { input: List<string>; output: List<string> };
+      results: List<{ input: List<string>; output: string }>;
+    };
 
 const imageDataset = ref<NamedImageDataset>();
 const tabularDataset = ref<Dataset<Tabular>>();
@@ -185,53 +178,36 @@ const dataset = computed<TypedNamedDataset | undefined>(() => {
   return undefined;
 });
 
-const featuresNames = ref<string[]>([]);
-const dataWithPred = ref<DataWithPrediction[]>();
-const downloadLink = ref<HTMLAnchorElement>();
-
-const validator = computed(() => {
-  return new Validator(
-    props.task,
-    new ConsoleLogger(),
-    memory.value,
-    props.modelID,
-  );
-});
-const memory = computed(() =>
-  useIndexedDB ? new IndexedDB() : new EmptyMemory(),
-);
-
-type InferenceResults = Array<{ features: number[]; pred: number }>;
-const inferenceGenerator = ref<AsyncGenerator<InferenceResults, void>>();
-
-const visitedSamples = computed(() => validator.value?.visitedSamples ?? 0);
+const generator = ref<AsyncGenerator<number, void>>();
+const predictions = ref<Results>();
 
-function getLabelName(labelIndex: number | undefined): string {
-  if (labelIndex === undefined) return "";
+watch(tabularDataset, async (dataset) => {
+  if (dataset === undefined) return;
 
-  const labelList = props.task.trainingInformation.LABEL_LIST;
-  if (labelList !== undefined && labelList.length > labelIndex)
-    return labelList[labelIndex];
+  const { inputColumns } = props.task.trainingInformation;
+  if (inputColumns === undefined)
+    throw new Error("tabular task without input columns");
+  const wantedColumns = Set(inputColumns);
 
-  return labelIndex.toString();
-}
-
-function dropName(dataset: TypedNamedDataset): TypedDataset {
-  switch (dataset[0]) {
-    case "image":
-      return ["image", dataset[1].map(({ image }) => image)];
-    default:
-      return dataset;
-  }
-}
+  try {
+    for await (const [columns, i] of toRaw(dataset)
+      .map((row) => Set(Object.keys(row)))
+      .zip(Range(1)))
+      if (!columns.isSuperset(wantedColumns))
+        throw new Error(
+          `row ${i} is missing columns ${wantedColumns.subtract(columns).join(", ")}`,
+        );
+  } catch (e) {
+    let msg = "Error when loading CSV";
+    if (e instanceof Error) msg = `${msg}: ${e.message}`;
+    toaster.error(msg);
 
-async function modelInference(): Promise<void> {
-  const v = validator.value;
-  if (v === undefined) {
-    toaster.error("No model found");
+    tabularDataset.value = undefined;
     return;
   }
+});
 
+async function startInference(): Promise<void> {
   if (dataset.value === undefined) {
     toaster.error("Upload a dataset first");
     return;
@@ -239,112 +215,131 @@ async function modelInference(): Promise<void> {
 
   toaster.info("Model inference started");
   try {
-    const generator = await v.inference(dropName(toRaw(dataset.value)));
-    inferenceGenerator.value = generator;
-
-    let runningPredictions: InferenceResults = [];
-    for await (const predictions of generator) {
-      runningPredictions = runningPredictions.concat(predictions);
-      switch (dataset.value[0]) {
-        case "image":
-          dataWithPred.value = await arrayFromAsync(
-            toRaw(dataset.value[1])
-              .zip(runningPredictions)
-              .map(([{ filename, image }, prediction]) => ({
-                data: {
-                  filename,
-                  image: new ImageData(
-                    new Uint8ClampedArray(image.data),
-                    image.width,
-                    image.height,
-                  ),
-                },
-                prediction: prediction.pred,
-              })),
-          );
-          break;
-        case "tabular": {
-          const { inputColumns, outputColumns } =
-            props.task.trainingInformation;
-          if (inputColumns === undefined || outputColumns === undefined)
-            throw new Error("no input and output columns but CSV needs it");
-          featuresNames.value = [
-            ...inputColumns,
-            ...outputColumns.map((c) => `Predicted_${c}`),
-          ];
-          dataWithPred.value = runningPredictions.map((pred) => ({
-            data: [...pred.features, pred.pred],
-          }));
-          break;
-        }
-        case "text":
-          throw new Error("TODO implement");
-      }
+    switch (dataset.value[0]) {
+      case "image":
+        await startImageInference(toRaw(dataset.value)[1]);
+        break;
+      case "tabular":
+        await startTabularInference(toRaw(dataset.value)[1]);
+        break;
+      case "text":
+        throw new Error("disabled dataset type");
     }
+
     toaster.success("Model inference finished successfully!");
   } catch (e) {
     debug("while infering: %o", e);
     toaster.error("Something went wrong during model inference");
+  }
+}
+
+async function startImageInference(dataset: NamedImageDataset): Promise<void> {
+  const labels = List(props.task.trainingInformation.LABEL_LIST);
+  const validator = new Validator(props.task, toRaw(props.model));
+
+  let results: (Results & { type: "image" })["results"] = List();
+  try {
+    generator.value = await validator.infer([
+      "image",
+      dataset.map(({ image }) => image),
+    ]);
+    for await (const [{ filename, image }, prediction] of dataset.zip(
+      toRaw(generator.value),
+    )) {
+      results = results.push({
+        input: {
+          filename,
+          image: new ImageData(
+            new Uint8ClampedArray(image.data),
+            image.width,
+            image.height,
+          ),
+        },
+        output: labels.get(prediction) ?? prediction.toString(),
+      });
+
+      predictions.value = { type: "image", results };
+    }
   } finally {
-    inferenceGenerator.value = undefined;
+    generator.value = undefined;
+  }
+}
+
+async function startTabularInference(dataset: Dataset<Tabular>): Promise<void> {
+  const { inputColumns, outputColumns } = props.task.trainingInformation;
+  if (inputColumns === undefined || outputColumns === undefined)
+    throw new Error("no input and output columns but tabular task needs it");
+  const labels = {
+    input: List(inputColumns),
+    output: List(outputColumns).map((c) => `Predicted_${c}`),
+  };
+  const validator = new Validator(props.task, toRaw(props.model));
+
+  let results: (Results & { type: "tabular" })["results"] = List();
+  try {
+    generator.value = await validator.infer(["tabular", dataset]);
+    for await (const [input, prediction] of dataset.zip(
+      toRaw(generator.value),
+    )) {
+      results = results.push({
+        input: labels.input.map((label) => {
+          const ret = input[label];
+          if (ret === undefined)
+            throw new Error("row doesn't have expected label");
+          return ret;
+        }),
+        output: prediction.toString(),
+      });
+
+      predictions.value = { type: "tabular", labels, results };
+    }
+  } finally {
+    generator.value = undefined;
   }
 }
 
 async function stopInference(): Promise<void> {
-  const generator = inferenceGenerator.value;
-  if (generator === undefined) return;
+  const g = generator.value;
+  if (g === undefined) return;
 
-  generator.return();
+  generator.value = undefined;
+  g.return();
 }
 
 function saveCsv() {
-  if (downloadLink.value === undefined)
-    throw new Error("asked to download CSV but page yet rendered");
-  if (dataWithPred.value === undefined)
-    throw new Error("asked to save CSV without having training results");
-  if (dataset.value === undefined) {
-    toaster.error("Upload a dataset first");
-    return;
-  }
+  if (predictions.value === undefined)
+    throw new Error("asked to save CSV without predictions results");
 
-  const csvData = format(dataset.value[0], dataWithPred.value);
+  const csvData = format(predictions.value);
   const blob = new Blob([csvData], { type: "text/csv;charset=utf-8;" });
-  downloadLink.value.setAttribute("href", URL.createObjectURL(blob));
-  downloadLink.value.setAttribute(
+  const downloadLink = document.createElement("a");
+  downloadLink.setAttribute("href", URL.createObjectURL(blob));
+  downloadLink.setAttribute(
     "download",
     `predictions_${props.task.id}_${Date.now()}.csv`,
   );
-  downloadLink.value.click();
+  downloadLink.click();
 
-  function format(
-    datasetType: TypedDataset[0],
-    dataWithPred: DataWithPrediction[],
-  ): string {
-    switch (datasetType) {
+  function format(predictions: Results): string {
+    switch (predictions.type) {
       case "image": {
-        const rows = dataWithPred.map((el) => [
-          (el.data as ImageWithUrl).filename,
-          String(el.prediction),
+        return d3.csvFormatRows([
+          ["Filename", "Prediction"],
+          ...predictions.results.map(({ input, output }) => [
+            input.filename,
+            output,
+          ]),
         ]);
-        return d3.csvFormatRows([["Filename", "Prediction"], ...rows]);
       }
       case "tabular": {
-        const featureNames = Object.values(featuresNames.value) as string[];
-        const predictionsWithLabels = dataWithPred.map((el) =>
-          Object.values(el.data).map(String),
-        );
-        return d3.csvFormatRows([featureNames, ...predictionsWithLabels]);
+        return d3.csvFormatRows([
+          predictions.labels.input.concat(predictions.labels.output).toArray(),
+          ...predictions.results.map(({ input, output }) =>
+            input.push(output).toArray(),
+          ),
+        ]);
       }
-      case "text":
-        throw new Error("TODO implement formatter");
     }
   }
 }
-
-// Array.fromAsync not yet widely used (2024)
-async function arrayFromAsync<T>(iter: AsyncIterable<T>): Promise<T[]> {
-  const ret: T[] = [];
-  for await (const e of iter) ret.push(e);
-  return ret;
-}
 </script>
diff --git a/webapp/src/components/testing/TestSteps.vue b/webapp/src/components/testing/TestSteps.vue
index 076a2df03..88e6a20e5 100644
--- a/webapp/src/components/testing/TestSteps.vue
+++ b/webapp/src/components/testing/TestSteps.vue
@@ -17,7 +17,7 @@
     <LabeledImageDatasetInput
       v-if="task.trainingInformation.dataType === 'image'"
       v-model="imageDataset"
-      :labels="labels"
+      :labels="imageLabels"
     />
     <TabularDatasetInput
       v-if="task.trainingInformation.dataType === 'tabular'"
@@ -34,12 +34,12 @@
     <IconCard class="mx-auto mt-10 lg:w-1/2" title-placement="left">
       <template #title> Test &amp; validate your model </template>
 
-      <div v-if="testGenerator === undefined">
+      <div v-if="generator === undefined">
         By clicking the button below, you will be able to validate your model
         against a chosen dataset of yours. Below, once you assessed the model,
         you can compare the ground truth and the predicted values
         <div class="flex justify-center mt-4">
-          <CustomButton @click="testModel()"> test </CustomButton>
+          <CustomButton @click="startTest()"> test </CustomButton>
         </div>
       </div>
       <div v-else>
@@ -68,53 +68,48 @@
       </div>
     </div>
 
-    <div v-if="dataWithPred !== undefined">
+    <div v-if="tested !== undefined">
       <div class="mx-auto lg:w-1/2 text-center pb-8">
         <CustomButton @click="saveCsv()"> download as csv </CustomButton>
-        <a ref="downloadLink" class="hidden" />
       </div>
-      <!-- Image prediction gallery -->
-      <div v-if="dataset?.[0] === 'image'" class="grid grid-cols-6 gap-6">
+
+      <!-- Image gallery -->
+      <div v-if="tested.type === 'image'" class="grid grid-cols-6 gap-6">
         <ImageCard
-          v-for="(value, key) in dataWithPred"
+          v-for="(result, key) in tested.results"
           :key="key"
-          :image="(value.data as ImageWithUrl).image"
+          :image="result.input.image"
         >
           <template #title>
             <p
               class="font-bold"
-              :class="
-                value.prediction !== value.groundTruth
-                  ? 'text-red-700'
-                  : 'text-green-500'
-              "
-            >
-              {{ getLabelName(value.prediction) }}
-            </p>
-          </template>
-          <template #subtitle>
-            <p
-              v-if="value.groundTruth !== value.prediction"
-              class="text-disco-blue"
+              :class="result.output.correct ? 'text-red-700' : 'text-green-500'"
             >
-              <span class="font-bold">{{
-                getLabelName(value.groundTruth)
-              }}</span>
+              {{ result.output.truth }}
             </p>
           </template>
         </ImageCard>
       </div>
+
       <div
-        v-else-if="task.trainingInformation.dataType === 'tabular'"
+        v-else-if="tested.type === 'tabular'"
         class="mx-auto lg:w-3/4 h-full bg-white rounded-md max-h-128 overflow-x-scroll overflow-y-hidden"
       >
         <TableLayout
-          :columns="featuresNames"
-          :rows="dataWithPred.map((pred) => pred.data)"
+          :columns="
+            tested.labels.input
+              .concat(tested.labels.output.truth)
+              .push(tested.labels.output.correct)
+          "
+          :rows="
+            tested.results.map(({ input, output }) =>
+              input.concat(output.truth).push(output.correct.toString()),
+            )
+          "
         />
       </div>
       <div
-        v-else-if="task.trainingInformation.dataType === 'text'"
+        v-else-if="tested.type === 'text'"
         class="mx-auto lg:w-3/4 h-full bg-white rounded-md max-h-128 overflow-x-scroll overflow-y-hidden"
       >
         <!-- Display nothing for now -->
@@ -126,21 +121,11 @@
 import * as d3 from "d3";
 import createDebug from "debug";
 import { List, Range, Set } from "immutable";
-import { storeToRefs } from "pinia";
 import { computed, ref, toRaw, watch } from "vue";
 
-import type {
-  Dataset,
-  Tabular,
-  Task,
-  Text,
-  TypedDataset,
-  TypedLabeledDataset,
-} from "@epfml/discojs";
-import { ConsoleLogger, EmptyMemory, Validator } from "@epfml/discojs";
-import { IndexedDB } from "@epfml/discojs-web";
-
-import { useMemoryStore } from "@/store/memory";
+import type { Dataset, Model, Tabular, Task, Text } from "@epfml/discojs";
+import { Validator } from "@epfml/discojs";
+
 import { useToaster } from "@/composables/toaster";
 import { useValidationStore } from "@/store/validation";
 
@@ -158,28 +143,55 @@ import type {
   TypedNamedLabeledDataset,
 } from "@/components/dataset_input/types.js";
 
-const debug = createDebug("webapp:Tester");
-const { useIndexedDB } = storeToRefs(useMemoryStore());
+const debug = createDebug("webapp:testing:TestSteps");
 const toaster = useToaster();
 const validationStore = useValidationStore();
 
 const props = defineProps<{
   task: Task;
-  modelID: string;
+  model: Model;
 }>();
 
-interface ImageWithUrl {
-  filename: string;
-  image: ImageData;
-}
-
-interface DataWithPrediction {
-  data: ImageWithUrl | number[];
-  prediction?: number;
-  groundTruth?: number;
-}
-
-const labels = computed(() => Set(props.task.trainingInformation.LABEL_LIST));
+type Results =
+  | {
+      type: "image";
+      results: List<{
+        input: {
+          filename: string;
+          image: ImageData;
+        };
+        output: {
+          truth: string;
+          correct: boolean;
+        };
+      }>;
+    }
+  | {
+      type: "tabular";
+      labels: {
+        input: List<string>;
+        output: {
+          truth: List<string>;
+          correct: string;
+        };
+      };
+      results: List<{
+        input: List<string>;
+        output: {
+          truth: List<string>;
+          correct: boolean;
+        };
+      }>;
+    }
+  | {
+      type: "text";
+      // TODO what to show?
+      results: List<{ output: { correct: boolean } }>;
+    };
+
+const imageLabels = computed(() =>
+  Set(props.task.trainingInformation.LABEL_LIST),
+);
 const imageDataset = ref<NamedLabeledImageDataset>();
 const tabularDataset = ref<Dataset<Tabular>>();
 const textDataset = ref<Dataset<Text>>();
@@ -203,6 +215,30 @@ const dataset = computed<TypedNamedLabeledDataset | undefined>(() => {
   return undefined;
 });
 
+const tested = ref<Results>();
+const visitedSamples = computed(() => tested.value?.results.size ?? 0);
+const currentAccuracy = computed(() => {
+  if (tested.value === undefined) return "0";
+
+  let hits: number | undefined;
+  switch (tested.value.type) {
+    case "image":
+      hits = tested.value.results.filter(({ output }) => output.correct).size;
+      break;
+    case "tabular":
+      hits = tested.value.results.filter(({ output }) => output.correct).size;
+      break;
+    case "text":
+      hits = tested.value.results.filter(({ output }) => output.correct).size;
+      break;
+  }
+  const accuracy = hits / tested.value.results.size;
+
+  return (accuracy * 100).toFixed(2);
+});
+
+const generator = ref<AsyncGenerator<boolean, void>>();
+
 watch(tabularDataset, async (dataset) => {
   if (dataset === undefined) return;
 
@@ -229,61 +265,7 @@ watch(tabularDataset, async (dataset) => {
   }
 });
 
-const featuresNames = ref<string[]>([]);
-const dataWithPred = ref<DataWithPrediction[]>();
-const downloadLink = ref<HTMLAnchorElement>();
-
-const validator = computed(() => {
-  return new Validator(
-    props.task,
-    new ConsoleLogger(),
-    memory.value,
-    props.modelID,
-  );
-});
-const memory = computed(() =>
-  useIndexedDB ? new IndexedDB() : new EmptyMemory(),
-);
-
-type TestResults = Array<{
-  groundTruth: number;
-  pred: number;
-  features: number[];
-}>;
-const testGenerator = ref<AsyncGenerator<TestResults, void>>();
-
-const currentAccuracy = computed(() => {
-  const r = validator.value?.accuracy;
-  return r !== undefined ? (r * 100).toFixed(2) : "0";
-});
-const visitedSamples = computed(() => validator.value?.visitedSamples ?? 0);
-
-function getLabelName(labelIndex: number | undefined): string {
-  if (labelIndex === undefined) return "";
-
-  const labelList = props.task.trainingInformation.LABEL_LIST;
-  if (labelList !== undefined && labelList.length > labelIndex)
-    return labelList[labelIndex];
-
-  return labelIndex.toString();
-}
-
-function dropName(dataset: TypedNamedLabeledDataset): TypedLabeledDataset {
-  switch (dataset[0]) {
-    case "image":
-      return ["image", dataset[1].map(({ image, label }) => [image, label])];
-    default:
-      return dataset;
-  }
-}
-
-async function testModel(): Promise<void> {
-  const v = validator.value;
-  if (v === undefined) {
-    toaster.error("No model found");
-    return;
-  }
-
+async function startTest(): Promise<void> {
   if (dataset.value === undefined) {
     toaster.error("Upload a dataset first");
     return;
@@ -291,120 +273,170 @@ async function testModel(): Promise<void> {
 
   toaster.info("Model testing started");
   try {
-    const generator = v.test(dropName(toRaw(dataset.value)));
-    testGenerator.value = generator;
-    let runningResults: TestResults = [];
-    for await (const assessmentResults of generator) {
-      runningResults = runningResults.concat(assessmentResults);
-      switch (dataset.value[0]) {
-        case "image":
-          dataWithPred.value = List(
-            await arrayFromAsync(toRaw(dataset.value[1])),
-          )
-            .zip(List(runningResults))
-            .map(([{ filename, image }, prediction]) => ({
-              data: {
-                filename,
-                image: new ImageData(
-                  new Uint8ClampedArray(image.data),
-                  image.width,
-                  image.height,
-                ),
-              },
-              prediction: prediction.pred,
-              groundTruth: prediction.groundTruth,
-            }))
-            .toArray();
-          break;
-        case "tabular": {
-          const { inputColumns, outputColumns } =
-            props.task.trainingInformation;
-          if (inputColumns === undefined || outputColumns === undefined)
-            throw new Error("no input and output columns but CSV needs it");
-          featuresNames.value = [
-            ...inputColumns,
-            ...outputColumns.map((c) => `Predicted_${c}`),
-            ...outputColumns.map((c) => `Target_${c}`),
-          ];
-          dataWithPred.value = runningResults.map((pred) => ({
-            data: [...pred.features, pred.pred, pred.groundTruth],
-          }));
-          break;
-        }
-        case "text":
-          // TODO nothing to show for now
-          dataWithPred.value = [];
-      }
+    switch (dataset.value[0]) {
+      case "image":
+        await startImageTest(toRaw(dataset.value)[1]);
+        break;
+      case "tabular":
+        await startTabularTest(toRaw(dataset.value)[1]);
+        break;
+      case "text":
+        startTextTest(toRaw(dataset.value)[1]);
+        break;
     }
 
     toaster.success("Model testing finished successfully!");
   } catch (e) {
     debug("while testing: %o", e);
     toaster.error("Something went wrong during model testing");
+  }
+}
+
+async function startImageTest(
+  dataset: NamedLabeledImageDataset,
+): Promise<void> {
+  const validator = new Validator(props.task, toRaw(props.model));
+  let results: (Results & { type: "image" })["results"] = List();
+
+  try {
+    generator.value = validator.test([
+      "image",
+      dataset.map(({ image, label }) => [image, label]),
+    ]);
+    for await (const [{ filename, image, label }, correct] of dataset.zip(
+      toRaw(generator.value),
+    )) {
+      results = results.push({
+        input: {
+          filename,
+          image: new ImageData(
+            new Uint8ClampedArray(image.data),
+            image.width,
+            image.height,
+          ),
+        },
+        output: {
+          truth: label,
+          correct,
+        },
+      });
+
+      tested.value = { type: "image", results };
+    }
+  } finally {
+    generator.value = undefined;
+  }
+}
+
+async function startTabularTest(dataset: Dataset<Tabular>): Promise<void> {
+  const { inputColumns, outputColumns } = props.task.trainingInformation;
+  if (inputColumns === undefined || outputColumns === undefined)
+    throw new Error("no input and output columns but CSV needs it");
+  const labels = {
+    input: List(inputColumns),
+    output: {
+      truth: List(outputColumns).map((c) => `Truth_${c}`),
+      correct: "Correct",
+    },
+  };
+  const validator = new Validator(props.task, toRaw(props.model));
+
+  let results: (Results & { type: "tabular" })["results"] = List();
+  try {
+    generator.value = validator.test(["tabular", dataset]);
+    for await (const [row, correct] of dataset.zip(toRaw(generator.value))) {
+      // TODO we only really support a single output, change Task to reflect that
+      const truth = List(outputColumns).map((column) => {
+        const ret = row[column];
+        if (ret === undefined)
+          throw new Error("row doesn't have expected output column");
+        return ret;
+      });
+
+      results = results.push({
+        input: labels.input.map((label) => {
+          const ret = row[label];
+          if (ret === undefined)
+            throw new Error("row doesn't have expected label");
+          return ret;
+        }),
+        output: {
+          truth,
+          correct,
+        },
+      });
+
+      tested.value = { type: "tabular", labels, results };
+    }
+  } finally {
+    generator.value = undefined;
+  }
+}
+
+async function startTextTest(dataset: Dataset<Text>): Promise<void> {
+  const validator = new Validator(props.task, toRaw(props.model));
+  let results: (Results & { type: "text" })["results"] = List();
+
+  try {
+    generator.value = validator.test(["text", dataset]);
+    for await (const correct of toRaw(generator.value)) {
+      results = results.push({ output: { correct } });
+      tested.value = { type: "text", results };
+    }
   } finally {
-    testGenerator.value = undefined;
+    generator.value = undefined;
   }
 }
 
 async function stopTest(): Promise<void> {
-  const generator = testGenerator.value;
-  if (generator === undefined) return;
+  const g = generator.value;
+  if (g === undefined) return;
 
-  generator.return();
+  generator.value = undefined;
+  g.return();
 }
 
-function saveCsv() {
-  if (downloadLink.value === undefined)
-    throw new Error("asked to download CSV but page yet rendered");
-  if (dataWithPred.value === undefined)
-    throw new Error("asked to save CSV without having training results");
-  if (dataset.value === undefined) {
-    toaster.error("Upload a dataset first");
-    return;
-  }
+function saveCsv(): void {
+  if (tested.value === undefined)
+    throw new Error("asked to save CSV without testing results");
 
-  const csvData = format(dataset.value[0], dataWithPred.value);
+  const csvData = format(tested.value);
   const blob = new Blob([csvData], { type: "text/csv;charset=utf-8;" });
-  downloadLink.value.setAttribute("href", URL.createObjectURL(blob));
-  downloadLink.value.setAttribute(
+  const downloadLink = document.createElement("a");
+  downloadLink.setAttribute("href", URL.createObjectURL(blob));
+  downloadLink.setAttribute(
     "download",
-    `predictions_${props.task.id}_${Date.now()}.csv`,
+    `tests_${props.task.id}_${Date.now()}.csv`,
   );
-  downloadLink.value.click();
-
-  function format(
-    datasetType: TypedDataset[0],
-    dataWithPred: DataWithPrediction[],
-  ): string {
-    switch (datasetType) {
-      case "image": {
-        const rows = dataWithPred.map((el) => [
-          (el.data as ImageWithUrl).filename,
-          String(el.prediction),
-          String(el.groundTruth),
+  downloadLink.click();
+
+  function format(tested: Results): string {
+    switch (tested.type) {
+      case "image":
+        return d3.csvFormatRows([
+          ["Filename", "Truth", "Correct"],
+          ...tested.results.map(({ input, output }) => [
+            input.filename,
+            output.truth,
+            output.correct.toString(),
+          ]),
         ]);
+      case "tabular":
         return d3.csvFormatRows([
-          ["Filename", "Prediction", "Ground Truth"],
-          ...rows,
+          tested.labels.input
+            .concat(tested.labels.output.truth)
+            .push(tested.labels.output.correct)
+            .toArray(),
+          ...tested.results.map((result) =>
+            result.input
+              .concat(result.output.truth)
+              .push(result.output.correct.toString())
+              .toArray(),
+          ),
         ]);
-      }
-      case "tabular": {
-        const featureNames = Object.values(featuresNames.value) as string[];
-        const predictionsWithLabels = dataWithPred.map((el) =>
-          Object.values(el.data).map(String),
-        );
-        return d3.csvFormatRows([featureNames, ...predictionsWithLabels]);
-      }
-      case "text":
-        throw new Error("TODO implement formatter");
     }
-  }
-}
 
-// Array.fromAsync not yet widely used (2024)
-async function arrayFromAsync<T>(iter: AsyncIterable<T>): Promise<T[]> {
-  const ret: T[] = [];
-  for await (const e of iter) ret.push(e);
-  return ret;
+    throw new Error("TODO implement formatter");
+  }
 }
 </script>
diff --git a/webapp/src/components/testing/Testing.vue b/webapp/src/components/testing/Testing.vue
index 61e4b8ac5..1840f42db 100644
--- a/webapp/src/components/testing/Testing.vue
+++ b/webapp/src/components/testing/Testing.vue
@@ -24,8 +24,8 @@
               <ButtonsCard
                 :buttons="
                   List.of(
-                    ['test', () => selectModel(id, false)],
-                    ['predict', () => selectModel(id, true)],
+                    ['test', () => selectModel(id, 'test')],
+                    ['predict', () => selectModel(id, 'predict')],
                   )
                 "
                 class="shadow shadow-disco-cyan"
@@ -135,14 +135,14 @@
   <div v-if="selection !== undefined">
     <div v-if="validationStore.step !== 0">
       <TestSteps
-        v-if="selection.type === 'test'"
+        v-if="selection.mode === 'test'"
         :task="selection.task"
-        :modelID="selection.modelID"
+        :model="selection.model"
       />
       <PredictSteps
-        v-if="selection.type === 'predict'"
+        v-if="selection.mode === 'predict'"
         :task="selection.task"
-        :modelID="selection.modelID"
+        :model="selection.model"
       />
     </div>
   </div>
@@ -157,7 +157,7 @@ import { computed, ref, onActivated } from "vue";
 import { RouterLink } from "vue-router";
 import { VueSpinner } from "vue3-spinners";
 
-import type { Task } from "@epfml/discojs";
+import type { Model, Task } from "@epfml/discojs";
 import { EmptyMemory, client as clients, aggregator } from "@epfml/discojs";
 import { IndexedDB } from "@epfml/discojs-web";
 
@@ -183,10 +183,10 @@ const tasksStore = useTasksStore();
 const toaster = useToaster();
 
 const selection = ref<{
-  type: "predict" | "test";
+  mode: "predict" | "test";
   task: Task;
   // same as in validation store but not undef
-  modelID: string;
+  model: Model;
 }>();
 
 const federatedTasks = computed(() =>
@@ -204,7 +204,7 @@ onActivated(async () => {
   // handle test after training or from library
   // TODO encode model ID inside the URL instead of relying on store
   if (validationStore.modelID !== undefined)
-    selectModel(validationStore.modelID, false);
+    selectModel(validationStore.modelID, "test");
 });
 
 async function downloadModel(task: Task): Promise<void> {
@@ -243,25 +243,24 @@ async function downloadModel(task: Task): Promise<void> {
     });
 }
 
-function selectModel(modelID: string, isOnlyPrediction: boolean): void {
+async function selectModel(
+  modelID: string,
+  mode: "predict" | "test",
+): Promise<void> {
   const taskID = memory.value.getModelInfo(modelID)?.taskID;
-  if (taskID === undefined)
-    throw new Error(
-      `Model task id for found in memory for model id: ${modelID}`,
-    );
+  if (taskID === undefined) throw new Error("task ID for model ID not found");
 
   const task = tasksStore.tasks.get(taskID);
-  if (task === undefined)
-    throw new Error("Task not found in the task store for task id: " + taskID);
+  if (task === undefined) throw new Error("task not found");
+
+  if (!(await memory.value.contains(modelID)))
+    throw new Error("model ID not present in memory");
+  const model = await memory.value.getModel(modelID);
 
-  selection.value = {
-    type: isOnlyPrediction ? "predict" : "test",
-    task,
-    modelID,
-  };
+  selection.value = { mode, model, task };
+  validationStore.mode = mode;
+  validationStore.modelID = modelID;
   validationStore.step = 1;
-  validationStore.modelID = selection.value.modelID;
-  validationStore.mode = selection.value.type;
 }
 
 function taskTitle(taskID: string): string | undefined {