TokenDagger – um tokenizador mais rápido que o Tiktoken da OpenAI

• Tokenizador de alto desempenho 100% compatível com o TikToken da OpenAI, oferecendo mais de 2x de throughput e tokenização de código 4x mais rápida no processamento de grandes volumes de texto
• Motor de parsing de expressões regulares de alta velocidade baseado em PCRE2 para maximizar a velocidade de correspondência de padrões de tokens
• Algoritmo BPE simplificado para minimizar a queda de desempenho ao lidar com grandes volumes de tokens especiais
• Em benchmarks reais, a tokenização de código é mais de 4x mais rápida, e é possível substituir diretamente o código existente que usa TikToken
• Suporta Python 3.8+, pode ser instalado facilmente via PyPI com pip install tokendagger e tem dependência de PCRE2

[GN⁺]

Comentários no Hacker News

• Acho que, no curto prazo, ainda há muito espaço para otimização de desempenho na infraestrutura de AI/ML resolvendo gargalos centrais assim com C++, não reescrevendo tudo em C++, mas com trade-offs de engenharia inteligentes que muitas vezes levam a ganhos reais de performance, especialmente engenheiros chineses parecem ser bons nesse tipo de trabalho
  • Meu mentor me ensinava uma forma de pensar sobre desenvolvimento de software: etapa 1, fazer funcionar; etapa 2, fazer rápido; etapa 3, deixar bonito. Transformers e LLMs já chegaram bem à fase de “funciona”, então parece que agora estamos num período em que os maiores avanços estão acontecendo no lado de performance
  • No longo prazo, acho que faz sentido se afastar do foco em Python; se ML engineers usam só por familiaridade, isso não parece muito orientado ao futuro
  • Na prática, o TikToken é escrito em Rust, então fico curioso se essa melhoria realmente vem da migração para C++
  • Na verdade, tokenização não é o maior gargalo, e a maior parte do cálculo acontece na execução real dos kernels CUDA. O overhead de Python é muito pequeno (o VLLM também é escrito majoritariamente em Python), então reescrever em C++ quase sempre significa reescrever os kernels CUDA de forma mais eficiente
• Acho algo bem bonito quando se cria uma substituição drop-in para um sistema existente que melhora bastante a performance; me lembra o ScyllaDB
  • Na prática, acho que ninguém usaria se não fosse esse tipo de substituto
• Fico curioso sobre a importância do tokenizer para a performance geral de LLMs. Tenho interesse em otimização de tokenizers, mas ainda não medi isso na prática. Eu supunha que a maior parte do custo vinha de matmul, mas, pelos comentários, o tokenizer parece ter relevância
  • Tokenização normalmente roda na CPU e raramente vira gargalo em treino ou inferência; quase todo o tempo vai para os kernels da GPU, e só modelos muito pequenos são exceção. A latência do tokenizer também pode ser “escondida” pela forma como a CPU prepara o próximo batch
  • Performance de tokenização é um tema meio complexo, mas organizações com competência e recursos reais costumam escrever tokenizers muito rápidos por conta própria. SentencePiece e tiktoken são exemplos clássicos de escolhas feitas mesmo com o custo de complexidade e distribuição. Especialistas de verdade olham os gargalos com flame graphs, e em treinamentos grandes às vezes fazem tokenização antecipada. Também dá para sentir uma certa tensão no setor com o C++ voltando a ganhar espaço, ao contrário da narrativa do Rust; espero que haja um motivo novo ou algum insight interessante aí
  • Como em outros comentários, na prática o tokenizer não é o gargalo; só foi trabalhado primeiro porque é a etapa inicial do pipeline de inferência
  • Tokenização de texto representa uma fração realmente mínima do cálculo total, mas ao processar dados em escala de petabytes, qualquer ganho de velocidade já vale a pena
• Para acompanhar o tiktoken, é preciso converter o formato do vocab, e seria ainda melhor como substituto drop-in totalmente compatível se essa exigência também desaparecesse. Também seria bom haver um exemplo bidirecional em que se inicializa o tiktoken, depois o tokendagger, e se verifica se os resultados são iguais
  • Desde a versão 0.1.1, ele virou um substituto drop-in de verdade, e em breve planejam adicionar esse exemplo também
  • Você apontou um ótimo ponto, então atualizei isso imediatamente
• Também tenho curiosidade sobre como esse projeto se compara ao crate BPE, cujo ponto forte é a capacidade de retokenizar texto incrementalmente e ser mais rápido que o tiktoken
  • Na próxima etapa, planejo adicionar a função de retokenização incremental e também fazer benchmark com esse crate
• Queria saber como conseguir tokenizers locais para outros LLMs; por exemplo, o Gemini só expõe API remota. Fico em dúvida se isso é algo proprietário ou se haveria alguma forma de inferir o mapeamento de tokens fazendo muitas chamadas à API
  • O Gemini usa SentencePiece e compartilha o mesmo vocabulary de tokenizer do Gemma (referência 1, referência 2, referência 3); entre os grandes laboratórios, só a Anthropic, com Claude 3 ou superior, não tem tokenizer local
  • Os tokenizers de cada modelo em geral compartilham algoritmos centrais como SentencePiece ou Byte-pair encoding (BPE), e diferem mais no nível do wrapper, por exemplo no tratamento de tokens especiais. Tiktoken e TokenDagger também são implementações de BPE. Se a biblioteca refletir as características específicas de cada modelo, dá para obter pequenos ganhos de performance e facilitar a integração. Não é algo muito grande, então adaptar a novos modelos também não pesa tanto. Exemplos de referência: tokenizer.py do llama, tokenizer do Mistral
  • Eu já achava que o Gemini usava SentencePiece
• Já tentei algo parecido antes: tokie. Na prática, a maior parte do custo de rodar um tokenizer vem da pré-tokenização (regex). Parece que vocês encontraram uma abordagem de regex mais rápida; fico curioso se também compararam a mudança de performance ao trocar só o engine de regex e manter o BPE do tiktoken como está. Se sim, talvez isso até pudesse virar contribuição upstream
  • Projeto muito legal; já entrei em contato com a pessoa que mantém o tiktoken
• Seria bom também comparar a performance com Huggingface tokenizers; pelos benchmarks do readme do tiktoken, a informação já está bem desatualizada
  • Pessoalmente, o tiktoken sempre me pareceu mais lento que o huggingface tokenizers. Ainda não investiguei o tiktoken a fundo para saber por quê, mas essa é a sensação que tenho como usuário dos tokenizers em Rust da HF
• Também queria ver bindings WASM para o tiktoken, ou saber se o ganho de performance do “b” também poderia ser aplicado na implementação pura em js
• Projeto realmente muito legal. Nós também usamos tiktoken, e se houver ganho de performance com compatibilidade drop-in, fico curioso sobre o impacto. A escolha pelo formato drop-in foi excelente