Show HN: Wordllama – O que dá para fazer com embeddings de tokens de LLM

 (github.com/dleemiller)
1 pontos por GN⁺ 2024-09-16 | 1 comentários | Compartilhar no WhatsApp
  • WordLlama é um toolkit de NLP rápido e leve que reutiliza embeddings de tokens de LLM para fazer deduplicação fuzzy, cálculo de similaridade, ranking, clustering e segmentação semântica de texto
  • A inferência funciona principalmente com busca de tokens e average pooling, destacando um pipeline leve executável apenas com NumPy e otimizado para CPU
  • O modelo padrão tem 256 dimensões e 16 MB, pode ter a dimensionalidade reduzida com representações Matryoshka, e embeddings binários permitem cálculos mais rápidos com similaridade de Hamming
  • Na tabela MTEB, WL64~WL1024 superam GloVe 300d e Komninos em várias métricas, mas em geral ficam abaixo de all-MiniLM-L6-v2
  • Pode ser usado com pip install wordllama e WordLlama.load(), e .key(query) retorna um callable que pode ser passado para funções da biblioteca padrão como sorted, min e max

O que o WordLlama faz

  • WordLlama é um toolkit leve para tarefas utilitárias de NLP como deduplicação fuzzy, cálculo de similaridade, ranking, clustering e segmentação semântica de texto
  • Ele extrai codebooks de embeddings de tokens de LLMs modernos como LLaMA 2 e LLaMA 3 70B para criar representações compactas de palavras, semelhantes a GloVe, Word2Vec e FastText
  • Na inferência, tem poucas dependências e é otimizado para hardware de CPU, sendo adequado para implantação em ambientes com recursos limitados
  • Por ser rápido e pequeno, pode ser usado em análises exploratórias, avaliadores de saída de LLM e tarefas preparatórias em workflows multi-hop ou agentic

Instalação e uso básico

  • A instalação é feita com pip
pip install wordllama
  • O modelo padrão de 256 dimensões é carregado com WordLlama.load()
from wordllama import WordLlama

wl = WordLlama.load()
  • .key(query) retorna Callable[[str], float], permitindo ordenar strings candidatas por similaridade com a consulta ou escolher a de maior pontuação
query = "Machine learning methods"
candidates = [
    "Foundations of neural science",
    "Introduction to neural networks",
    "Cooking delicious pasta at home",
    "Introduction to philosophy: logic",
]

sim_key = wl.key(query)

sorted_candidates = sorted(candidates, key=sim_key, reverse=True)
best_candidate = max(candidates, key=sim_key)
  • No resultado de exemplo, "Introduction to neural networks" se torna a candidata com a maior pontuação, 0.3414

Principais recursos

  • Geração de embeddings: cria embeddings de texto rapidamente com simples busca de tokens e average pooling
  • Cálculo de similaridade: calcula a cosine similarity entre dois textos
  • Ranking de documentos: ordena documentos com base na similaridade entre consulta e candidatos
  • Deduplicação fuzzy: remove textos duplicados com base em um limiar de similaridade
  • Clustering: agrupa documentos com KMeans
  • Filtragem: mantém apenas documentos cuja similaridade com a consulta esteja acima do critério
  • Busca Top-K: retorna os K documentos mais similares à consulta
  • Segmentação semântica de texto: divide o texto em chunks semanticamente coesos
  • Embeddings binários: permitem cálculos mais rápidos com similaridade de Hamming
  • Representações Matryoshka: cortam a dimensionalidade do embedding conforme necessário para ajustar tamanho do modelo e desempenho

Estrutura do modelo e desempenho

  • O WordLlama treina modelos pequenos sem contexto dentro de um framework geral de embeddings
  • O modelo padrão tem 256 dimensões e 16 MB
  • A tabela MTEB do README compara WL64, WL128, WL256, WL512 e WL1024 com GloVe 300d, Komninos e all-MiniLM-L6-v2
    • O WL256 registra Clustering 33.25, Reranking 52.03, Classification 58.21, Pair Classification 78.22, STS 67.91, CQA DupStack 24.12, SummEval 30.99
    • O GloVe 300d registra, nos mesmos itens, 27.73, 43.29, 57.29, 70.92, 61.85, 15.47, 28.87
    • all-MiniLM-L6-v2 registra Clustering 42.35, Reranking 58.04, Classification 63.05, Pair Classification 82.37, STS 78.90, CQA DupStack 41.32, SummEval 30.81
  • l2_supercat é um modelo com vocabulário do LLaMA 2
    • Ele é treinado conectando codebooks de vários modelos, como LLaMA 2 70B e phi 3 medium, após remover tokens especiais adicionais
    • É possível treiná-lo conectando juntos os codebooks de vários modelos que usam o tokenizer do LLaMA 2
    • Apresenta desempenho semelhante ao treino com codebook do LLaMA 3 70B, mas com vocabulário 4 vezes menor: 32k contra 128k
  • Um modelo baseado em LLaMA 3, l3_supercat, também é oferecido
  • Resultados adicionais estão em Results

