Integrating AI and Machine Learning

Running a model in the browser

Nearly every "AI feature" I've been asked about in the last year turned out to be a fetch call to somebody else's API. That's usually the right answer. But there is a smaller category of feature where the model has to run on the user's machine, and React has no opinion about how you do that.

I wrote a while back about how AI changes the way we write code. This is the other half of it: shipping a model as part of the frontend rather than using one to produce the frontend. The tooling is more mature than I expected and the trade-offs are sharper than the marketing suggests.

The two libraries

TensorFlow.js is the serious one. It runs Keras and TensorFlow models in the browser or in Node, either on WebGL or, with @tensorflow/tfjs-backend-wasm, on WebAssembly. You can train in it, though mostly you won’t; the realistic use is loading a model somebody else trained and running inference on it.

Brain.js is the small one. Neural networks in a few lines, trained in the browser on data you have to hand:

import { NeuralNetwork } from 'brain.js';

const net = new NeuralNetwork();

net.train([
  { input: { r: 0.03, g: 0.7, b: 0.5 }, output: { black: 1 } },
  { input: { r: 0.16, g: 0.09, b: 0.2 }, output: { white: 1 } },
]);

net.run({ r: 1, g: 0.4, b: 0 });

That’s the whole API surface, more or less. If your problem fits in a small feedforward network and you can describe the training data in JSON, Brain.js will get you there faster than anything else. If it doesn’t fit, it won’t stretch, and you’ll be back on TensorFlow.js within a week.

A sentiment component, properly

The usual worked example is sentiment analysis, so let’s do that one and be honest about the sharp edges. TensorFlow.js publishes a sentiment example built on a small CNN trained against the IMDB reviews set, which is a reasonable thing to point at.

npm install @tensorflow/tfjs

First, loading. Almost every tutorial I’ve read loads the model inside the prediction function, which means a fresh download and a fresh parse on every keystroke. Load it once and hold the promise:

import * as tf from '@tensorflow/tfjs';

let modelPromise;

const loadModel = () => {
  modelPromise ??= tf.loadLayersModel('/models/sentiment/model.json');

  return modelPromise;
};

Then the prediction. The thing that trips people up here is that model.predict() returns a tensor, not a number. Render it straight into JSX and React will throw at you. You have to pull the value out with data(), which is async, and dispose of the tensors afterwards because the WebGL backend does not garbage collect them for you:

const predictSentiment = async (text) => {
  const model = await loadModel();
  // Tokenise to word indices and pad to the length the model expects.
  const input = preprocess(text);
  const output = model.predict(input);
  const [score] = await output.data();

  tf.dispose([input, output]);

  return score;
};

preprocess is doing a lot of quiet work in that snippet. A text model doesn’t take text; it takes a fixed-length tensor of integers, mapped through the exact vocabulary the model was trained with, padded or truncated to the exact sequence length. Get the vocabulary wrong and nothing errors. You just get confident nonsense, which is a considerably worse failure mode than a stack trace.

Finally, the component:

import { useState } from 'react';

const SentimentAnalyzer = () => {
  const [text, setText] = useState('');
  const [score, setScore] = useState(null);
  const [pending, setPending] = useState(false);

  const handleAnalyze = async () => {
    setPending(true);
    setScore(await predictSentiment(text));
    setPending(false);
  };

  return (
    <div>
      <textarea value={text} onChange={(event) => setText(event.target.value)} />
      <button onClick={handleAnalyze} disabled={pending}>
        Analyze
      </button>
      {score !== null && (
        <p>
          {score > 0.5 ? 'Positive' : 'Negative'} ({score.toFixed(2)})
        </p>
      )}
    </div>
  );
};

Nothing about that is React-specific, which is rather the point. The model is a module-level singleton, the prediction is an async function, and the component is a form. There’s no hook to learn.

What it costs

Three things, in roughly the order they’ll annoy you.

The library is a heavy dependency, and the model weights are a separate download on top of it, fetched at runtime as model.json plus binary shards. Neither belongs in your initial bundle. Lazy-load the whole feature behind a dynamic import and don’t touch it until the user does something that needs it.

Inference on the main thread will drop frames on a mid-range phone. WebGL helps, but the first call after loading is always the slow one while shaders compile, so warm the model up with a dummy prediction while the user is still typing.

And tensors leak. tf.dispose() and tf.tidy() exist for a reason, and forgetting them in a component that re-renders is how you end up with a tab consuming a gigabyte.

When to do this at all

My honest position: run the model in the browser when the data must not leave the device, when the feature has to work offline, or when a round trip per keystroke would be absurd. Local search ranking, an offline classifier, live pose or gesture detection from a webcam. Those earn it.

Everything else is better served by a request to a server, where the model is bigger, the hardware is yours, you can swap it without shipping a release, and nobody downloads several megabytes of weights to find out their review was positive.

The nice thing about 2024 is that both options are genuinely easy now. The hard part was never the integration; it’s deciding which side of the wire the model belongs on, and that decision has almost nothing to do with React.

Next

I want to try the same component with the WASM backend on an older phone and see how far off WebGL it lands. If the gap is small it changes the calculus quite a bit for anything that has to work on hardware people actually own 📱

‘Till next time!