Mettre en cache les modèles d'IA dans le navigateur

La plupart des modèles d'IA ont une caractéristique commune: ils sont assez volumineux pour une ressource transférée sur Internet. Le plus petit modèle de détection d'objets MediaPipe (SSD MobileNetV2 float16) pèse 5,6 Mo, et le plus grand environ 25 Mo.

Le LLM Open Source gemma-2b-it-gpu-int4.bin fait 1,35 Go, ce qui est considéré comme très faible pour un LLM. Les modèles d'IA générative peuvent être énormes. C'est pourquoi l'IA est aujourd'hui largement utilisée dans le cloud. De plus en plus d'applications exécutent des modèles hautement optimisés directement sur l'appareil. Bien qu'il existe des démos de LLM exécutées dans le navigateur, voici quelques exemples d'autres modèles exécutés dans le navigateur, prêts à être utilisés en production:

Adobe Photoshop sur le Web avec l'outil de sélection d'objets optimisé par l'IA ouvert et trois objets sélectionnés: deux girafes et une lune.

Pour accélérer les lancements futurs de vos applications, vous devez mettre en cache explicitement les données du modèle sur l'appareil, plutôt que de vous appuyer sur le cache HTTP implicite du navigateur.

Bien que ce guide utilise gemma-2b-it-gpu-int4.bin model pour créer un chatbot, l'approche peut être généralisée pour s'adapter à d'autres modèles et à d'autres cas d'utilisation sur l'appareil. La méthode la plus courante pour connecter une application à un modèle consiste à diffuser le modèle avec le reste des ressources de l'application. Il est essentiel d'optimiser la diffusion.

Configurer les bons en-têtes de cache

Si vous diffusez des modèles d'IA à partir de votre serveur, il est important de configurer l'en-tête Cache-Control approprié. L'exemple suivant montre un paramètre par défaut solide, sur lequel vous pouvez vous appuyer pour répondre aux besoins de votre application.

Cache-Control: public, max-age=31536000, immutable

Chaque version publiée d'un modèle d'IA est une ressource statique. Les contenus qui ne changent jamais doivent être associés à un max-age long combiné à un vidage du cache dans l'URL de la requête. Si vous devez mettre à jour le modèle, vous devez lui attribuer une nouvelle URL.

Lorsque l'utilisateur actualise la page, le client envoie une requête de revalidation, même si le serveur sait que le contenu est stable. La directive immutable indique explicitement que la revalidation n'est pas nécessaire, car le contenu ne change pas. La directive immutable n'est pas largement compatible avec les navigateurs et les serveurs proxy ou de cache intermédiaires, mais en la combinant à la directive max-age, qui est universellement comprise, vous pouvez assurer une compatibilité maximale. L'instruction de réponse public indique que la réponse peut être stockée dans un cache partagé.

Les outils de développement Chrome affichent les en-têtes Cache-Control de production envoyés par Hugging Face lors de la demande d'un modèle d'IA. (Source)

Mettre en cache des modèles d'IA côté client

Lorsque vous diffusez un modèle d'IA, il est important de le mettre en cache explicitement dans le navigateur. Cela garantit que les données du modèle sont facilement disponibles après qu'un utilisateur a actualisé l'application.

Pour ce faire, vous pouvez utiliser plusieurs techniques. Pour les exemples de code suivants, supposons que chaque fichier de modèle soit stocké dans un objet Blob nommé blob en mémoire.

Pour comprendre les performances, chaque exemple de code est annoté avec les méthodes performance.mark() et performance.measure(). Ces mesures dépendent de l'appareil et ne sont pas généralisables.

Dans les outils pour les développeurs Chrome, Application > Stockage, examinez le diagramme d'utilisation avec des segments pour IndexedDB, le stockage en cache et le système de fichiers. Chaque segment consomme 1 354 mégaoctets de données, ce qui représente un total de 4 063 mégaoctets.

Vous pouvez choisir d'utiliser l'une des API suivantes pour mettre en cache des modèles d'IA dans le navigateur : l'API Cache, l'API Origin Private File System et l'API IndexedDB. Nous vous recommandons d'utiliser l'API Cache, mais ce guide présente les avantages et les inconvénients de toutes les options.

API Cache

L'API Cache fournit un stockage persistant pour les paires d'objets Request et Response mises en cache dans une mémoire durable. Bien qu'elle soit définie dans la spécification des services workers, vous pouvez utiliser cette API à partir du thread principal ou d'un worker standard. Pour l'utiliser en dehors d'un contexte de service worker, appelez la méthode Cache.put() avec un objet Response synthétique associé à une URL synthétique au lieu d'un objet Request.