Segmentação semântica de texto

  • .split() divide textos longos em chunks semânticos
long_text = "Your very long text goes here... " * 100
chunks = wl.split(long_text, target_size=1536)

print(list(map(len, chunks)))

# Output: [1055, 1055, 1187]
  • target_size é ao mesmo tempo o tamanho desejado e o tamanho máximo
  • O processo de divisão tenta preservar a ordem do texto, a estrutura das frases e, quando possível, a estrutura dos parágrafos
  • Usa embeddings do WordLlama para encontrar índices de divisão mais naturais
  • O tamanho dos chunks de saída pode variar dentro do limite de target_size
  • O tamanho recomendado é de 512 a 2048 caracteres, e o padrão é 1536
  • Se forem necessários chunks maiores, a recomendação é agrupar vários semantic chunks em batch após a divisão
  • Mais detalhes estão na technical overview

Model2Vec e inferência direta

  • Na atualização de 2025-01-04, foi adicionado suporte a Model2Vec static embeddings
  • WordLlama.load_m2v() pode ser usado para carregar modelos Model2Vec
wl = WordLlama.list_configs()

wl = WordLlama.load_m2v("potion_base_8m")  # 256-dim model
wl = WordLlama.load_m2v("m2v_multilingual")  # multilingual model
  • Model2Vec é outra abordagem para criar static embeddings usando PCA
  • O lado do Model2Vec criou um modelo multilíngue e modelos baseados em glove, e afirma ter bons resultados em tarefas de similaridade entre palavras
  • Pode ser conferido em minishlab no Hugging Face
  • WordLlamaInference pode ser usado diretamente com um array de static embeddings no formato (n_vocab, dim) e um tokenizer, em vez do loader
from wordllama import WordLlamaInference
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_pretrained(...)
wl = WordLlamaInference(np_embeddings_ar, tokenizer)

Treinamento e extração de embeddings

  • Modelos de embeddings binários mostraram ganhos mais nítidos em dimensões altas, e para embeddings binários são recomendadas 512 ou 1024 dimensões
  • O modelo L2 Supercat foi treinado por 12 horas em uma única GPU A100 com batch size 512
  • Para extrair embeddings de tokens de modelos LLaMA, é necessário aceitar o contrato de usuário e fazer login no Hugging Face CLI
from wordllama.extract.extract_safetensors import extract_safetensors

extract_safetensors("llama3_70B", "path/to/saved/model-0001-of-00XX.safetensors")
  • Os embeddings geralmente estão no primeiro arquivo safetensors, mas nem sempre
    • pode haver um manifest
    • pode ser necessário inspecionar manualmente para encontrá-los
  • O treinamento usa os scripts do repositório, e é preciso copiar ou modificar uma configuração existente para adicionar um configuration file
pip install wordllama[train]
python train.py train --config your_new_config
python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/
  • A etapa de salvamento grava um modelo para cada dimensão Matryoshka

Atualizações, roadmap e licença

  • Na atualização de 2025-02-01, foi adicionado suporte a callables utilizáveis em funções da biblioteca padrão como sorted, min e max
  • Na atualização de 2024-10-04, foi adicionado o algoritmo de inferência para semantic splitting
  • O roadmap inclui notebooks de exemplo para DSPy evaluator e para pipelines de Retrieval-Augmented Generation, ou RAG
  • Há projetos da comunidade como Gradio Demo HF Space e CPU-ish RAG
  • A licença do projeto é a MIT License

Leituras relacionadas

