Realize buscas vetoriais no Elasticsearch com o cliente Elasticsearch Go.

Desenvolver software em qualquer linguagem de programação, incluindo Go, é comprometer-se com uma vida inteira de aprendizado. Ao longo de sua trajetória acadêmica e profissional, Carly explorou diversas linguagens de programação e tecnologias, incluindo as mais recentes e melhores implementações de busca vetorial. Mas isso não foi suficiente! Recentemente, Carly também começou a jogar Go.

Assim como os animais, as linguagens de programação e seu amigável autor, a busca passou por uma evolução de diferentes práticas, o que pode dificultar a escolha da melhor opção para o seu caso específico. Neste blog, compartilharemos uma visão geral da busca vetorial, juntamente com exemplos de cada abordagem usando o Elasticsearch e o cliente Elasticsearch Go. Estes exemplos mostrarão como encontrar esquilos-da-terra e determinar o que eles comem usando a busca vetorial no Elasticsearch e em Go.

Pré-requisitos

Para dar continuidade a este exemplo, certifique-se de que os seguintes pré-requisitos sejam atendidos:

1. Instalação do Go versão 1.21 ou posterior

2. Criação do seu próprio repositório Go com o

3. Criação do seu próprio cluster Elasticsearch, preenchido com um conjunto de páginas baseadas em roedores, incluindo para o nosso amigável Gopher, da Wikipédia:

Conectando-se ao Elasticsearch

Em nossos exemplos, utilizaremos a API tipada oferecida pelo cliente Go. Para estabelecer uma conexão segura para qualquer consulta, é necessário configurar o cliente usando um dos seguintes métodos:

1. ID da nuvem e chave de API, caso esteja utilizando o Elastic Cloud.

2. URL do cluster, nome de usuário, senha e certificado.

A conexão com nosso cluster localizado no Elastic Cloud seria feita da seguinte forma:

func GetElasticsearchClient() (*elasticsearch.TypedClient, error) {
	var cloudID = os.Getenv("ELASTIC_CLOUD_ID")
	var apiKey = os.Getenv("ELASTIC_API_KEY")

	var es, err = elasticsearch.NewTypedClient(elasticsearch.Config{
		CloudID: cloudID,
		APIKey:  apiKey,
		Logger:  &elastictransport.ColorLogger{os.Stdout, true, true},
	})

	if err != nil {
		return nil, fmt.Errorf("unable to connect: %w", err)
	}

	return es, nil
}

A conexão client pode então ser usada para busca vetorial, como mostrado nas seções subsequentes.

Busca vetorial

A busca vetorial tenta resolver esse problema convertendo-o em uma comparação matemática usando vetores. O processo de incorporação de documentos possui uma etapa adicional de conversão do documento, utilizando um modelo, em uma representação vetorial densa, ou simplesmente em um fluxo de números. A vantagem dessa abordagem é que você pode pesquisar documentos não textuais, como imagens e áudio, convertendo-os em um vetor juntamente com uma consulta.

Em termos simples, a busca vetorial é um conjunto de cálculos de distância vetorial. Na ilustração abaixo, a representação vetorial da nossa consulta Go Gopheré comparada com os documentos no espaço vetorial, e os resultados mais próximos (denotados pela constante k) são retornados:

Dependendo da abordagem usada para gerar os embeddings dos seus documentos, existem duas maneiras diferentes de descobrir o que os gophers comem.

Abordagem 1: Traga seu próprio modelo

Com uma licença Platinum, é possível gerar os embeddings dentro do Elasticsearch, carregando o modelo e usando a API de inferência. A configuração do modelo envolve seis etapas:

1. Selecione um modelo PyTorch para carregar a partir de um repositório de modelos. Neste exemplo, estamos usando o sentence-transformers/msmarco-MiniLM-L-12-v3 da Hugging Face para gerar os embeddings.

