Almacena en caché los modelos de IA en el navegador

La mayoría de los modelos de IA tienen al menos un aspecto en común: son bastante grandes para un recurso que se transfiere a través de Internet. El modelo de detección de objetos de MediaPipe más pequeño (SSD MobileNetV2 float16) pesa 5.6 MB y el más grande es de alrededor de 25 MB.

El LLM de código abierto gemma-2b-it-gpu-int4.bin llega a los 1.35 GB, lo que se considera muy pequeño para un LLM. Los modelos de IA generativa pueden ser enormes. Por eso, hoy en día, gran parte del uso de la IA ocurre en la nube. Cada vez más, las apps ejecutan modelos altamente optimizados directamente en el dispositivo. Si bien existen demostraciones de LLM que se ejecutan en el navegador, estos son algunos ejemplos de nivel de producción de otros modelos que se ejecutan en el navegador:

Adobe Photoshop en la Web con la herramienta de selección de objetos potenciada por IA abierta, con tres objetos seleccionados: dos jirafas y una luna.

Para que los lanzamientos futuros de tus aplicaciones sean más rápidos, debes almacenar en caché de manera explícita los datos del modelo en el dispositivo, en lugar de depender de la caché implícita del navegador HTTP.

Si bien en esta guía se usa gemma-2b-it-gpu-int4.bin model para crear un chatbot, el enfoque se puede generalizar para adaptarse a otros modelos y otros casos de uso en el dispositivo. La forma más común de conectar una app a un modelo es entregarlo junto con el resto de los recursos de la app. Es crucial optimizar la entrega.

Configura los encabezados de caché correctos

Si entregas modelos de IA desde tu servidor, es importante configurar el encabezado Cache-Control correcto. En el siguiente ejemplo, se muestra una configuración predeterminada sólida que puedes usar según las necesidades de tu app.

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

Cada versión lanzada de un modelo de IA es un recurso estático. El contenido que nunca cambia debe recibir un max-age largo combinado con la impulsión de caché en la URL de la solicitud. Si necesitas actualizar el modelo, deberás proporcionarle una URL nueva.

Cuando el usuario vuelve a cargar la página, el cliente envía una solicitud de revalidación, aunque el servidor sepa que el contenido es estable. La directiva immutable indica de forma explícita que la revalidación es innecesaria, ya que el contenido no cambiará. La directiva immutable no es ampliamente compatible con los navegadores y los servidores proxy o caché intermedios, pero si la combinas con la directiva max-age entendida universalmente, puedes garantizar la máxima compatibilidad. La directiva de respuesta public indica que la respuesta se puede almacenar en una caché compartida.

Las Herramientas para desarrolladores de Chrome muestran los encabezados de producción Cache-Control que envía Hugging Face cuando se solicita un modelo de IA. (Fuente)

Almacena en caché los modelos de IA del cliente

Cuando entregas un modelo de IA, es importante almacenar en caché de forma explícita el modelo en el navegador. Esto garantiza que los datos del modelo estén disponibles cuando un usuario vuelva a cargar la app.

Existen varias técnicas que puedes usar para lograrlo. Para las siguientes muestras de código, supongamos que cada archivo de modelo se almacena en un objeto Blob llamado blob en la memoria.

Para comprender el rendimiento, cada muestra de código se anota con los métodos performance.mark() y performance.measure(). Estas mediciones dependen del dispositivo y no son generalizables.

En Aplicación > Almacenamiento de las Herramientas para desarrolladores de Chrome, revisa el diagrama de uso con segmentos para IndexedDB, el almacenamiento en caché y el sistema de archivos. Se muestra que cada segmento consume 1,354 megabytes de datos, lo que equivale a 4,063 megabytes.

Puedes elegir usar una de las siguientes APIs para almacenar en caché los modelos de IA en el navegador: la API de Cache, la API de Origin Private File System y la API de IndexedDB. La recomendación general es usar la API de Cache, pero en esta guía se analizan las ventajas y desventajas de todas las opciones.

API de Cache

La API de Cache proporciona almacenamiento continuo para los pares de objetos Request y Response que se almacenan en caché en la memoria de larga duración. Aunque está definida en las especificaciones de los service workers, puedes usar esta API desde el subproceso principal o un trabajador normal. Para usarlo fuera del contexto de un service worker, llama al método Cache.put() con un objeto Response sintético vinculado con una URL sintética en lugar de un objeto Request.

En esta guía, se supone que hay un blob en la memoria. Usa una URL falsa como clave de caché y un Response sintético basado en blob. Si descargaras directamente el modelo, usarías el Response que obtendrías si haces una solicitud fetch().

Por ejemplo, aquí se muestra cómo almacenar y restablecer un archivo de modelo con la API de 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 de Origin Private File System

El sistema de archivos privados de origen (OPFS) es un estándar comparativamente nuevo para un extremo de almacenamiento. Es privada para el origen de la página y, por lo tanto, es invisible para el usuario, a diferencia del sistema de archivos normal. Proporciona acceso a un archivo especial que está altamente optimizado para el rendimiento y ofrece acceso de escritura a su contenido.

