Almacena en caché los modelos de IA en el navegador

La mayoría de los modelos de IA tienen algo en común: son bastante grandes para un recurso que se transfiere a través de Internet. El modelo de detección de objetos 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 tiene un tamaño de 1.35 GB, lo que se considera muy pequeño para un LLM. Los modelos de IA generativa pueden ser enormes. Por eso, gran parte del uso de la IA en la actualidad se realiza en la nube. Cada vez más, las apps ejecutan modelos altamente optimizados directamente en el dispositivo. Si bien existen demos de LLM que se ejecutan en el navegador, estos son algunos ejemplos de otros modelos de nivel de producción que se ejecutan en el navegador:

Adobe Photoshop en la Web con la herramienta de selección de objetos potenciada por IA abierta y 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 forma 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 publicar el modelo junto con el resto de los recursos de la app. Es fundamental optimizar la publicación.

Configura los encabezados de caché correctos

Si publicas 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 compilar según las necesidades de tu app.

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

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

Cuando el usuario vuelve a cargar la página, el cliente envía una solicitud de validación, aunque el servidor sepa que el contenido es estable. La directiva immutable indica explícitamente que no es necesario volver a validar, ya que el contenido no cambiará. Los navegadores y los servidores proxy o de caché intermedios no admiten ampliamente la directiva immutable, pero si la combinas con la directiva max-age, que se comprende de forma universal, puedes garantizar la máxima compatibilidad. La directiva de respuesta public indica que la respuesta se puede almacenar en un caché compartido.

Chrome DevTools muestra los encabezados Cache-Control de producción que envía Hugging Face cuando se solicita un modelo de IA. (Fuente)

Almacena en caché modelos de IA del cliente

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

Existen varias técnicas que puedes usar para lograrlo. En las siguientes muestras de código, se supone 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 está anotada con los métodos performance.mark() y performance.measure(). Estas medidas dependen del dispositivo y no son generalizables.

En las Herramientas para desarrolladores de Chrome, Application > Storage, 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 suma un total de 4,063 megabytes.

Puedes usar una de las siguientes APIs para almacenar en caché modelos de IA en el navegador: Cache API, Origin Private File System API y IndexedDB API. 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 persistente para los pares de objetos Request y Response que se almacenan en caché en la memoria duradera. Aunque se define en la especificación de los trabajadores de servicio, puedes usar esta API desde el subproceso principal o un trabajador normal. Para usarlo fuera de un contexto de trabajador de servicio, llama al método Cache.put() con un objeto Response sintético, vinculado a 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 el modelo directamente, usarías el Response que obtendrías si realizaras una solicitud fetch().

Por ejemplo, a continuación 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 Origin (OPFS) es un estándar relativamente nuevo para un extremo de almacenamiento. Es privado 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á muy optimizado para el rendimiento y ofrece acceso de escritura a su contenido.

Por ejemplo, a continuación, se muestra cómo almacenar y restablecer un archivo de modelo en el sistema de archivos 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 bien establecido para almacenar datos arbitrarios de forma persistente en el navegador. Es conocida por su API algo compleja, pero si usas una biblioteca de wrapper, como idb-keyval, puedes tratar 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;
  }
};

Marca el almacenamiento como persistente

Llama a navigator.storage.persist() al final de cualquiera de estos métodos de almacenamiento en caché para solicitar permiso para usar el almacenamiento persistente. Este método muestra una promesa que se resuelve en true si se otorga el permiso y en false de lo contrario. El navegador puede aceptar o rechazar 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 modelos de IA directamente desde el disco duro de un usuario como alternativa al almacenamiento del navegador. Esta técnica puede ayudar a que las apps centradas en la investigación muestren la viabilidad de ejecutar modelos determinados en el navegador o permitir que los artistas usen modelos autoentrenados en apps de creatividad para 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 necesita otorgar acceso al archivo del modelo una vez. Gracias a los permisos persistentes, el usuario puede optar por otorgar acceso permanente al archivo. Después de volver a cargar la app y un gesto del usuario obligatorio, como un clic del mouse, se puede restablecer FileSystemFileHandle desde IndexedDB con acceso al archivo en el disco duro.

Los permisos de acceso a los archivos se consultan y solicitan si es necesario, lo que facilita las cargas futuras. En el siguiente ejemplo, se muestra cómo obtener un control 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. Puede haber un caso en el 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 normales y el método de disco duro implementados en la demo de LLM de MediaPipe.

Bonificación: Descarga un archivo grande en fragmentos

Si necesitas descargar un modelo de IA grande de Internet, realiza la descarga en paralelo en segmentos separados y, luego, vuelve a unirlos en el cliente.

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

Puedes copiar la siguiente función en tu proyecto o instalar el paquete fetch-in-chunks desde 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 que es 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, garantizar el acceso rápido a los modelos de IA, reducir los tiempos de carga y mejorar la capacidad de respuesta.

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

En el caso de las aplicaciones de nicho, la API de acceso al sistema de archivos ofrece acceso directo a los archivos en el 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, guárdalo en el servidor. Una vez almacenados en el cliente, es trivial extraer los datos de la caché y IndexedDB con DevTools o la extensión de DevTools de OFPS. Estas APIs de almacenamiento son inherentemente iguales en seguridad. Es posible que te sientas tentado a almacenar una versión encriptada del modelo, pero luego debes obtener la clave de desencriptación para el cliente, que podría interceptarse. Esto significa que el intento de un usuario malintencionado 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 varias condiciones de red y restricciones del sistema.


Agradecimientos

Esta versión fue revisada por 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.