milvus-logo
LFAI
Home
  • Integraciones

Búsqueda semántica con Milvus y OpenAI

Open In Colab GitHub Repository

Esta guía muestra cómo la API de incrustación de OpenAI puede utilizarse con la base de datos vectorial Milvus para realizar búsquedas semánticas en texto.

Para empezar

Antes de empezar, asegúrese de que tiene preparada la clave API de OpenAI, o consiga una en el sitio web de OpenAI.

Los datos utilizados en este ejemplo son títulos de libros. Puedes descargar el conjunto de datos aquí y ponerlo en el mismo directorio donde ejecutes el siguiente código.

Primero, instala el paquete para Milvus y OpenAI:

pip install --upgrade openai pymilvus

Si utilizas Google Colab, para habilitar las dependencias que acabas de instalar, es posible que tengas que reiniciar el tiempo de ejecución. (Haga clic en el menú "Runtime" en la parte superior de la pantalla, y seleccione "Reiniciar sesión" en el menú desplegable).

Con esto, ya estamos listos para generar los embeddings y utilizar la base de datos vectorial para realizar la búsqueda semántica.

Búsqueda de títulos de libros con OpenAI y Milvus

En el siguiente ejemplo, cargamos los datos de los títulos de los libros desde el archivo CSV descargado, utilizamos el modelo de incrustación de OpenAI para generar representaciones vectoriales y las almacenamos en la base de datos vectorial Milvus para realizar búsquedas semánticas.

from openai import OpenAI
from pymilvus import MilvusClient

MODEL_NAME = "text-embedding-3-small"  # Which model to use, please check https://platform.openai.com/docs/guides/embeddings for available models
DIMENSION = 1536  # Dimension of vector embedding

# Connect to OpenAI with API Key.
openai_client = OpenAI(api_key="<YOUR_OPENAI_API_KEY>")

docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

vectors = [
    vec.embedding
    for vec in openai_client.embeddings.create(input=docs, model=MODEL_NAME).data
]

# Prepare data to be stored in Milvus vector database.
# We can store the id, vector representation, raw text and labels such as "subject" in this case in Milvus.
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(docs))
]


# Connect to Milvus, all data is stored in a local file named "milvus_openai_demo.db"
# in current directory. You can also connect to a remote Milvus server following this
# instruction: https://milvus.io/docs/install_standalone-docker.md.
milvus_client = MilvusClient(uri="milvus_openai_demo.db")
COLLECTION_NAME = "demo_collection"  # Milvus collection name
# Create a collection to store the vectors and text.
if milvus_client.has_collection(collection_name=COLLECTION_NAME):
    milvus_client.drop_collection(collection_name=COLLECTION_NAME)
milvus_client.create_collection(collection_name=COLLECTION_NAME, dimension=DIMENSION)

# Insert all data into Milvus vector database.
res = milvus_client.insert(collection_name="demo_collection", data=data)

print(res["insert_count"])

En cuanto al argumento de MilvusClient:

  • Establecer el uri como un archivo local, por ejemplo./milvus.db, es el método más conveniente, ya que utiliza automáticamente Milvus Lite para almacenar todos los datos en este archivo.
  • Si tiene una gran escala de datos, puede configurar un servidor Milvus más eficiente en docker o kubernetes. En esta configuración, por favor utilice la uri del servidor, por ejemplohttp://localhost:19530, como su uri.
  • Si desea utilizar Zilliz Cloud, el servicio en la nube totalmente gestionado para Milvus, ajuste uri y token, que corresponden al punto final público y a la clave Api en Zilliz Cloud.

Con todos los datos en la base de datos vectorial de Milvus, ahora podemos realizar la búsqueda semántica generando la incrustación vectorial para la consulta y realizar la búsqueda vectorial.

queries = ["When was artificial intelligence founded?"]

query_vectors = [
    vec.embedding
    for vec in openai_client.embeddings.create(input=queries, model=MODEL_NAME).data
]

res = milvus_client.search(
    collection_name=COLLECTION_NAME,  # target collection
    data=query_vectors,  # query vectors
    limit=2,  # number of returned entities
    output_fields=["text", "subject"],  # specifies fields to be returned
)

for q in queries:
    print("Query:", q)
    for result in res:
        print(result)
    print("\n")

Debería ver lo siguiente como salida:

[
    {
        "id": 0,
        "distance": -0.772376537322998,
        "entity": {
            "text": "Artificial intelligence was founded as an academic discipline in 1956.",
            "subject": "history",
        },
    },
    {
        "id": 1,
        "distance": -0.58596271276474,
        "entity": {
            "text": "Alan Turing was the first person to conduct substantial research in AI.",
            "subject": "history",
        },
    },
]

Traducido porDeepL

Try Managed Milvus for Free

Zilliz Cloud is hassle-free, powered by Milvus and 10x faster.

Get Started
Feedback

¿Fue útil esta página?