Orquestrando Small Language Models (SLM) usando JavaScript e a API de Inferência do Hugging Face

English Version

Olá!

Vou mostrar como usei a API de Inferência do Hugging Face, um Space, Docker e menos de 500 linhas de código JavaScript para orquestrar vários pequenos LLMs, fazendo com que eles me respondam sobre tentativas em um pequeno simulador interativo de redes neurais.

Você pode ver o resultado nest post do blog IATalking. Mova o slider e tente diminuir o erro; você verá que uma IA gerará comentários com base na sua tentativa (e histórico de tentativas). Esses comentários vêm de modelos como o Phi3, Llama, Mystral, etc., todos pequenos, com alguns poucos bilhoes de parâmetros.

Sempre que um novo request é enviado, o código que vou apresentar aqui escolhe um desses modelos. Se um modelo começa a falhar, as chances de ser escolhido diminuem na próxima vez que um request for feito. Assim, consigo ter uma espécie de "alta disponibilidade" de LLMs, usando diferentes modelos e criando variações nas respostas, além de ajustar a temperatura de um único LLM para ser um pouco mais criativo a cada requisição.

Este post mostrará os detalhes de como fiz isso, explicando os arquivos envolvidos.

O Space

Você encontrará o código deste modelo no Space: Jay Trainer. A API responde nesse link: https://iatalking-jaytrainer.hf.space/error?error=123