Ce guide suppose que vous utilisez une blob en mémoire. Utilisez une URL fictive comme clé de cache et un Response synthétique basé sur blob. Si vous téléchargez directement le modèle, vous devez utiliser l'Response que vous obtiendrez en envoyant une requête fetch().

Par exemple, voici comment stocker et restaurer un fichier de modèle avec l'API Cache.

const storeFileInSWCache = async (blob) => {
  try {
    performance.mark('start-sw-cache-cache');
    const modelCache = await caches.open('models');
    await modelCache.put('model.bin', new Response(blob));
    performance.mark('end-sw-cache-cache');

    const mark = performance.measure(
      'sw-cache-cache',
      'start-sw-cache-cache',
      'end-sw-cache-cache'
    );
    console.log('Model file cached in sw-cache.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromSWCache = async () => {
  try {
    performance.mark('start-sw-cache-restore');
    const modelCache = await caches.open('models');
    const response = await modelCache.match('model.bin');
    if (!response) {
      throw new Error(`File model.bin not found in sw-cache.`);
    }
    const file = await response.blob();
    performance.mark('end-sw-cache-restore');
    const mark = performance.measure(
      'sw-cache-restore',
      'start-sw-cache-restore',
      'end-sw-cache-restore'
    );
    console.log(mark.name, mark.duration.toFixed(2));
    console.log('Cached model file found in sw-cache.');
    return file;
  } catch (err) {    
    throw err;
  }
};

API Origin Private File System

Le Origin Private File System (OPFS) est une norme relativement récente pour un point de terminaison de stockage. Il est privé par rapport à l'origine de la page et est donc invisible pour l'utilisateur, contrairement au système de fichiers standard. Il permet d'accéder à un fichier spécial hautement optimisé pour les performances et offre un accès en écriture à son contenu.

Par exemple, voici comment stocker et restaurer un fichier de modèle dans l'OPFS.

const storeFileInOPFS = async (blob) => {
  try {
    performance.mark('start-opfs-cache');
    const root = await navigator.storage.getDirectory();
    const handle = await root.getFileHandle('model.bin', { create: true });
    const writable = await handle.createWritable();
    await blob.stream().pipeTo(writable);
    performance.mark('end-opfs-cache');
    const mark = performance.measure(
      'opfs-cache',
      'start-opfs-cache',
      'end-opfs-cache'
    );
    console.log('Model file cached in OPFS.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromOPFS = async () => {
  try {
    performance.mark('start-opfs-restore');
    const root = await navigator.storage.getDirectory();
    const handle = await root.getFileHandle('model.bin');
    const file = await handle.getFile();
    performance.mark('end-opfs-restore');
    const mark = performance.measure(
      'opfs-restore',
      'start-opfs-restore',
      'end-opfs-restore'
    );
    console.log('Cached model file found in OPFS.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

API IndexedDB

IndexedDB est une norme bien établie pour stocker des données arbitraires de manière persistante dans le navigateur. Il est tristement célèbre pour son API quelque peu complexe, mais en utilisant une bibliothèque de wrapper telle que idb-keyval, vous pouvez traiter IndexedDB comme un magasin de clés-valeurs classique.

Exemple :

import { get, set } from 'https://cdn.jsdelivr.net/npm/idb-keyval@latest/+esm';

const storeFileInIDB = async (blob) => {
  try {
    performance.mark('start-idb-cache');
    await set('model.bin', blob);
    performance.mark('end-idb-cache');
    const mark = performance.measure(
      'idb-cache',
      'start-idb-cache',
      'end-idb-cache'
    );
    console.log('Model file cached in IDB.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromIDB = async () => {
  try {
    performance.mark('start-idb-restore');
    const file = await get('model.bin');
    if (!file) {
      throw new Error('File model.bin not found in IDB.');
    }
    performance.mark('end-idb-restore');
    const mark = performance.measure(
      'idb-restore',
      'start-idb-restore',
      'end-idb-restore'
    );
    console.log('Cached model file found in IDB.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

Marquer l'espace de stockage comme persistant

Appelez navigator.storage.persist() à la fin de l'une de ces méthodes de mise en cache pour demander l'autorisation d'utiliser un stockage persistant. Cette méthode renvoie une promesse qui renvoie true si l'autorisation est accordée et false dans le cas contraire. Le navigateur peut ou non accepter la requête, en fonction des règles spécifiques au navigateur.

if ('storage' in navigator && 'persist' in navigator.storage) {
  try {
    const persistent = await navigator.storage.persist();
    if (persistent) {
      console.log("Storage will not be cleared except by explicit user action.");
      return;
    }
    console.log("Storage may be cleared under storage pressure.");  
  } catch (err) {
    console.error(err.name, err.message);
  }
}

Cas particulier: utiliser un modèle sur un disque dur

Vous pouvez référencer des modèles d'IA directement à partir du disque dur d'un utilisateur comme alternative au stockage du navigateur. Cette technique peut aider les applications axées sur la recherche à montrer la faisabilité de l'exécution de modèles donnés dans le navigateur ou à permettre aux artistes d'utiliser des modèles auto-entraînés dans des applications de créativité experte.

API File System Access

Avec l'API File System Access, vous pouvez ouvrir des fichiers à partir du disque dur et obtenir un FileSystemFileHandle que vous pouvez conserver dans IndexedDB.

Avec ce modèle, l'utilisateur n'a besoin d'accorder l'accès au fichier de modèle qu'une seule fois. Grâce aux autorisations persistantes, l'utilisateur peut choisir d'accorder un accès permanent au fichier. Après avoir actualisé l'application et effectué un geste utilisateur requis, tel qu'un clic de souris, FileSystemFileHandle peut être restauré à partir d'IndexedDB avec accès au fichier sur le disque dur.

Les autorisations d'accès aux fichiers sont interrogées et demandées si nécessaire, ce qui facilite les futurs rechargements. L'exemple suivant montre comment obtenir un handle pour un fichier à partir du disque dur, puis le stocker et le restaurer.

import { fileOpen } from 'https://cdn.jsdelivr.net/npm/browser-fs-access@latest/dist/index.modern.js';
import { get, set } from 'https://cdn.jsdelivr.net/npm/idb-keyval@latest/+esm';

button.addEventListener('click', async () => {
  try {
    const file = await fileOpen({
      extensions: ['.bin'],
      mimeTypes: ['application/octet-stream'],
      description: 'AI model files',
    });
    if (file.handle) {
      // It's an asynchronous method, but no need to await it.
      storeFileHandleInIDB(file.handle);
    }
    return file;
  } catch (err) {
    if (err.name !== 'AbortError') {
      console.error(err.name, err.message);
    }
  }
});

const storeFileHandleInIDB = async (handle) => {
  try {
    performance.mark('start-file-handle-cache');
    await set('model.bin.handle', handle);
    performance.mark('end-file-handle-cache');
    const mark = performance.measure(
      'file-handle-cache',
      'start-file-handle-cache',
      'end-file-handle-cache'
    );
    console.log('Model file handle cached in IDB.', mark.name, mark.duration.toFixed(2));
  } catch (err) {
    console.error(err.name, err.message);
  }
};

const restoreFileFromFileHandle = async () => {
  try {
    performance.mark('start-file-handle-restore');
    const handle = await get('model.bin.handle');
    if (!handle) {
      throw new Error('File handle model.bin.handle not found in IDB.');
    }
    if ((await handle.queryPermission()) !== 'granted') {
      const decision = await handle.requestPermission();
      if (decision === 'denied' || decision === 'prompt') {
        throw new Error(Access to file model.bin.handle not granted.');
      }
    }
    const file = await handle.getFile();
    performance.mark('end-file-handle-restore');
    const mark = performance.measure(
      'file-handle-restore',
      'start-file-handle-restore',
      'end-file-handle-restore'
    );
    console.log('Cached model file handle found in IDB.', mark.name, mark.duration.toFixed(2));
    return file;
  } catch (err) {    
    throw err;
  }
};

Ces méthodes ne s'excluent pas mutuellement. Il peut arriver que vous mettiez explicitement en cache un modèle dans le navigateur et que vous utilisiez un modèle à partir du disque dur d'un utilisateur.

Démo

Vous pouvez voir les trois méthodes de stockage de cas standards et la méthode de disque dur implémentées dans la démonstration du LLM MediaPipe.

Bonus: Télécharger un fichier volumineux par fragments

Si vous devez télécharger un grand modèle d'IA depuis Internet, parallélisez le téléchargement en segments distincts, puis assemblez-les à nouveau sur le client.

Voici une fonction d'assistance que vous pouvez utiliser dans votre code. Il vous suffit de lui transmettre url. Les valeurs chunkSize (par défaut: 5 Mo), maxParallelRequests (par défaut: 6), la fonction progressCallback (qui fournit des rapports sur le downloadedBytes et le fileSize total) et signal pour un signal AbortSignal sont toutes facultatives.

Vous pouvez copier la fonction suivante dans votre projet ou installer le package fetch-in-chunks à partir du package npm.

async function fetchInChunks(
  url,
  chunkSize = 5 * 1024 * 1024,
  maxParallelRequests = 6,
  progressCallback = null,
  signal = null
) {
  // Helper function to get the size of the remote file using a HEAD request
  async function getFileSize(url, signal) {
    const response = await fetch(url, { method: 'HEAD', signal });
    if (!response.ok) {
      throw new Error('Failed to fetch the file size');
    }
    const contentLength = response.headers.get('content-length');
    if (!contentLength) {
      throw new Error('Content-Length header is missing');
    }
    return parseInt(contentLength, 10);
  }

  // Helper function to fetch a chunk of the file
  async function fetchChunk(url, start, end, signal) {
    const response = await fetch(url, {
      headers: { Range: `bytes=${start}-${end}` },
      signal,
    });
    if (!response.ok && response.status !== 206) {
      throw new Error('Failed to fetch chunk');
    }
    return await response.arrayBuffer();
  }

  // Helper function to download chunks with parallelism
  async function downloadChunks(
    url,
    fileSize,
    chunkSize,
    maxParallelRequests,
    progressCallback,
    signal
  ) {
    let chunks = [];
    let queue = [];
    let start = 0;
    let downloadedBytes = 0;

    // Function to process the queue
    async function processQueue() {
      while (start < fileSize) {
        if (queue.length < maxParallelRequests) {
          let end = Math.min(start + chunkSize - 1, fileSize - 1);
          let promise = fetchChunk(url, start, end, signal)
            .then((chunk) => {
              chunks.push({ start, chunk });
              downloadedBytes += chunk.byteLength;

              // Update progress if callback is provided
              if (progressCallback) {
                progressCallback(downloadedBytes, fileSize);
              }

              // Remove this promise from the queue when it resolves
              queue = queue.filter((p) => p !== promise);
            })
            .catch((err) => {              
              throw err;              
            });
          queue.push(promise);
          start += chunkSize;
        }
        // Wait for at least one promise to resolve before continuing
        if (queue.length >= maxParallelRequests) {
          await Promise.race(queue);
        }
      }

      // Wait for all remaining promises to resolve
      await Promise.all(queue);
    }

    await processQueue();

    return chunks.sort((a, b) => a.start - b.start).map((chunk) => chunk.chunk);
  }

  // Get the file size
  const fileSize = await getFileSize(url, signal);

  // Download the file in chunks
  const chunks = await downloadChunks(
    url,
    fileSize,
    chunkSize,
    maxParallelRequests,
    progressCallback,
    signal
  );

  // Stitch the chunks together
  const blob = new Blob(chunks);

  return blob;
}

export default fetchInChunks;

Choisir la méthode qui vous convient

Ce guide a exploré différentes méthodes pour mettre en cache efficacement des modèles d'IA dans le navigateur, une tâche cruciale pour améliorer l'expérience utilisateur et les performances de votre application. L'équipe de stockage Chrome recommande l'API Cache pour des performances optimales, afin de garantir un accès rapide aux modèles d'IA, de réduire les temps de chargement et d'améliorer la réactivité.

Les options OPFS et IndexedDB sont moins utilisables. Les API OPFS et IndexedDB doivent sérialiser les données avant qu'elles ne puissent être stockées. IndexedDB doit également désérialiser les données lors de leur récupération, ce qui en fait le pire endroit pour stocker de grands modèles.

Pour les applications de niche, l'API File System Access offre un accès direct aux fichiers sur l'appareil de l'utilisateur, ce qui est idéal pour les utilisateurs qui gèrent leurs propres modèles d'IA.

Si vous devez sécuriser votre modèle d'IA, conservez-le sur le serveur. Une fois stockées sur le client, il est facile d'extraire les données du cache et d'IndexedDB avec les outils de développement ou l'extension OFPS DevTools. Ces API de stockage sont intrinsèquement équivalentes en termes de sécurité. Vous pourriez être tenté de stocker une version chiffrée du modèle, mais vous devez ensuite transmettre la clé de déchiffrement au client, qui peut être interceptée. Cela signifie que la tentative de vol de votre modèle par un acteur malveillant est légèrement plus difficile, mais pas impossible.

Nous vous encourageons à choisir une stratégie de mise en cache qui correspond aux exigences de votre application, au comportement de l'audience cible et aux caractéristiques des modèles d'IA utilisés. Cela garantit que vos applications sont réactives et robustes dans différentes conditions réseau et contraintes système.


Remerciements

Cet article a été examiné par Joshua Bell, Reilly Grant, Evan Stade, Nathan Memmott, Austin Sullivan, Etienne Noël, André Bandarra, Alexandra Klepper, François Beaufort, Paul Kinlan et Rachel Andrew.