Realiza búsqueda vectorial en Elasticsearch con el cliente Elasticsearch Go

Desarrollar software en cualquier lenguaje de programación, incluido Go, es comprometer con una vida entera de aprendizaje. A lo largo de su carrera universitaria y profesional, Carly experimentó con muchos lenguajes de programación y tecnologías, incluyendo las últimas y mejores implementaciones de búsqueda vectorial. ¡Pero eso no fue suficiente! Así que recientemente Carly también empezó a jugar con Go.

Al igual que los animales, los lenguajes de programación y tu autor amable, la búsqueda experimentó una evolución de diferentes prácticas que pueden ser difíciles de elegir para tu propio caso de uso. En este blog, compartiremos una visión general de la búsqueda vectorial junto con ejemplos de cada enfoque usando Elasticsearch y el cliente Elasticsearch Go. Estos ejemplos te mostrarán cómo encontrar topos y determinar qué comen usando búsqueda vectorial en Elasticsearch y Go.

Prerrequisitos

Para seguir con este ejemplo, cerciorar de cumplir los siguientes requisitos:

1. Instalación de Go versión 1.21 o posterior

2. Creación de tu propio repositorio Go con el

3. Creación de tu propio clúster Elasticsearch, poblado con un conjunto de páginas basadas en roedores, incluyendo para nuestro amable Gopher, de Wikipedia:

Conexión con Elasticsearch

En nuestros ejemplos, emplearemos la API tipada que ofrece el cliente Go. Establecer una conexión segura para cualquier consulta requiere configurar al cliente usando cualquiera de:

1. Cloud ID y clave API si se emplea Elastic Cloud.

2. URL del clúster, nombre de usuario, contraseña y el certificado.

Conectarse a nuestro clúster ubicado en Elastic Cloud sería así:

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
}

La conexión client puede entonces emplear para búsqueda vectorial, como se muestra en secciones posteriores.

Búsqueda de vectores

La búsqueda vectorial intenta resolver este problema convirtiendo el problema de búsqueda en una comparación matemática usando vectores. El proceso de incrustación de documentos tiene una etapa adicional de convertir el documento usando un modelo en una representación vectorial densa, o simplemente en un flujo de números. El beneficio de este enfoque es que puedes buscar documentos no textuales, como imágenes y audio, traduciéndolos a un vector junto con una consulta.

En términos simples, la búsqueda vectorial es un conjunto de cálculos de distancia vectorial. En la ilustración siguiente, la representación vectorial de nuestro Go Gopherde consulta se compara con los documentos en el espacio vectorial, y se devuelven los resultados más cercanos (denotados por kconstantes):

Dependiendo del método que se emplee para generar los embeddings de tus documentos, hay dos formas diferentes de saber qué comen los topos.

Enfoque 1: Trae tu propio modelo

Con una licencia Platinum, es posible generar las incrustaciones dentro de Elasticsearch subiendo el modelo y usando la API de inferencia. Hay seis pasos implicados para establecer el modelo:

1. Selecciona un modelo de PyTorch para subirlo desde un repositorio de modelos. Para este ejemplo, estamos usando los transformadores de frase/msmarco-MiniLM-L-12-v3 de Hugging Face para generar los embeddings.

2. Carga el modelo en Elastic usando el cliente Eland Machine Learning para Python usando las credenciales de nuestro clúster Elasticsearch y tipo de tarea text_embeddings. Si no tienes Eland instalado, puedes ejecutar el paso de importación usando Docker, como se muestra a continuación:

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. Una vez cargado, prueba rápidamente el sentence-transformers__msmarco-minilm-l-12-v3 del modelo con un documento de muestra para cerciorarte de que las incrustaciones se generan como se espera:

1. Crear una tubería de ingesta que contenga un procesador de inferencia. Esto permitirá generar la representación vectorial usando el modelo subido:

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. Crea un nuevo índice que contenga el text_embedding.predicted_value de campo de tipo dense_vector para almacenar las incrustaciones vectoriales generadas 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. Reindexa los documentos usando la nueva tubería de ingesta para generar las incrustaciones de texto como el campo adicional text_embedding.predicted_value en cada documento:

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

Ahora podemos usar la opción Knn en la misma API de búsqueda usando el nuevo índice vector-search-rodents, como se muestra en el siguiente ejemplo:

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

La conversión del objeto de resultado JSON mediante desmaravillamiento se realiza exactamente de la misma manera que en el ejemplo de búsqueda por palabra clave. Las constantes K y NumCandidates nos permiten configurar el número de documentos vecinos que devolver y el número de candidatos a considerar por fragmento. Ten en cuenta que aumentar el número de candidatos incrementa la precisión de los resultados pero conduce a una consulta de mayor duración a medida que se realizan más comparaciones.

Cuando el código se ejecuta usando la consulta What do Gophers eat?, los resultados devuelto se ven similares a los siguientes, destacando que el artículo de Gopher contiene la información aplicar, a diferencia de la búsqueda por palabra clave anterior:

[
  {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
]

Enfoque 2: API de inferencia de rostros de abrazo

Otra opción es generar estas mismas incrustaciones fuera de Elasticsearch e ingerirlas como parte de tu documento. Como esta opción no emplea un nodo de aprendizaje automático de Elasticsearch, puede hacer en la capa libre.

Hugging Face expone una API de inferencia gratis y limitada a la velocidad que, con una cuenta y un token de API, puede usar para generar manualmente las mismas incrustaciones para experimentación y prototipado que te ayude a empezar. No se recomienda para uso en producción. Invocar tus propios modelos localmente para generar embeddings o usar la API de pago también puede hacer con un enfoque similar.

En la función GetTextEmbeddingForQuery siguiente usamos la API de inferencia contra nuestra cadena de consulta para generar el vector devuelto de una petición POST al punto final:

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

El vector resultante, de tipo []float32 , se pasa entonces como QueryVector en lugar de usar la opción QueryVectorBuilder para aprovechar el modelo previamente subido a 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)
}

Cabe señalar que las opciones de K y NumCandidates siguen siendo las mismas independientemente de las dos opciones y que se generan los mismos resultados, ya que estamos usando el mismo modelo para generar las incrustaciones

Conclusión

Aquí discutimos cómo realizar búsqueda vectorial en Elasticsearch usando el cliente Elasticsearch Go. Echa un vistazo al repositorio de GitHub para ver todo el código de este serial. Sigue la parte 3 para obtener una visión general de cómo combinar la búsqueda vectorial con las capacidades de búsqueda por palabras clave que se tratan en la primera parte de Go.

Hasta entonces, ¡feliz caza de topos!

Recursos

1. Guía Elasticsearch

2. Cliente Elasticsearch Go

3. ¿Qué es la búsqueda vectorial? | Elástico