Armazenar modelos de IA em cache no navegador

Thomas Steiner

A maioria dos modelos de IA tem uma coisa em comum: eles são bastante grandes para um recurso transferido pela Internet. O modelo de detecção de objetos do MediaPipe mais leve (SSD MobileNetV2 float16) pesa 5,6 MB e o maior tem cerca de 25 MB.

O LLM de código aberto gemma-2b-it-gpu-int4.bin tem 1,35 GB, o que é considerado muito pequeno para um LLM. Os modelos de IA generativa podem ser enormes. É por isso que muito uso da IA hoje acontece na nuvem. Cada vez mais, os apps executam modelos altamente otimizados diretamente no dispositivo. Embora existam demonstrações de LLMs em execução no navegador, confira alguns exemplos no nível de produção de outros modelos executados no navegador:

• O Adobe Photoshop executa uma variante do modelo Conv2D no dispositivo para a ferramenta de seleção de objetos inteligente.
• O Google Meet executa uma versão otimizada do modelo MobileNetV3-small para segmentação de pessoas com o recurso de desfoque do plano de fundo.
• A Tokopedia executa o modelo MediaPipeFaceDetector-TFJS para a detecção facial em tempo real e evitar inscrições inválidas no serviço.
• O Google Colab permite que os usuários usem modelos do disco rígido nos notebooks do Colab.

Para tornar as inicializações futuras dos aplicativos mais rápidas, armazene os dados do modelo em cache explicitamente no dispositivo, em vez de depender do cache HTTP implícito do navegador.

Embora este guia use o gemma-2b-it-gpu-int4.bin model para criar um bot de bate-papo, a abordagem pode ser generalizada para se adequar a outros modelos e outros casos de uso no dispositivo. A maneira mais comum de conectar um app a um modelo é veicular o modelo com o restante dos recursos do app. É fundamental otimizar a entrega.

Configurar os cabeçalhos de cache corretos

Se você disponibilizar modelos de IA do seu servidor, é importante configurar o cabeçalho Cache-Control correto. O exemplo a seguir mostra uma configuração padrão sólida, que pode ser usada para atender às necessidades do seu app.

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

Cada versão lançada de um modelo de IA é um recurso estático. Conteúdo que nunca muda precisa ter um longo max-age combinado com quebra de cache no URL da solicitação. Se você precisar atualizar o modelo, forneça um novo URL.

Quando o usuário recarrega a página, o cliente envia uma solicitação de nova validação, mesmo que o servidor saiba que o conteúdo está estável. A diretiva immutable indica explicitamente que a revalidação é desnecessária, porque o conteúdo não vai mudar. A diretiva immutable não tem ampla compatibilidade com navegadores e servidores de proxy ou cache intermediários. No entanto, ao combinar a diretiva com a max-age, que é universalmente compreendida, você pode garantir a compatibilidade máxima. A diretiva de resposta public indica que a resposta pode ser armazenada em um cache compartilhado.

O Chrome DevTools mostra os cabeçalhos Cache-Control de produção enviados pelo Hugging Face ao solicitar um modelo de IA. (Fonte)

Armazenar modelos de IA em cache do lado do cliente

Ao veicular um modelo de IA, é importante armazenar em cache o modelo explicitamente no navegador. Isso garante que os dados do modelo estejam disponíveis assim que o usuário recarrega o app.

Há várias técnicas que podem ser usadas para isso. Para os exemplos de código a seguir, suponha que cada arquivo de modelo esteja armazenado em um objeto Blob chamado blob na memória.

Para entender a performance, cada exemplo de código é anotado com os métodos performance.mark() e performance.measure(). Essas medidas dependem do dispositivo e não são generalizáveis.

No Chrome DevTools Application > Storage, analise o diagrama de uso com segmentos para IndexedDB, armazenamento em cache e sistema de arquivos. Cada segmento consome 1.354 megabytes de dados, totalizando 4.063 megabytes.

É possível usar uma das seguintes APIs para armazenar em cache modelos de IA no navegador: API Cache, a API Origin Private File System e a API IndexedDB. A recomendação geral é usar a API Cache, mas este guia discute as vantagens e desvantagens de todas as opções.

API Cache

A API Cache oferece armazenamento persistente para pares de objetos Request e Response que são armazenados em cache na memória de longa duração. Embora ela seja definida na especificação de Service Workers, é possível usar essa API na linha de execução principal ou em um worker regular. Para usá-lo fora de um contexto de service worker, chame o método Cache.put() com um objeto Response sintético, associado a um URL sintético em vez de um objeto Request.

Este guia pressupõe um blob na memória. Use um URL falso como a chave de cache e um Response sintético com base no blob. Se você fizer o download direto do modelo, vai usar o Response que receberia ao fazer uma solicitação fetch().

Por exemplo, confira como armazenar e restaurar um arquivo de modelo com a 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

O Origin Private File System (OPFS, na sigla em inglês) é um padrão relativamente novo para um endpoint de armazenamento. Ele é privado para a origem da página e, portanto, fica invisível para o usuário, ao contrário do sistema de arquivos normal. Ele fornece acesso a um arquivo especial altamente otimizado para desempenho e acesso de gravação ao conteúdo.

Por exemplo, confira como armazenar e restaurar um arquivo de modelo no 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

