Complete Apps Script Google Sheets code for ml-classify-text-js

Code.js

// Import ml-classify-text
eval(UrlFetchApp.fetch('https://cdn.jsdelivr.net/npm/[email protected]/lib/index.js').getContentText());
let classifier = new Classifier()

// Add menu items
function onOpen() {
  var ui = SpreadsheetApp.getUi();
  ui.createMenu('ML Train Menu')
    .addItem('Get training set', 'captureRange')
    .addItem('Predict', 'predictRange')
    .addToUi();
}

// Copy range data for training/prediction later
function captureRange(e) {
  var sheet = SpreadsheetApp.getActiveSpreadsheet().getActiveSheet();
  let range = this.resolveRange(sheet)
  var data = sheet.getRange(range).getValues();
  let processed = this.processData(data)
  this.setTrainData(processed)

  Browser.msgBox('Captured training set!')
}

// Predict click
function predictRange() {

  // Get saved data and train it
  let processed = this.getTrainData()
  this.trainModel(processed)

  // Get selected range and data
  let sheet = SpreadsheetApp.getActiveSpreadsheet().getActiveSheet();
  let resolved = this.resolveRange(sheet)
  let range = sheet.getRange(resolved)
  var data = range.getValues();

  // Predict
  let predicted = []
  for (let [text, prediction] of data) {
    let predictions = classifier.predict(text);
    if (predictions.length) {
      prediction = predictions[0].label
    }
    predicted.push([text, prediction]);
  }

  // Write out
  range.setValues(predicted)

}

// Deal with column vs range select
function resolveRange(sheet) {
  let range = sheet.getActiveRange().getA1Notation();
  let [first, second] = range.split(':')
  if (!/\d/.test(second)) {
    range = `${first}1:${second}${sheet.getLastRow()}`
  }

  return range
}

// Formats data into key:array
function processData(arrayData) {
  let processed = {};

  for (let [i, [text, label]] of arrayData.entries()) {
    // TODO: Deal with headers
    if (i == 0 && this.excludeFirstRow) continue;
    // Do not add empty rows
    if (text.trim() == '' || label.trim() == '') continue;

    //console.log(`Processing: ${text}, ${label}`);
    if (label in processed) {
      processed[label].push(text);
    } else {
      processed[label] = [text];
    }
  }

  return processed;
}

function setTrainData(trainData) {
  PropertiesService.getScriptProperties().setProperty('trainData', JSON.stringify(trainData));
}

function getTrainData() {
  return JSON.parse(PropertiesService.getScriptProperties().getProperty('trainData'));
}

function trainModel(data) {
  classifier = new Classifier({
    nGramMin: 1,
    nGramMax: 3,
  });

  for (let key in data) {
    classifier.train(data[key], key);
  }
}