Chunk embedding, Embedding de bloques de textos


from sentence_transformers import SentenceTransformer
import json

def get_vectors(chunk):
#model_name = "sentence-transformers/all-mpnet-base-v2"
model_name = "all-MiniLM-L6-v2"
model = SentenceTransformer(model_name)

texts = [data['text'] for data in chunk] # Obtener la lista de textos del chunk

embeddings = model.encode(texts, batch_size=100).tolist() # Obtener los embeddings

vector_data = {
"vectors": [{
"id": data['id'],
"metadata": data['metadata'],
"values": embedding
} for data, embedding in zip(chunk, embeddings)]
}

json_data = json.dumps(vector_data)
return json_data

 

  1. Importa las bibliotecas necesarias: SentenceTransformer para obtener los embeddings de las oraciones y json para trabajar con formato JSON.
  2. Define el modelo de transformación de oraciones que se utilizará. En este caso, el modelo se llama «all-MiniLM-L6-v2».
  3. Extrae los textos de cada elemento en el chunk y los guarda en una lista llamada texts.
  4. Utiliza el modelo de transformación de oraciones para generar los embeddings de los textos en el chunk. Los embeddings se generan en lotes de tamaño 100.
  5. Convierte los embeddings generados en una lista de Python utilizando el método tolist().
  6. Crea un diccionario llamado vector_data que contiene la lista de vectores generados. Cada vector está representado por un diccionario que contiene las claves «id», «metadata» y «values».
  7. Convierte el diccionario vector_data a formato JSON utilizando el método json.dumps().
  8. Devuelve el resultado como una cadena JSON.

En resumen, esta función toma un chunk de datos que contiene textos y metadatos, utiliza un modelo de transformación de oraciones para generar embeddings para cada texto y devuelve los resultados en formato JSON.

https://colab.research.google.com/drive/1aO_itojfM1u4IyHKwuz9oHEBq_pKk12i#scrollTo=jzatIfVCj7_Y