Add tabular preprocessing

discojs/src/default_tasks/titanic.ts

@@ -46,8 +46,14 @@ export const titanic: TaskProvider<"tabular", "federated"> = {
           'SibSp',
           'Parch',
           'Fare',
-          'Pclass'
+          'Pclass',
+          'Sex',
+          'Embarked'
         ],
+        categoricalColumns: {
+          Sex: ["male", "female"],
+          Embarked: ["C", "S", "Q", "Missing"]
+        },
         outputColumn: 'Survived',
         scheme: 'federated',
         aggregationStrategy: 'mean',
@@ -62,7 +68,7 @@ export const titanic: TaskProvider<"tabular", "federated"> = {
 
     model.add(
       tf.layers.dense({
-        inputShape: [5],
+        inputShape: [11],
         units: 124,
         activation: 'relu',
         kernelInitializer: 'leCunNormal'

discojs/src/index.ts

@@ -17,6 +17,7 @@ export {
   EpochLogs,
   Tokenizer,
   ValidationMetrics,
+  ModelMetadata,
 } from "./models/index.js";
 export * as models from './models/index.js'
 

discojs/src/models/index.ts

@@ -1,4 +1,4 @@
-export { Model } from './model.js'
+export { Model, ModelMetadata } from './model.js'
 export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js";
 export { Tokenizer } from "./tokenizer.js";
 

discojs/src/models/model.ts

@@ -7,6 +7,11 @@ import type {
 } from "../index.js";
 
 import type { BatchLogs, EpochLogs } from "./logs.js";
+import type { StandardizationStats } from "../processing/tabular.js";
+
+export type ModelMetadata = {
+  tabularStandardization?: StandardizationStats;
+};
 
 /**
  * Trainable predictor
@@ -21,6 +26,9 @@ export abstract class Model<D extends DataType> implements Disposable {
   /** Set training state */
   abstract set weights(ws: WeightsContainer);
 
+  /** Optional metadata for tabular task data standardization */
+  metadata?: ModelMetadata;
+
   /**
    * Improve predictor
    *

discojs/src/models/tfjs.ts

@@ -12,17 +12,20 @@ import {
 import { BatchLogs } from './index.js'
 import { Model } from './index.js'
 import { EpochLogs } from './logs.js'
+import { ModelMetadata } from "./model.js";
 
-type Serialized<D extends DataType> = [D, tf.io.ModelArtifacts];
+type Serialized<D extends DataType> = [D, tf.io.ModelArtifacts, ModelMetadata?];
 
 /** TensorFlow JavaScript model with standard training */
 export class TFJS<D extends "image" | "tabular"> extends Model<D> {
   /** Wrap the given trainable model */
   constructor (
     public readonly datatype: D,
-    private readonly model: tf.LayersModel
+    private readonly model: tf.LayersModel,
+    metadata?: ModelMetadata,
   ) {
     super()
+    this.metadata = metadata;
 
     if (model.loss === undefined) {
       throw new Error('TFJS models need to be compiled to be used')
@@ -176,12 +179,14 @@ export class TFJS<D extends "image" | "tabular"> extends Model<D> {
   static async deserialize<D extends "image" | "tabular">([
     datatype,
     artifacts,
+    metadata
   ]: Serialized<D>): Promise<TFJS<D>> {
     return new this(
       datatype,
       await tf.loadLayersModel({
         load: () => Promise.resolve(artifacts),
       }),
+      metadata
     );
   }
 
@@ -204,7 +209,7 @@ export class TFJS<D extends "image" | "tabular"> extends Model<D> {
       includeOptimizer: true // keep model compiled
     })
 
-    return [this.datatype, await ret]
+    return [this.datatype, await ret, this.metadata]
   }
 
   [Symbol.dispose](): void{

discojs/src/processing/index.spec.ts

@@ -23,6 +23,7 @@ describe("preprocess", () => {
         batchSize: 1,
         validationSplit: 0,
         inputColumns: ["a", "b"],
+        categoricalColumns: {},
         outputColumn: "c",
       },
     };

discojs/src/processing/index.ts

@@ -9,6 +9,7 @@ import type {
 	Tabular,
 	Task,
 	Network,
+  ModelMetadata,
 } from "../index.js";
 
 import * as processing from "./index.js";
@@ -19,6 +20,7 @@ export * from "./tabular.js";
 export function preprocess<D extends DataType, N extends Network>(
   task: Task<D, N>,
   dataset: Dataset<DataFormat.Raw[D]>,
+  metadata?: ModelMetadata,
 ): Dataset<DataFormat.ModelEncoded[D]> {
   switch (task.dataType) {
     case "image": {
@@ -36,13 +38,18 @@ export function preprocess<D extends DataType, N extends Network>(
     case "tabular": {
       // cast as typescript doesn't reduce generic type
       const d = dataset as Dataset<DataFormat.Raw["tabular"]>;
-      const { inputColumns, outputColumn } = task.trainingInformation;
+      const { inputColumns, outputColumn, categoricalColumns } = task.trainingInformation;
+      const stats = metadata?.tabularStandardization;
 
       return d.map((row) => {
         const output = processing.extractColumn(row, outputColumn);
 
+        const inputs = List(
+          processing.encodeTabularRow(row, inputColumns, categoricalColumns, stats)
+        );
+
         return [
-          extractToNumbers(inputColumns, row),
+          inputs,
           // TODO sanitization doesn't care about column distribution
           output !== "" ? processing.convertToNumber(output) : 0,
         ];
@@ -68,6 +75,7 @@ export function preprocess<D extends DataType, N extends Network>(
 export function preprocessWithoutLabel<D extends DataType>(
   task: Task<D, Network>,
   dataset: Dataset<DataFormat.RawWithoutLabel[D]>,
+  metadata?: ModelMetadata,
 ): Dataset<DataFormat.ModelEncoded[D][0]> {
   switch (task.dataType) {
     case "image": {
@@ -84,9 +92,14 @@ export function preprocessWithoutLabel<D extends DataType>(
     case "tabular": {
       // cast as typescript doesn't reduce generic type
       const d = dataset as Dataset<DataFormat.Raw["tabular"]>;
-      const { inputColumns } = task.trainingInformation;
+      const { inputColumns, categoricalColumns } = task.trainingInformation;
+      const stats = metadata?.tabularStandardization;
 
-      return d.map((row) => extractToNumbers(inputColumns, row));
+      return d.map((row) => 
+        List(
+          processing.encodeTabularRow(row, inputColumns, categoricalColumns, stats)
+        )
+      );
     }
     case "text": {
       // cast as typescript doesn't reduce generic type

discojs/src/processing/tabular.ts

@@ -1,5 +1,10 @@
 import { List } from "immutable";
 
+export type StandardizationStats = {
+  means: Record<string, number>;
+  stds: Record<string, number>;
+};
+
 /**
  * Convert a string to a number
  *
@@ -38,3 +43,104 @@ export function indexInList(
   if (ret === -1) throw new Error(`${element} not found in list`);
   return ret;
 }
+
+/**
+ * Return the mean, std value of each column
+ */
+export function computeStandardizationStats(
+  rows: Array<Partial<Record<string, string>>>,
+  columns: Array<string>,
+): StandardizationStats{
+  const means: Record<string, number> = {};
+  const stds: Record<string, number> = {};
+
+  for (const col of columns){
+    const values = rows.map((row)=> {
+      const rawValue = extractColumn(row, col);
+      return convertToNumber(rawValue !== "" ? rawValue : "0");
+    });
+    const mean = values.reduce((a, b)=> a+b, 0) / values.length;
+    const variance = values.reduce((acc, val) => acc + (val-mean)**2, 0) / values.length;
+
+    const std = Math.sqrt(variance);
+
+    means[col] = mean;
+    stds[col] = std;
+  }
+
+  return {means, stds};
+}
+
+/**
+ * Apply standardization for a single value
+ */
+export function standardizeValue(
+  value: number,
+  mean: number,
+  std: number,
+): number{
+  if (std == 0) return 0; // avoid divide by 0
+  return (value - mean) / std;
+}
+
+/**
+ * Apply one hot encoding for a row
+ * 
+ * One hot encoding function is called for each row in dataset
+ */
+export function oneHotEncode(
+  value: string,
+  categories: Array<string>,
+): Array<number> {
+  // Get the index of the value among the possible categories
+  const index = categories.indexOf(value);
+
+  // If the value does not exist, raise an error
+  if (index === -1) {
+    throw new Error(`"${value}" is not a valid category for this column`);
+  }
+
+  return categories.map((_, categoryIndex) => 
+    categoryIndex === index ? 1 : 0
+  );
+}
+
+/**
+ * Apply standardization for numerical columns and 
+ * apply one hot encoding for categorical columns and return the final row
+ */
+export function encodeTabularRow(
+  row: Partial<Record<string, string>>,
+  inputColumns: Array<string>,
+  categoricalColumns: Record<string, Array<string>>,
+  stats?: StandardizationStats,
+): Array<number> {
+  const outputRow = inputColumns.flatMap((column) => {
+    const raw = extractColumn(row, column);
+    const categories = categoricalColumns[column];
+
+    // If the column exists in the list of categorical columns, apply one hot encoding
+    if (categories !== undefined){
+      return oneHotEncode(raw, categories);
+    }
+
+    // If the column is numerical column, apply standardization
+    const value = convertToNumber(raw !== "" ? raw : "0");
+
+    if (stats === undefined) {
+      return [value];
+    }
+
+    const mean = stats.means[column];
+    const std = stats.stds[column];
+
+    // Raise an error when stats is not defined
+    if (mean === undefined || std === undefined){
+      throw new Error(`Standardization statistics is not defined for column ${column}`);
+    }
+
+    return [standardizeValue(value, mean, std)];
+  });
+
+  return outputRow;
+}

discojs/src/serialization/model.ts

@@ -1,6 +1,6 @@
 import type tf from '@tensorflow/tfjs'
 
-import type { DataType, Model } from '../index.js'
+import type { DataType, Model, ModelMetadata } from '../index.js'
 import { models, serialization } from '../index.js'
 import { GPTConfig } from '../models/index.js'
 
@@ -41,11 +41,11 @@ export async function decode(encoded: Encoded): Promise<Model<DataType>> {
   const rawModel = raw[1] as unknown
   switch (type) {
     case Type.TFJS: {
-      if (raw.length !== 3)
+      if (raw.length !== 3 && raw.length !== 4)
         throw new Error(
-          "invalid TFJS model encoding: should be an array of length 3",
+          "invalid TFJS model encoding: should be an array of length 3 or 4",
         );
-      const [rawDatatype, rawModel] = raw.slice(1) as unknown[];
+      const [rawDatatype, rawModel, rawMetadata] = raw.slice(1) as unknown[];
 
       let datatype;
       switch (rawDatatype) {
@@ -63,6 +63,8 @@ export async function decode(encoded: Encoded): Promise<Model<DataType>> {
         datatype,
         // TODO totally unsafe casting
         rawModel as tf.io.ModelArtifacts,
+        // metadata for tabular task standardization
+        rawMetadata as ModelMetadata,
       ]);
     }
     case Type.GPT: {  

discojs/src/task/training_information.ts

@@ -83,6 +83,8 @@ export namespace TrainingInformation {
 		tabular: z.object({
 			// the columns to be chosen as input data for the model
 			inputColumns: z.array(z.string()),
+			// categorical columns to be chosen as input data for the model
+			categoricalColumns: z.record(z.string(), z.array(z.string()).min(1)).optional().default({}),
 			// the columns to be predicted by the model
 			outputColumn: z.string(),
 		}),