2. Carregue o modelo no Elastic usando o cliente Eland Machine Learning para Python usando as credenciais para nosso cluster Elasticsearch e tipo de tarefa text_embeddings. Se você não tiver o Eland instalado, poderá executar a etapa de importação usando o Docker, conforme mostrado abaixo:

docker run -it --rm --network host \
    docker.elastic.co/eland/eland \
    eland_import_hub_model \
      --cloud-id $ELASTIC_CLOUD_ID \
      --es-api-key $ELASTIC_API_KEY \
      --hub-model-id sentence-transformers/msmarco-MiniLM-L-12-v3 \
      --task-type text_embedding

1. Após o carregamento, teste rapidamente o modelo sentence-transformers__msmarco-minilm-l-12-v3 com um documento de exemplo para garantir que os embeddings sejam gerados conforme o esperado:

1. Crie um pipeline de ingestão contendo um processador de inferência. Isso permitirá que a representação vetorial seja gerada usando o modelo carregado:

PUT _ingest/pipeline/search-rodents-vector-embedding-pipeline
{
  "processors": [
    {
      "inference": {
        "model_id": "sentence-transformers__msmarco-minilm-l-12-v3",
        "target_field": "text_embedding",
        "field_map": {
          "body_content": "text_field"
        }
      }
    }
  ]
}

1. Crie um novo índice contendo o campo text_embedding.predicted_value do tipo dense_vector para armazenar os vetores de incorporação gerados para cada documento:

PUT vector-search-rodents
{
  "mappings": {
    "properties": {
      "text_embedding.predicted_value": {
        "type": "dense_vector",
        "dims": 384,
        "index": true,
        "similarity": "cosine"
      },
      "text": {
        "type": "text"
      }
    }
  }
}

1. Reindexe os documentos usando o pipeline de ingestão recém-criado para gerar os embeddings de texto como o campo adicional text_embedding.predicted_value em cada documento:

POST _reindex
{
  "source": {
    "index": "search-rodents"
  },
  "dest": {
    "index": "vector-search-rodents",
    "pipeline": "search-rodents-vector-embedding-pipeline"
  }
}

Agora podemos usar a opção Knn na mesma API de pesquisa usando o novo índice vector-search-rodents, como mostrado no exemplo abaixo:

func VectorSearch(client *elasticsearch.TypedClient, term string) ([]Rodent, error) {
  var k = 10
	var numCandidates = 10

	res, err := client.Search().
		Index("vector-search-rodents").
		Knn(types.KnnSearch{
      # Field in document containing vector
			Field:         "text_embedding.predicted_value",
      # Number of neighbors to return
			K:             &k,
      # Number of candidates to evaluate in comparison
			NumCandidates: &numCandidates,
      # Generate query vector using the same model used in the inference processor
			QueryVectorBuilder: &types.QueryVectorBuilder{
				TextEmbedding: &types.TextEmbedding{
					ModelId:   "sentence-transformers__msmarco-minilm-l-12-v3",
					ModelText: term,
				},
			}}).Do(context.Background())

	if err != nil {
		return nil, fmt.Errorf("error in rodents vector search: %w", err)
	}

	return getRodents(res.Hits.Hits)
}

A conversão do objeto JSON resultante por meio de desserialização é feita exatamente da mesma forma que no exemplo de pesquisa por palavra-chave. As constantes K e NumCandidates permitem configurar o número de documentos vizinhos a serem retornados e o número de candidatos a serem considerados por fragmento. Note que aumentar o número de candidatos aumenta a precisão dos resultados, mas leva a uma consulta mais demorada, pois mais comparações são realizadas.

Quando o código é executado usando a consulta What do Gophers eat?, os resultados retornados são semelhantes aos abaixo, destacando que o artigo do Gopher contém as informações solicitadas, diferentemente da pesquisa anterior por palavra-chave:

