import * as tf from '@tensorflow/tfjs-node';
import { getModel, setBackend, } from '@jsonstack/jsonm';

//set tensorflow
setBackend(tf);

//Iris Dataset e.g from https://raw.githubusercontent.com/repetere/modelx-model/master/src/test/mock/data/iris_data.csv
const type = 'ai-classification';
const dataset = [
  {
    "sepal_length_cm": 5.1,
    "sepal_width_cm": 3.5,
    "petal_length_cm": 1.4,
    "petal_width_cm": 0.2,
    "plant": "Iris-setosa",
  },
//  ...
  {
    "sepal_length_cm": 7.0,
    "sepal_width_cm": 3.2,
    "petal_length_cm": 4.7,
    "petal_width_cm": 1.4,
    "plant": "Iris-versicolor",
  },
  // ...
  {
    "sepal_length_cm": 5.9,
    "sepal_width_cm": 3.0,
    "petal_length_cm": 5.1,
    "petal_width_cm": 1.8,
    "plant": "virginica",
  }
]
const inputs = ['sepal_length_cm','sepal_width_cm','petal_length_cm','petal_width_cm', ];
const outputs = [ 'plant',];
const on_progress = ({ completion_percentage, loss, epoch, status, logs, defaultLog, }) => { 
  console.log({ completion_percentage, loss, epoch, status, logs, defaultLog, });
}
const IrisModel = await getModel({
  type,
  dataset,
  inputs,
  outputs,
  on_progress,
}); 
await IrisModel.trainModel()
const predictions = await IrisModel.predictModel({ 
  prediction_inputs:[
    { sepal_length_cm: 5.1, sepal_width_cm: 3.5, petal_length_cm: 1.4, petal_width_cm: 0.2, },
    { sepal_length_cm: 5.9, sepal_width_cm: 3.0, petal_length_cm: 5.1, petal_width_cm: 1.8, },
  ],
}); // => [ { plant:'Iris-setosa' }, { plant:'Iris-virginica' }, ]