O IndexedDB é um padrão bem estabelecido para armazenar dados arbitrários de maneira persistente no navegador. Ela é conhecida por ter uma API um tanto complexa, mas, usando uma biblioteca wrapper como idb-keyval, é possível tratar o IndexedDB como um armazenamento clássico de chave-valor.

Exemplo:

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 o armazenamento como persistente

Chame navigator.storage.persist() no final de qualquer um desses métodos de armazenamento em cache para solicitar permissão para usar o armazenamento persistente. Esse método retorna uma promessa que é resolvida para true se a permissão for concedida e para false caso contrário. O navegador pode ou não aceitar a solicitação, dependendo das regras específicas do 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);
  }
}

Dica: use a API Storage Buckets em navegadores compatíveis para marcar o modelo como persistido:

// Create a persisted storage bucket specifically for models.
const modelsBucket = await navigator.storageBuckets.open(models', {
  durability: 'strict'
  persisted: true,
});

Caso especial: usar um modelo em um disco rígido

É possível fazer referência a modelos de IA diretamente do disco rígido de um usuário como uma alternativa ao armazenamento do navegador. Essa técnica pode ajudar apps focados em pesquisa a mostrar a viabilidade de executar determinados modelos no navegador ou permitir que os artistas usem modelos autotreinados em apps especializados de criatividade.

API File System Access

Com a API File System Access, é possível abrir arquivos do disco rígido e receber um FileSystemFileHandle que pode ser mantido no IndexedDB.

Com esse padrão, o usuário só precisa conceder acesso ao arquivo do modelo uma vez. Graças às permissões persistentes, o usuário pode conceder acesso permanente ao arquivo. Depois de recarregar o app e um gesto do usuário necessário, como um clique do mouse, o FileSystemFileHandle pode ser restaurado do IndexedDB com acesso ao arquivo no disco rígido.

As permissões de acesso a arquivos são consultadas e solicitadas, se necessário, o que facilita as recarregares futuras. O exemplo a seguir mostra como conseguir um handle para um arquivo do disco rígido e, em seguida, armazenar e restaurar o handle.

Observação: se o navegador não oferecer suporte à API File System Access, o browser-fs-access usado na biblioteca de snippet de código abaixo vai usar <input type="file">. Isso permite que o app acesse o arquivo de modelo. No entanto, o usuário precisa abrir manualmente os arquivos de modelo sempre que recarrega o app.

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

Esses métodos não são mutuamente exclusivos. Pode haver um caso em que você armazenará um modelo em cache explicitamente no navegador e usará um modelo do disco rígido de um usuário.

Demonstração

É possível conferir os três métodos regulares de armazenamento de casos e o método de disco rígido implementado na demonstração do LLM do MediaPipe.

Observação: no código-fonte, a demonstração compara todos os métodos de armazenamento entre si, para armazenamento e restauração. Na pior das hipóteses, o modelo é armazenado de maneira redundante quatro vezes. Isso é apenas para demonstração e não deve ser feito na produção.

Bônus: fazer o download de um arquivo grande em partes

Se você precisar fazer o download de um modelo de IA grande da Internet, faça o download em paralelo em partes separadas e junte tudo novamente no cliente.

Confira uma função auxiliar que pode ser usada no código. Você só precisa transmitir o url. O chunkSize (padrão: 5 MB), o maxParallelRequests (padrão: 6), a função progressCallback (que informa o downloadedBytes e o total de fileSize) e o signal para um sinal AbortSignal são opcionais.

Você pode copiar a função a seguir no seu projeto ou instalar o pacote fetch-in-chunks do 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;

Escolha o método ideal para você

Este guia explorou vários métodos para armazenar em cache modelos de IA de maneira eficaz no navegador, uma tarefa crucial para melhorar a experiência do usuário e o desempenho do app. A equipe de armazenamento do Chrome recomenda a API Cache para garantir o acesso rápido aos modelos de IA, reduzindo os tempos de carregamento e melhorando a capacidade de resposta.

O OPFS e o IndexedDB são opções menos úteis. As APIs OPFS e IndexedDB precisam serializar os dados antes de serem armazenados. O IndexedDB também precisa desserializar os dados quando eles forem recuperados, o que faz dele o pior lugar para armazenar modelos grandes.

Para aplicativos de nicho, a API File System Access oferece acesso direto a arquivos no dispositivo de um usuário, ideal para usuários que gerenciam os próprios modelos de IA.

Se você precisar proteger seu modelo de IA, mantenha-o no servidor. Depois de armazenados no cliente, é fácil extrair os dados do cache e do IndexedDB com as DevTools ou a extensão das DevTools para OFPS. Essas APIs de armazenamento são inerentemente iguais em segurança. Você pode querer armazenar uma versão criptografada do modelo, mas precisa enviar a chave de descriptografia ao cliente, que pode ser interceptada. Isso significa que a tentativa de um usuário de má-fé de roubar o modelo é um pouco mais difícil, mas não impossível.

Recomendamos que você escolha uma estratégia de armazenamento em cache que se alinhe aos requisitos do app, ao comportamento do público-alvo e às características dos modelos de IA usados. Isso garante que seus aplicativos sejam responsivos e robustos em várias condições de rede e restrições do sistema.

---

Agradecimentos

Este artigo foi revisado por Joshua Bell, Reilly Grant, Evan Stade, Nathan Memmott, Austin Sullivan, Etienne Noël, André Bandarra, Alexandra Klepper, François Beaufort, Paul Kinlan e Rachel Andrew.