1 comentários

 
GN⁺ 2024-09-16
Opiniões no Hacker News

  • Gosto muito do tamanho pequeno. Ele já tem vantagens até sobre o menor modelo do SBERT.
    Tecnicamente, porém, parece uma abordagem bem antiga, e entendo que há um compromisso com desempenho. Ainda assim, fico curioso se seria possível oferecer alternância entre tipos de similaridade, como similaridade semântica, inferência em linguagem natural (NLI) e abstração de substantivos.
    Por exemplo, ao agrupar artigos de jornal em uma categoria como “eventos ambientais extremos”, eu gostaria que “Freezing” e “Burning” aparecessem como muito parecidos. É o caso de algo que se comporta como MTEB/Sentence-Similarity ou Word2Vec/GloVe clássicos. Mas, se for um artigo de química, os dois deveriam aparecer quase como opostos, e às vezes quero usar embeddings de inferência em linguagem natural para observar relações causais entre duas coisas.
    Os dois últimos tipos de embedding são abordagens relativamente recentes, pós-2019, então vejo uma oportunidade técnica maior aí. A família mais antiga de MTEB/similaridade semântica já era suficiente para vários usos desde 2014 e melhorou bastante em 2019 com coisas como mini-lm-v2.
    Esses três tipos de embedding também são possíveis com SBERT, mas as dimensões são grandes e os modelos também, então carregar vários modelos por tipo pesa nos recursos. Modelos de embedding generativos, E5 e modelos de inferência em linguagem natural são grandes e muitas vezes exigem algo em torno de 6 GB.

    • Boa ideia. Vou fazer alguns experimentos e verificar a viabilidade.
      Quero ver como fica o desempenho ao treinar com um único tipo de similaridade. Não tenho certeza se há outra forma de lidar com isso sem calcular contexto. Talvez seja preciso alternar modelos, mas isso por si só não é um grande problema.
    • Este é um modelo de 17 MB e, nos benchmarks, naturalmente fica abaixo do MiniLM v2, ou seja, do SBERT. Eu rodo o V3 no ONNX com um modelo de 23 MB em quase todas as plataformas.
      Não digo isso para diminuir o trabalho; é importante entender esse tipo de projeto no seu contexto. Aqui, o contexto é que, ao entender LLMs em profundidade, você descobre que LLMs também têm embeddings e, dessa perspectiva, é mais natural mexer nesses embeddings e avançar um passo do que revisar o estado de toda a área de embeddings.
    • Se “embeddings do ChatGPT” significa os modelos de embedding da OpenAI, “burning” e “freezing” não são opostos de forma alguma. Ao rodar com as 1.024 dimensões do text-embedding-large-3, a similaridade de cosseno fica por volta de 0,46. Se fossem embeddings completamente opostos, a similaridade teria que ser -1.
      É um equívoco comum achar que palavras de sentido oposto têm embeddings opostos. Na prática, palavras de sentido oposto também têm muitas coisas em comum. “burning” e “freezing” têm ambos relação com temperatura e física, são palavras em inglês, podem ser usados como verbo, substantivo e adjetivo, e estão escritos corretamente. Todas essas características entram nos embeddings.

  • Embeddings carregam muita informação semântica, dependendo dos dados de treinamento e da função objetivo, e podem ser usados de forma independente em várias tarefas úteis.
    No passado, usei embeddings do codificador de texto de um modelo CLIP para reforçar prompts de modo que combinassem melhor com as imagens correspondentes. Por exemplo, se o prompt tivesse “building”, eu procurava vizinhos mais próximos na matriz de embeddings, como “concrete” e “underground”, e substituía ou acrescentava esses termos depois da palavra. Em experimentos limitados, isso aumentou o recall na maioria das consultas.

    • Sim. É possível ensinar esse tipo de relação contextual dentro do domínio a um modelo de embeddings.
      https://www.marqo.ai/blog/generalized-contrastive-learning-f...
    • Ideia muito legal. Parece possível nesta implementação também, então vou pensar mais sobre isso.
      Pelo tamanho dos embeddings de tokens no wordllama, acho que eles também podem ajudar a identificar tokens importantes para enriquecimento. Dito isso, se treinar com dados selecionados para essa tarefa, provavelmente funcionará muito melhor.

  • Fico curioso se há planos para idiomas além do inglês. Para o francês, parece que seria uma ferramenta perfeita.

    • É totalmente possível. Seria preciso montar um corpus de treinamento, mas não sei bem quais materiais existem em francês.
      Já fiz um pouco de treinamento com modelos da família Mistral, então acho que começaria por eles em um corpus em francês.
      Se você abrir uma issue, posso trabalhar nisso quando tiver tempo.

  • Para um corpus grande, por exemplo com mais de 10.000 frases tratando cada frase como um documento, dá para obter resultados parecidos agrupando vetores de matriz esparsa TF-IDF com k-means.
    Dito isso, esta ferramenta parece ter várias utilidades para acelerar a parte do k-means, como binarização. Pretendo fazer benchmarks nas próximas semanas.

  • Alguns anos atrás, criei uma coleção de jogos de linguagem que usa funções parecidas: https://github.com/Hellisotherpeople/Language-games

  • Fico imaginando se alguém já pensou em resolver Little Alchemy com embeddings. #sample-use

  • Parece bacana. Fico curioso se há vantagem em relação aos modelos mini-lm. Na maioria das tarefas do MTEB, o mini-lm parece melhor; queria saber se há pontos em que este seja melhor, como velocidade de inferência.

    • O Mini-lm é um modelo de embeddings melhor. Este modelo não faz cálculo de atenção e, depois do treinamento, nem usa framework de deep learning. Por isso, ele não obtém as vantagens contextuais de um modelo Transformer.
      A meta também não era ser o modelo de ponta com melhor desempenho. É um modelo com restrições bastante fortes para reduzir dependências, tamanho e requisitos de hardware, além de aumentar a velocidade.
      Mesmo como modelo de embeddings de palavras, ele é bem leve. Normalmente, modelos desse tipo têm vocabulários muito maiores e chegam a alguns GB.
    • Parece ser uma diferença no tamanho do modelo em si. É mais leve e rápido. O mini-lm tem 80 MB, e o menor modelo aqui tem 16 MB.

  • Parece muito útil para criação de jogos.

  • Mostra bem quanto conteúdo semântico existe nos próprios tokens.

  • Será que daria para transformá-lo em uma extensão do PostgreSQL?