Eu o nomeei de Jay Trainer em referência a Jay Alammar (https://jalammar.github.io/), que fez o post original e trouxe a ideia de gerar um simulador simples de rede neural para facilitar o aprendizado. O post dele é incrível: Visual Interactive Guide Basics Neural Networks.

Olhando os arquivos, você notará o seguinte:

• Um Dockerfile
• docker-compose.yml
• server.js
• README.md

Este é um Docker Space, e escolhi Docker pela simplicidade. O Hugging Face permite que você suba APIs usando Docker, o que é sensacional, pois abre possibilidades infinitas: você pode subir uma API na sua linguagem de preferência!

Se olhar no Dockerfile, verá um servidor Node.js padrão:

O docker-compose.yml é apenas para facilitar meus testes locais. Antes de enviar pro Hugging Face, eu subo o mesmo código com apenas um docker compose up. Isso me poupa o trabalho de ficar digitando "docker -p etc. etc."

O README.md contém apenas o necessário exigido pelo Hugging Face e a grande estrela é o server.js. É aqui que está toda a lógica e e código dos endpoints!

API disponibilizada

A API, que é implementada no server.js, é uma aplicação Express. Os seguintes endpoints são criados:

• /error
  Aqui é o mais usado: você passa o número de erro e o histórico de tentativas, e ele devolve um texto!
• /
  Aqui é o endpoint principal. Ele apenas vai exibir uma mensagem simples, indicando que o serviço está no ar!
• /models
  Aqui é um endpoint de debug. Com ele, eu consigo ver o histórico dos LLMs envolvidos e as estatísticas de execução. Isso me permite saber quem está gerando mais erros e quem está sendo mais executado, junto com algumas outras informações apenas para acompanhamento e debug!
• /test
  Este é um endpoint apenas para testar se o Express está funcionando mesmo e respondendo.

Inicialização do servidor

Ao iniciar, o servidor faz algumas configurações, e a principal delas é definir um objeto contendo a lista de modelos que quero usar:

A constante global MODELS vai conter o meu modelo! Cada objetivo é um "id" de um LLM que será usado. id aqui é algo interno que eu defini para este código, não um Id do Hugging Face. É só uma espécie de "alias" para facilitar a identificação.

Então, para cada um deles, preciso definir: o nome do modelo e o template do prompt. O nome é o nome único e exclusivo no Hugging Face, no formato ORG/MODEL. Esse é o nome que você encontra no Model Card:

prompt é um método que irei chamar quando precisar gerar um prompt para este modelo. Cada modelo pode ter um formato diferente de prompt, então não posso usar o mesmo prompt para todos. Através dessa função, consigo criar um mecanismo dinâmico de gerar prompts usando a mesma chamada: apenas preciso chamar o método prompt() e passar o texto que quero.

Eu nunca estudei o código dos Transformers do Hugging Face, mas é muito provável que isso siga a mesma ideia por trás dos métodos apply_chat_template. Aqui, é apenas uma versão bem simples do processo!

Portanto, com apenas estes dois membros, consigo adicionar ou remover LLMs facilmente! Optei por deixar hard-coded para ter controle e simplificar as coisas. Mas eu poderia ter jogado isso externamente para um arquivo ou API, de forma a conseguir incluir ou remover LLM em runtime. Como isso é uma simples PoC, preferi manter simples por agora.

Por último, é importante lembrar que o código também tem uma validação na variável de ambiente HF_TOKEN, que é o token do Hugging Face. Esse token é configurado como uma secret no Space. Nos testes locais, eu gero um token de testes e uso no meu docker compose up, criando uma variável de ambiente aqui no meu shell. Outra facilidade que deixar um docker compose pronto me traz!

Inicialização dos modelos

Ainda no startup do servidor, para cada um dos modelos em MODELS, eu vou gerar um terceiro atributo: stats. Aqui, teremos algumas métricas de execução do LLM, como o total de execuções, o total de erros e o percentual de chances dele gerar erros.

A descrição das propriedades é:

• total
  Total de execuções.
• erros
  Um número que indica o quanto do total foi de baixa qualidade.
• parcela
  Parcela de erro. Isso será calculado com base no valor de erros.
• errop
  % de erros: erros/total.
• pok
  % de ok (contrário de errop): 1 - errop.

E, nesse trecho, monto um array com todos os LLMs, que está na variável ModelList.

O restante do código se dedica aos endpoints do Express e às funções com a lógica, que serão detalhadas melhor a seguir.

Endpoint /error

Este é o endpoint principal. É este endpoint que a página acima chama. Ele espera os seguintes parâmetros:

• error
  Valor de erro da tentativa atual do usuário.

• tentativas
  Histórico de tentativas. São os números de erros, separados por ",".

O código começa fazendo algumas validações nos valores dos parâmetros. Basicamente, estou garantindo que error seja um número e que tentativas seja uma lista de números separados por vírgula. Qualquer coisa diferente disso, eu retorno para o usuário. A ideia é evitar um prompt injection, já que (como você vai ver já já), eu vou concatenar esses valores diretamente.

Então, ele vai chamar a função Prompt() (que falaremos dela logo em seguida), que é quem vai montar o prompt e pedir para o LLM usando a API de Inferência do Hugging Face.

A função vai retornar a resposta do LLM. Aqui serão feitas algumas validações.

Dentre as validações, as mais importantes são extrair o texto do LLM somente até a marca de "fim". Você verá que peço ao LLM que sempre encerre o texto com uma marca que chamei de |fim|. Isso é uma tentativa de fazer ele me devolver um ponto sinalizando que até ali ele gerou o texto correto. Dali para frente, não há garantia. Então, eu vou pegar somente o resultado até antes do |fim|.

Outra validação que faço é na quantidade de caracteres. Aqui, assumo um padrão de no máximo 8 caracteres por palavra, e como deixei um default de 20 palavras, faço a conta simples de 20*8.

E, mais adiante você entenderá o porquê, mas note que há alguns trechos em que estamos alterando a variável de erros do modelo que respondeu. É um mecanismo que criei para "penalizar" as respostas que fugiram do padrão de qualidade que eu esperava.

Essa validação tem muitas brechas e poderia ser melhor. Mas, novamente: isso é uma pequena PoC para um blog, com o objetivo de aprender a usar mais a API e interagir com LLMs pequenos. Por isso, não fiz muito além do básico.

Função prompt()

A função prompt é a função que será chamada sempre que o endpoint /error for usado.

A missão dessa função é montar o prompt com base no erro que tenho. Meu objetivo é gerar mensagens com base na tentativa do usuário. Ele está tentando produzir um valor menor do que 450, então gero um prompt com base no valor atual. A ideia é tentar fazer o LLM ser engraçado e brincar com o usuário se ele estiver muito abaixo ou muito acima do valor. Se ele estiver próximo, gero uma mensagem mais motivacional do que engraçada.

Eu poderia ter colocado tudo em um único prompt, mas acho que isso seria ineficiente em vários níveis. Primeiro, eu geraria exemplos para condições que não seriam usadas. Por exemplo, se o erro é de 2000, não há necessidade de enviar o prompt explicando para ele gerar uma mensagem motivacional. Eu só preciso do prompt que gere uma mensagem de brincadeira (definido como qualquer coisa acima de 2000).

Com isso, economizo o contexto do meu LLM, que é pequeno e possui limitações! Outra vantagem de fazer isso é que ele fica mais preciso. Com menos tokens no meio dos exemplos, a probabilidade dele gerar tokens parecidos com os meus exemplos é muito melhor. A diferença foi absurda na resposta quando fiz isso. Ele passou a gerar muito mais próximo dos exemplos para a faixa de erro do que se eu colocasse tudo. Aqui, a engenharia de prompt pura me ajudou a extrair o melhor do LLM.

E, como já mencionei, além do limite de caracteres (que deixei fixo em 20 palavras, por enquanto), eu também peço que ele encerre com um "|fim|". Adicionar essa marcação filtrou ainda mais os casos em que ele alucinava após as 20 primeiras palavras. Ele geralmente começa a alucinar depois do "|fim|". E como eu pego tudo antes dessa marca, isso reduziu consideravelmente os casos de mensagens sem nenhuma relação.

E então, uma vez que o meu prompt está pronto, eu chamo a função GetModelAnswer, que cuida da parte técnica da coisa, isto é, escolher o melhor LLM e gerenciar quando um LLM não responde!

Função GetModelAnswer()

Esta função é responsável por enviar o meu prompt a um dos LLMs da lista em MODELS.

A ideia é simples: se o usuário não escolheu um modelo explicitamente no parâmetro, então eu vou tentar escolher o melhor modelo e invocá-lo. Se ele falhar, tento o próximo até acabar as opções!

Já já veremos sobre a UpdateProbs, que é quem calcula o melhor modelo para ser chamado. No loop, que vai repetir no máximo o número de vezes equivalente ao número de modelos que tenho na lista (o array ModelList, que carregamos na inicialização, como mostrado acima), ele obterá o objeto que representa nosso modelo:

Ele monta a URL para a API de inferência do Hugging Face (que poderia ser uma variável de ambiente, para o caso de um dia mudar). Lembra do "name" que definimos lá na constante MODELS? Então, é aqui que vamos usá-lo! Por isso, ele tem que ser exatamente o mesmo nome.

Então, ele vai montar os dados para enviar à API do Hugging Face:

A API de Inferência tem várias opções. No nosso caso, precisamos usar o inputs, que é o nosso prompt. Aqui, note que estou chamando o método prompt do modelo atual, que é quem vai montar o prompt específico para este modelo (lembra, falamos dele lá em cima!). Graças a esse método, posso formatar o prompt de acordo com as necessidades de cada LLM.

As chaves parameters e options são configurações que deixei hard coded: máximo de 70 tokens e temperatura de 50% para evitar muita alucinação, mas ainda ter variações entre as chamadas.

Por fim, um simples fetch para enviar a requisição para a API do Hugging Face e aguardar a resposta! Aqui, também valeria adicionar algumas tratativas, como timeout, etc., mas optei por manter simples.

Note que, neste momento, também incremento o contador de total de requisições nas estatísticas desse modelo. Isso ajudará a gerar as probabilidades de acerto.

Adicionalmente, eu vou medir o tempo que levou, em millisegundos, para isso rodar. E irei usar isso ja já para penalizar ou premiar o modelo.

Se a resposta for um erro (diferente de HTTP 200), incrementamos o contador de erros e tentamos o próximo modelo da lista. Eu vou repetir esse processo até receber um HTTP 200 ou até acabar os modelos que já tentei. Toda a lógica dentro deste IF é apenas para isso: pegar o próximo modelo da lista e repetir todo o loop!

Mais abaixo, eu faço alguns checks de qualidade: Se o tempo de resposta foi maior que 2.5 segundos, eu incremento um pouco a variável de erros. Ou, se o tempo de resposta for menor que 900ms, eu removo um pouco do erro. Com isso, eu consigo penalizar ou premiar o modelo baseado no tempo de resposta. Você vai ver que isso vai interferir no modelo mais escolhido. Eu poderia adicionar aqui mais checks. Basta manipular o valor da variável erros de cada modelo, e isso vai se refletir no algoritmo que define o melhor modelo.

E, se deu um HTTP 200, significa que tenho a resposta e então retorno direto, o que encerra meu loop!

Função UpdateProbs

Esta é a função que contém a lógica para definir qual seria o melhor LLM a ser usado. A ideia aqui é escolher o modelo que tem mais chances de produzir uma resposta com qualidade. As chances são controladas pela propriedade stats.erros de cada modelo. Esse valor é um percentual relativo ao stats.total, que é o total de tentativas. A ideia aqui é: Se stats.erros é igual ao stats.total, então, a chance do modelo produzir uma resposta de baixa qualidade (ou de falhar) é de 100%.

Para decidir qual do modelos é o melhor, nós calculamos o total de erros e dividimos um percentual entre todos os modelos. Por exemplo, vamos considerar 3 LLMs: Gemini, Phi3 e Llama. Vamos supor que as tentativas e erros foram estas:

• Gemini, erros = 2, tentativas = 2
• LAma, erros = 0, tentativas = 2
• Phi3, erros = 1, tentativas = 2

Note que o Gemini tem 0% de acerto (2 erros em 2 tentativas), Llama tem 100% de acerto e Phi3 tem 50% de acerto.
Logo, se você distribuir o total de acertos, terá o seguinte:

• Gemini: 0% ( 0/(0+100+50) )
• Phi3: 33% ( 50/(0+100+50) )
• Lama: 66% ( 100/(0+100+50) )

Escolha um número aleatório entre 0% e 100%. Pegue o primeiro modelo que cubra a faixa do número aleatório. Por exemplo, se pegar 30%, você pode escolher o Phi3, pois ele compreende a faixa de 0 a 33%. Se pegar 35%, somente o Lama estaria apto, pois ele compreende a faixa de 33 a 100%.

Com isso, nós consegumos priorizar quem mais acerta. Se 2 llm tem chances iguais, um pequeno truque vai fazer um deles ser o escolhido:

Então, nós conseguimos escolher aleatoriamente o melhor modelo baseado no seu desempenho, que nesse caso, é medido por alguns checks simples. Mas o código está pronto para permitir mais métricas (bastando incrementar o erro quando essas condições adicionais de erro forem atingidas).

Demais endpoints

Os demais endpoints, como o /models são apenas para debug. Então, acredito que não precise de nenhuma explanação adicional. Você pode conferir diretamente no código ou acessar em realtime:

Espero que tenha curtido e que essa implementação possa te ajudar a gerar mais ideias de uso de vários LLM no Hugging Face!!!

Qualquer dúvida, é só me procurar!

X: @IATalking
LinkedIn: @rodrigoribeirogomes