translate.js

//after npm or yarn install
import * as tf from '@tensorflow/tfjs';
//loading pretrained model
import * as loader from './loader';
//sets up basic dom elements
import * as ui from './ui';

//Load the pretrained models , versioning information, timestamps,
const HOSTED_URLS = {
  model:
      'https://storage.googleapis.com/tfjs-models/tfjs/translation_en_fr_v1/model.json',
  metadata:
      'https://storage.googleapis.com/tfjs-models/tfjs/translation_en_fr_v1/metadata.json'
};
//we have our models locally stored locally too! as JSON
const LOCAL_URLS = {
  model: 'http://localhost:1235/resources/model.json',
  metadata: 'http://localhost:1235/resources/metadata.json'
};

//translator class
class Translator {
  /**
   * Initializes the Translation demo.
   */
  //no waitiing for DOM elemnts to load
  async init(urls) {
    //init using our URL of choice (remote or local)
    this.urls = urls;
    //lets get our model!
    const model = await loader.loadHostedPretrainedModel(urls.model);
    //The await expression causes async function execution to pause until a Promise is resolved, 
    //that is fulfilled or rejected, and to resume execution of the async function after fulfillment
    await this.loadMetadata();
    //segment into encoder and decoder models
    this.prepareEncoderModel(model);
    this.prepareDecoderModel(model);
    //return both as our ready to go models
    return this;
  }

  async loadMetadata() {
    //retrieve from helper class
    const translationMetadata =
        await loader.loadHostedMetadata(this.urls.metadata);
    //max sequence lengths for encoder and decoder
    //exit conditions, either hit max length or find stop character
    this.maxDecoderSeqLength = translationMetadata['max_decoder_seq_length'];
    this.maxEncoderSeqLength = translationMetadata['max_encoder_seq_length'];
    //print them
    console.log('maxDecoderSeqLength = ' + this.maxDecoderSeqLength);
    console.log('maxEncoderSeqLength = ' + this.maxEncoderSeqLength);
    //index of input and target variables
    //iterating 
    this.inputTokenIndex = translationMetadata['input_token_index'];
    this.targetTokenIndex = translationMetadata['target_token_index'];
    //for generating next chars
    //retrieve a single value
    this.reverseTargetCharIndex =
        Object.keys(this.targetTokenIndex)
            .reduce(
                (obj, key) => (obj[this.targetTokenIndex[key]] = key, obj), {});
  }

  prepareEncoderModel(model) {
    //how many encoder tokens?
    this.numEncoderTokens = model.input[0].shape[2];
    //how many inputs?
    const encoderInputs = model.input[0];
    //hidden state
    const stateH = model.layers[2].output[1];
    //hidden state
    const stateC = model.layers[2].output[2];
    //hidden states
    const encoderStates = [stateH, stateC];
  //build model using inputs and hidden states
    this.encoderModel =
        tf.model({inputs: encoderInputs, outputs: encoderStates});
  }

  prepareDecoderModel(model) {
    //decoder token count
    this.numDecoderTokens = model.input[1].shape[2];
    console.log('numDecoderTokens = ' + this.numDecoderTokens);

    //hidden state to help define input
    const stateH = model.layers[2].output[1];
    //helps define hidden state 2's input
    const latentDim = stateH.shape[stateH.shape.length - 1];
    console.log('latentDim = ' + latentDim);
    
    //define inputs directly into both hidden states
    const decoderStateInputH =
        tf.input({shape: [latentDim], name: 'decoder_state_input_h'});
    const decoderStateInputC =
        tf.input({shape: [latentDim], name: 'decoder_state_input_c'});\
    //here are our hidden state inputs!
    const decoderStateInputs = [decoderStateInputH, decoderStateInputC];

    //retrieve our LSTM model!
    const decoderLSTM = model.layers[3];
    //our initial input
    const decoderInputs = decoderLSTM.input[0];
    //initialize model using first input, and hidden state inputs
    const applyOutputs =
        decoderLSTM.apply(decoderInputs, {initialState: decoderStateInputs});
    //last output (not activated)
    let decoderOutputs = applyOutputs[0];
    //we have both our hidden states! They are outputs of our inputs
    const decoderStateH = applyOutputs[1];
    const decoderStateC = applyOutputs[2];
    //hidden states get their own list
    const decoderStates = [decoderStateH, decoderStateC];

    //apply fully connected layer
    const decoderDense = model.layers[4];
    //get final output
    decoderOutputs = decoderDense.apply(decoderOutputs);
    //define our model using our initial and state inputs concatented
    //with our final and hidden state outputs concatented
    this.decoderModel = tf.model({
      inputs: [decoderInputs].concat(decoderStateInputs),
      outputs: [decoderOutputs].concat(decoderStates)
    });
  }

