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
- Importa las bibliotecas necesarias:
SentenceTransformer
para obtener los embeddings de las oraciones yjson
para trabajar con formato JSON. - Define el modelo de transformación de oraciones que se utilizará. En este caso, el modelo se llama «all-MiniLM-L6-v2».
- Extrae los textos de cada elemento en el
chunk
y los guarda en una lista llamadatexts
. - 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. - Convierte los embeddings generados en una lista de Python utilizando el método
tolist()
. - 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». - Convierte el diccionario
vector_data
a formato JSON utilizando el métodojson.dumps()
. - 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