[
  {ID:64f74ecd4acb3df024d91112 Title:Gopher - Wikipedia Url:https://en.wikipedia.org/wiki/Gopher} 
  {ID:64f74ed34acb3d71aed91fcd Title:Squirrel - Wikipedia Url:https://en.wikipedia.org/wiki/Squirrel} 
  //Other results omitted
]

Abordagem 2: API de inferência do Hugging Face

Outra opção é gerar esses mesmos embeddings fora do Elasticsearch e ingeri-los como parte do seu documento. Como essa opção não utiliza um nó de aprendizado de máquina do Elasticsearch, ela pode ser executada no nível gratuito.

A Hugging Face disponibiliza uma API de inferência gratuita e com limite de requisições que, mediante uma conta e um token de API, pode ser usada para gerar manualmente os mesmos embeddings para experimentação e prototipagem, ajudando você a começar. Não é recomendado para uso em produção. Invocar seus próprios modelos localmente para gerar embeddings ou usar a API paga também pode ser feito usando uma abordagem semelhante.

Na função abaixo GetTextEmbeddingForQuery usamos a API de inferência em nossa string de consulta para gerar o vetor retornado de uma solicitação POST ao endpoint:

// HuggingFace text embedding helper
func GetTextEmbeddingForQuery(term string) []float32 {
    // HTTP endpoint
    model := "sentence-transformers/msmarco-minilm-l-12-v3"
    posturl := fmt.Sprintf("https://api-inference.huggingface.co/pipeline/feature-extraction/%s", model)

    // JSON body
    body := []byte(fmt.Sprintf(`{
        "inputs": "%s",
        "options": {"wait_for_model":True}
    }`, term))

    // Create a HTTP post request
    r, err := http.NewRequest("POST", posturl, bytes.NewBuffer(body))

    if err != nil {
        log.Fatal(err)
        return nil
    }

    token := os.Getenv("HUGGING_FACE_TOKEN")
    r.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))

    client := &http.Client{}
    res, err := client.Do(r)
    if err != nil {
        panic(err)
    }

    defer res.Body.Close()

    var post []float32
    derr := json.NewDecoder(res.Body).Decode(&post)

    if derr != nil {
        log.Fatal(derr)
        return nil
    }

    return post
}

O vetor resultante, do tipo []float32 é então passado como QueryVector em vez de usar a opção QueryVectorBuilder para aproveitar o modelo previamente carregado no Elastic.

func VectorSearchWithGeneratedQueryVector(client *elasticsearch.TypedClient, term string) ([]Rodent, error) {
	vector, err := GetTextEmbeddingForQuery(term)
	if err != nil {
		return nil, err
	}

	if vector == nil {
		return nil, fmt.Errorf("unable to generate vector: %w", err)
	}

  var k = 10
	var numCandidates = 10

	res, err := client.Search().
		Index("vector-search-rodents").
		Knn(types.KnnSearch{
      # Field in document containing vector
			Field:         "text_embedding.predicted_value",
      # Number of neighbors to return
			K:             &k,
      # Number of candidates to evaluate in comparison
			NumCandidates: &numCandidates,
      # Query vector returned from Hugging Face inference API
			QueryVector:   vector,
		}).
		Do(context.Background())

	if err != nil {
		return nil, err
	}

	return getRodents(res.Hits.Hits)
}

Note que as opções K e NumCandidates permanecem as mesmas, independentemente das duas opções, e que os mesmos resultados são gerados, pois estamos usando o mesmo modelo para gerar os embeddings.

Conclusão

Aqui discutimos como realizar uma busca vetorial no Elasticsearch usando o cliente Elasticsearch Go. Confira o repositório do GitHub para acessar todo o código desta série. Continue para a parte 3 para obter uma visão geral da combinação da pesquisa vetorial com os recursos de pesquisa por palavra-chave abordados na parte 1 em Go.

Até lá, boa caçada aos esquilos!

Recursos

1. Guia do Elasticsearch

2. Cliente Go do Elasticsearch

3. O que é busca vetorial? | Elástico