  /**
   * Encode a string (e.g., a sentence) as a Tensor3D that can be fed directly
   * into the TensorFlow.js model.
   */
  encodeString(str) {
    const strLen = str.length;
    const encoded =
        tf.buffer([1, this.maxEncoderSeqLength, this.numEncoderTokens]);
    for (let i = 0; i < strLen; ++i) {
      if (i >= this.maxEncoderSeqLength) {
        console.error(
            'Input sentence exceeds maximum encoder sequence length: ' +
            this.maxEncoderSeqLength);
      }

      const tokenIndex = this.inputTokenIndex[str[i]];
      if (tokenIndex == null) {
        console.error(
            'Character not found in input token index: "' + tokenIndex + '"');
      }
      encoded.set(1, 0, i, tokenIndex);
    }
    return encoded.toTensor();
  }

  decodeSequence(inputSeq) {
    // Encode the inputs state vectors.
    let statesValue = this.encoderModel.predict(inputSeq);

    // Generate empty target sequence of length 1.
    let targetSeq = tf.buffer([1, 1, this.numDecoderTokens]);
    // Populate the first character of the target sequence with the start
    // character.
    targetSeq.set(1, 0, 0, this.targetTokenIndex['\t']);

    // Sample loop for a batch of sequences.
    // (to simplify, here we assume that a batch of size 1).
    let stopCondition = false;
    let decodedSentence = '';
    while (!stopCondition) {
      const predictOutputs =
          this.decoderModel.predict([targetSeq.toTensor()].concat(statesValue));
      const outputTokens = predictOutputs[0];
      const h = predictOutputs[1];
      const c = predictOutputs[2];

      // Sample a token.
      // We know that outputTokens.shape is [1, 1, n], so no need for slicing.
      const logits = outputTokens.reshape([outputTokens.shape[2]]);
      const sampledTokenIndex = logits.argMax().dataSync()[0];
      const sampledChar = this.reverseTargetCharIndex[sampledTokenIndex];
      decodedSentence += sampledChar;

      // Exit condition: either hit max length or find stop character.
      if (sampledChar === '\n' ||
          decodedSentence.length > this.maxDecoderSeqLength) {
        stopCondition = true;
      }

      // Update the target sequence (of length 1).
      targetSeq = tf.buffer([1, 1, this.numDecoderTokens]);
      targetSeq.set(1, 0, 0, sampledTokenIndex);

      // Update states.
      statesValue = [h, c];
    }

    return decodedSentence;
  }

  /** Translate the given English sentence into French. */
  translate(inputSentence) {
    ui.status('Translating...');
    const inputSeq = this.encodeString(inputSentence);
    const decodedSentence = this.decodeSequence(inputSeq);
    ui.status('');
    return decodedSentence;
  }
}

/**
 * Loads the pretrained model and metadata, and registers the translation
 * function with the UI.
 */
async function setupTranslator() {
  if (await loader.urlExists(HOSTED_URLS.model)) {
    ui.status('Model available: ' + HOSTED_URLS.model);
    const button = document.getElementById('load-pretrained-remote');
    button.addEventListener('click', async () => {
      const translator = await new Translator().init(HOSTED_URLS);
      ui.setTranslationFunction(x => translator.translate(x));
      ui.setEnglish('Go.', x => translator.translate(x));
    });
    button.style.display = 'inline-block';
  }

  if (await loader.urlExists(LOCAL_URLS.model)) {
    ui.status('Model available: ' + LOCAL_URLS.model);
    const button = document.getElementById('load-pretrained-local');
    button.addEventListener('click', async () => {
      const translator = await new Translator().init(LOCAL_URLS);
      ui.setTranslationFunction(x => translator.translate(x));
      ui.setEnglish('Go.', x => translator.translate(x));
    });
    button.style.display = 'inline-block';
  }

  ui.status('Standing by.');
}

setupTranslator();