Por ejemplo, aquí se muestra cómo almacenar y restablecer un archivo de modelo en 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 de IndexedDB

IndexedDB es un estándar establecido para almacenar datos arbitrarios de manera persistente en el navegador. Es infame por su API un tanto compleja, pero si usas una biblioteca de wrappers como idb-keyval, puedes tratar a IndexedDB como un almacén de pares clave-valor clásico.

Por ejemplo:

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;
  }
};

Marcar el almacenamiento como persistente

Llama a navigator.storage.persist() al final de cualquiera de estos métodos de almacenamiento en caché a fin de solicitar permiso para usar el almacenamiento persistente. Este método muestra una promesa que se resuelve en true si se otorga el permiso; de lo contrario, en false. El navegador puede o no cumplir con la solicitud, según las reglas específicas del navegador.

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);
  }
}

Caso especial: Usa un modelo en un disco duro

Puedes hacer referencia a los modelos de IA directamente desde el disco duro de un usuario como alternativa al almacenamiento del navegador. Esta técnica puede ayudar a las apps centradas en la investigación a mostrar la viabilidad de ejecutar modelos determinados en el navegador o permitir a los artistas usar modelos autoentrenados en apps de creatividad de expertos.

API de File System Access

Con la API de File System Access, puedes abrir archivos desde el disco duro y obtener un FileSystemFileHandle que puedes conservar en IndexedDB.

Con este patrón, el usuario solo debe otorgar acceso al archivo de modelo una vez. Gracias a los permisos persistentes, el usuario puede optar por otorgar acceso al archivo de forma permanente. Después de volver a cargar la app y de un gesto necesario del usuario, como un clic del mouse, se puede restablecer FileSystemFileHandle desde IndexedDB con acceso al archivo en el disco duro.

Los permisos de acceso a archivos se consultan y solicitan si es necesario, lo que hace que esto sea más sencillo para futuras recargas. En el siguiente ejemplo, se muestra cómo obtener un controlador para un archivo del disco duro y, luego, almacenarlo y restablecerlo.

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;
  }
};

Estos métodos no son mutuamente excluyentes. En algunos casos, es posible que almacenes en caché de forma explícita un modelo en el navegador y uses un modelo del disco duro de un usuario.

Demostración

Puedes ver los tres métodos de almacenamiento de casos regulares y el método de disco duro implementado en la demostración de MediaPipe LLM.

Contenido adicional: Descarga un archivo grande en partes

Si necesitas descargar un modelo de IA grande de Internet, paraleliza la descarga en fragmentos separados y vuelve a unirlos al cliente.

Esta es una función auxiliar que puedes usar en tu código. Solo debes pasarle url. chunkSize (valor predeterminado: 5 MB), maxParallelRequests (valor predeterminado: 6), la función progressCallback (que informa sobre downloadedBytes y el total de fileSize) y signal para un indicador AbortSignal son opcionales.

Puedes copiar la siguiente función en tu proyecto o instalar el paquete fetch-in-chunks del paquete 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;

Elige el método adecuado para ti

En esta guía, se exploraron varios métodos para almacenar en caché de manera eficaz los modelos de IA en el navegador, una tarea fundamental para mejorar la experiencia del usuario y el rendimiento de tu app. El equipo de almacenamiento de Chrome recomienda la API de Cache para obtener un rendimiento óptimo y garantizar un acceso rápido a los modelos de IA, lo que reduce los tiempos de carga y mejora la capacidad de respuesta.

OPFS e IndexedDB son opciones menos utilizables. Las APIs de OPFS y de IndexedDB deben serializar los datos antes de que se puedan almacenar. IndexedDB también necesita deserializar los datos cuando se recuperan, lo que lo convierte en el peor lugar para almacenar modelos grandes.

Para aplicaciones de nicho, la API de File System Access ofrece acceso directo a los archivos del dispositivo de un usuario, lo que es ideal para los usuarios que administran sus propios modelos de IA.

Si necesitas proteger tu modelo de IA, mantenlo en el servidor. Una vez almacenados en el cliente, es trivial extraer los datos de Cache e IndexedDB con DevOps o la extensión de OFPS para Herramientas para desarrolladores. La seguridad de estas APIs de almacenamiento es por naturaleza. Es posible que sientas la tentación de almacenar una versión encriptada del modelo, pero, luego, necesitarás enviar la clave de desencriptación al cliente, que podría ser interceptada. Esto significa que el intento de una persona que actúa de mala fe de robar tu modelo es un poco más difícil, pero no imposible.

Te recomendamos que elijas una estrategia de almacenamiento en caché que se alinee con los requisitos de tu app, el comportamiento del público objetivo y las características de los modelos de IA que se usan. Esto garantiza que tus aplicaciones sean responsivas y sólidas en diversas condiciones de red y restricciones del sistema.


Agradecimientos

Joshua Bell, Reilly Grant, Evan Stade, Nathan Memmott, Austin Sullivan, Etienne Noël, André Bandarra, Alexandra Klepper, François Beaufort, Paul Kinlan y Rachel Andrew