Extrae datos estructurados con llamadas a funciones

Ver en la IA de Google Ejecutar en Google Colab Ver el código fuente en GitHub

En este instructivo, analizarás un ejemplo de extracción de datos estructurados usando la API de Gemini para extraer las listas de personajes, relaciones, cosas y lugares de una historia.

Configuración

pip install -U -q google-generativeai
import pathlib
import textwrap

import google.generativeai as genai


from IPython.display import display
from IPython.display import Markdown

from google.api_core import retry

def to_markdown(text):
  text = text.replace('•', '  *')
  return Markdown(textwrap.indent(text, '> ', predicate=lambda _: True))

Una vez que tengas la clave de API, pásala al SDK. Puedes hacerlo de dos maneras:

  • Coloca la clave en la variable de entorno GOOGLE_API_KEY (el SDK la recogerá automáticamente desde allí).
  • Pasa la clave a genai.configure(api_key=...)
try:
    # Used to securely store your API key
    from google.colab import userdata

    # Or use `os.getenv('API_KEY')` to fetch an environment variable.
    GOOGLE_API_KEY=userdata.get('GOOGLE_API_KEY')
except ImportError:
    import os
    GOOGLE_API_KEY = os.environ['GOOGLE_API_KEY']

genai.configure(api_key=GOOGLE_API_KEY)

La tarea de ejemplo

Para este instructivo, extraerás entidades de historias en lenguaje natural. A continuación, verás una historia escrita por Gemini.

new_story = False

if new_story:
  model = genai.GenerativeModel(model_name='models/gemini-1.5-pro-latest')

  response = model.generate_content("""
      Write a long story about a girl with magic backpack, her family, and at
      least one other charater. Make sure everyone has names. Don't forget to
      describe the contents of the backpack, and where everyone and everything
      starts and ends up.""", request_options={'retry': retry.Retry()})
  story = response.text
  print(response.candidates[0].citation_metadata)
else:
  story = """In the quaint town of Willow Creek, nestled amidst rolling hills and whispering willows, resided a young girl named Anya. As she stepped out of the creaky wooden door of her modest cottage, her heart skipped a beat with excitement and anticipation. Today was her first day of school, and she couldn't wait to show off her prized possession - a magical backpack.\n\nHanded down to her from her grandmother, the backpack was no ordinary satchel. Its soft, emerald-green fabric shimmered with an ethereal glow, and its leather straps held secrets that only Anya knew. Within its capacious interior lay an enchanted world, filled with wonders that would ignite her imagination and change her life forever.\n\nAnya's parents, kind-hearted Elise and wise-bearded Edward, bid her farewell with warm embraces. "Remember, my dear," whispered her mother, "use your magic wisely and for good." Her father added, "Always seek knowledge, and let the backpack be your trusted companion."\n\nWith a skip in her step, Anya set off towards the town's only schoolhouse. On her way, she passed her best friend, Samuel, a curious and adventurous boy with a mischievous grin. "Hey, Anya," he called out. "Can I see your backpack?"\n\nAnya hesitated for a moment before unzipping the flap and revealing its contents. Samuel's eyes widened in amazement as he peered inside. There, nestled amidst pencils and notebooks, were a shimmering sword, a book of ancient spells, a tiny compass that always pointed north, and a magical key that could open any lock.\n\nTogether, they marveled at the backpack's wonders, promising to keep its secrets safe. As they approached the schoolhouse, Anya noticed a group of older children huddled together, their faces etched with fear. Curiosity getting the better of her, she cautiously approached.\n\n"What's wrong?" she asked.\n\nA tall, lanky boy stepped forward. "There's a monster in the forest," he stammered. "It's been terrorizing the town, attacking animals and even people."\n\nAnya's heart sank. The town of Willow Creek was small and peaceful, and the thought of a monster brought a shiver down her spine. She knew she had to do something to protect her family and friends.\n\nWithout a moment's hesitation, Anya opened her backpack and retrieved the shimmering sword. With a determined gleam in her eye, she turned to her terrified peers. "Don't worry," she said, her voice steady. "I'll take care of it."\n\nWith Samuel close behind her, Anya ventured into the shadowy depths of the forest. The trees seemed to whisper secrets as she passed, and the undergrowth rustled with unseen creatures. As they walked deeper into the forest, the air grew heavy and the ground beneath their feet trembled.\n\nSuddenly, they came to a clearing, and there before their eyes was the monster - a massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease. The creature roared, a thunderous sound that shook the forest to its core.\n\nFear surged through Anya, but she refused to let it consume her. She drew the sword from its sheath and charged towards the monster. The blade shimmered in the sunlight, and as it struck the beast's hide, a blinding light erupted, enveloping everything in its radiance.\n\nWhen the light faded, the monster was gone, and in its place was a pile of shattered crystals. Anya had defeated the creature with the magic of her backpack, proving that even the smallest of objects could hold the greatest of powers.\n\nAs she and Samuel returned to the town, they were greeted as heroes. The people of Willow Creek rejoiced, and the legend of Anya, the girl with the magic backpack, was passed down through generations. And so, Anya continued her adventures, using the backpack's wonders to make the world a better place, one magical step at a time."""
to_markdown(story)

En la pintoresca ciudad de Willow Creek, enclavada en medio de colinas ondulantes y sauces susurrantes, vivió una joven llamada Anya. Cuando salía por la crujiente puerta de madera de su sencilla casa de campo, su corazón le daba un brinco de emoción y expectativa. Hoy fue su primer día de escuela y no podía esperar para mostrar su preciada posesión: una mochila mágica.

La mochila que le entregó su abuela no era una mochila común. Su suave tela verde esmeralda brillaba con un brillo etéreo, y las correas de cuero contenían secretos que solo Anya sabía. En su espacio interior, había un mundo encantado, lleno de maravillas que encenderían su imaginación y cambiarían la vida para siempre.

Los padres de Anya, la bondadosa Elise y el sabio Edward, despiden su despedida con cariñosos abrazos. "Recuerda, querida", le susurró su madre, "usa tu magia con prudencia y para el bien". El padre agregó lo siguiente: "Busca siempre el conocimiento y deja que la mochila sea tu compañera de confianza".

Con un salto en su paso, Anya se dirige a la única escuela del pueblo. En el camino, se cruza con su mejor amigo, Samuel, un muchacho curioso y aventurero con una sonrisa traviesa. "Hola, Anya", dijo. "¿Puedo ver tu mochila?"

Anya dudó un momento antes de descomprimir la solapa y revelar su contenido. Los ojos de Samuel se ampliaron asombrados mientras miraba dentro. Allí, entre lápices y cuadernos, había una espada reluciente, un libro de hechizos antiguos, una pequeña brújula que siempre apuntaba al norte y una llave mágica que podía abrir cualquier cerradura.

Juntos, se maravillaron con las maravillas de la mochila y prometieron proteger sus secretos. Cuando se acercaron a la escuela, Anya observó que un grupo de niños mayores estaban juntados, con los rostros grabados de miedo. Con su curiosidad, se acercó más a ella con cautela.

"¿Qué sucede?", preguntó.

Un niño alto y larguirucho dio un paso hacia adelante. "Hay un monstruo en el bosque", tartamudó. "Ha estado aterrorizando al pueblo, atacando animales y hasta personas".

Se hundió el corazón a Anya. El pueblo de Willow Creek era pequeño y pacífico, y la idea de un monstruo le hizo temblar un estremecimiento. Sabía que tenía que hacer algo para proteger a su familia y sus amigos.

Sin dudar un momento, Anya abrió su mochila y recuperó la espada brillante. Con un brillo determinado en los ojos, se dirigió a sus aterrorizados compañeros. "No te preocupes", dijo con la voz firme. "Me ocuparé de ello".

Con Samuel cerca detrás de ella, Anya se aventuró en las sombrías profundidades del bosque. Los árboles parecían susurrar secretos al pasar, y la maleza crujía con criaturas ocultas. A medida que caminaban por el bosque, el aire creció y tembló el suelo bajo sus pies.

De repente, llegaron a un claro y, antes que sus ojos, se encontraba el monstruo, una enorme bestia con dientes afilados, ojos rojos brillantes y garras que podían aplastar a un humano con facilidad. La criatura rugió, con un estruendoso sonido que sacudió el bosque hasta el centro.

El miedo invadió a Anya, pero se negó a dejar que la consumira. Tomó la espada del envoltorio y se cargó hacia el monstruo. La hoja brillaba bajo la luz del sol y, cuando golpeó la piel de la bestia, erupcionó una luz ceguera que envolvió todo bajo su resplandor.

Cuando la luz se atenuó, el monstruo había desaparecido y en su lugar había una pila de cristales destrozados. Anya había derrotado a la criatura con la magia de su mochila, demostrando que incluso los objetos más pequeños podían contener el mayor de los poderes.

Cuando ella y Samuel regresaron al pueblo, los saludaron como héroes. Los habitantes de Willow Creek se alegraron, y la leyenda de Anya, la chica con la mochila mágica, se transmitió de generación en generación. Así, Anya continuó sus aventuras utilizando las maravillas de la mochila para hacer del mundo un lugar mejor, un paso mágico a la vez.

Usar el lenguaje natural

Los modelos grandes de lenguaje son una poderosa herramienta para realizar varias tareas a la vez. A menudo, puedes pedirle a Gemini lo que quieras, y todo saldrá bien.

La API de Gemini no tiene un modo JSON, así que debes tener en cuenta lo siguiente cuando generes estructuras de datos de esta manera:

  • A veces, el análisis falla.
  • El esquema no se puede aplicar de manera forzosa de forma estricta.

Resolverás esos problemas en la próxima sección. Primero, prueba una instrucción simple en lenguaje natural con el esquema escrito como texto. Esto no se optimizó:

model = model = model = genai.GenerativeModel(
    model_name='models/gemini-1.5-pro-latest')

response = model.generate_content(
  textwrap.dedent("""\
    Please return JSON describing the the people, places, things and relationships from this story using the following schema:

    {"people": list[PERSON], "places":list[PLACE], "things":list[THING], "relationships": list[RELATIONSHIP]}

    PERSON = {"name": str, "description": str, "start_place_name": str, "end_place_name": str}
    PLACE = {"name": str, "description": str}
    THING = {"name": str, "description": str, "start_place_name": str, "end_place_name": str}
    RELATIONSHIP = {"person_1_name": str, "person_2_name": str, "relationship": str}

    All fields are required.

    Important: Only return a single piece of valid JSON text.

    Here is the story:

    """) + story,
  generation_config={'response_mime_type':'application/json'}
)
response.text
'{"people": [\n    {\n        "name": "Anya",\n        "description": "A young girl who lives in the town of Willow Creek with her parents, Elise and Edward. She possesses a magical backpack that was handed down to her from her grandmother.",\n        "start_place_name": "Willow Creek",\n        "end_place_name": "Willow Creek"\n    },\n    {\n        "name": "Elise",\n        "description": "Anya\'s kind-hearted mother",\n        "start_place_name": "Willow Creek",\n        "end_place_name": "Willow Creek"\n    },\n    {\n        "name": "Edward",\n        "description": "Anya\'s wise-bearded father",\n        "start_place_name": "Willow Creek",\n        "end_place_name": "Willow Creek"\n    },\n    {\n        "name": "Samuel",\n        "description": "Anya\'s best friend, a curious and adventurous boy with a mischievous grin.",\n        "start_place_name": "Willow Creek",\n        "end_place_name": "Willow Creek"\n    },\n    {\n        "name": "Monster",\n        "description": "A massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease.",\n        "start_place_name": "Forest",\n        "end_place_name": "Forest"\n    }\n], "places": [\n    {\n        "name": "Willow Creek",\n        "description": "A quaint town nestled amidst rolling hills and whispering willows."\n    },\n    {\n        "name": "Forest",\n        "description": "A shadowy place with rustling undergrowth and whispering trees."\n    },\n    {\n        "name": "Schoolhouse",\n        "description": "The only school in the town of Willow Creek."\n    },\n    {\n        "name": "Anya\'s home",\n        "description": "A modest cottage with a creaky wooden door."\n    }\n], "things": [\n    {\n        "name": "Magic backpack",\n        "description": "A magical backpack that was handed down to Anya from her grandmother. Its soft, emerald-green fabric shimmered with an ethereal glow, and its leather straps held secrets that only Anya knew.",\n        "start_place_name": "Anya\'s home",\n        "end_place_name": "Forest"\n    },\n    {\n        "name": "Shimmering sword",\n        "description": "A sword that shimmered in the sunlight and could strike with blinding light.",\n        "start_place_name": "Magic backpack",\n        "end_place_name": "Forest"\n    },\n    {\n        "name": "Book of ancient spells",\n        "description": "A book that contained ancient spells.",\n        "start_place_name": "Magic backpack",\n        "end_place_name": "Forest"\n    },\n    {\n        "name": "Tiny compass",\n        "description": "A compass that always pointed north.",\n        "start_place_name": "Magic backpack",\n        "end_place_name": "Forest"\n    },\n    {\n        "name": "Magical key",\n        "description": "A key that could open any lock.",\n        "start_place_name": "Magic backpack",\n        "end_place_name": "Forest"\n    },\n    {\n        "name": "Shattered crystals",\n        "description": "The remains of the monster after it was defeated by Anya\'s magic backpack.",\n        "start_place_name": "Forest",\n        "end_place_name": "Forest"\n    }\n], "relationships": [\n    {\n        "person_1_name": "Anya",\n        "person_2_name": "Elise",\n        "relationship": "mother-daughter"\n    },\n    {\n        "person_1_name": "Anya",\n        "person_2_name": "Edward",\n        "relationship": "father-daughter"\n    },\n    {\n        "person_1_name": "Anya",\n        "person_2_name": "Samuel",\n        "relationship": "best friends"\n    }\n]}'

Se mostró una cadena JSON. Intenta analizarlo:

import json

print(json.dumps(json.loads(response.text), indent=4))
{
    "people": [
        {
            "name": "Anya",
            "description": "A young girl who lives in the town of Willow Creek with her parents, Elise and Edward. She possesses a magical backpack that was handed down to her from her grandmother.",
            "start_place_name": "Willow Creek",
            "end_place_name": "Willow Creek"
        },
        {
            "name": "Elise",
            "description": "Anya's kind-hearted mother",
            "start_place_name": "Willow Creek",
            "end_place_name": "Willow Creek"
        },
        {
            "name": "Edward",
            "description": "Anya's wise-bearded father",
            "start_place_name": "Willow Creek",
            "end_place_name": "Willow Creek"
        },
        {
            "name": "Samuel",
            "description": "Anya's best friend, a curious and adventurous boy with a mischievous grin.",
            "start_place_name": "Willow Creek",
            "end_place_name": "Willow Creek"
        },
        {
            "name": "Monster",
            "description": "A massive beast with sharp teeth, glowing red eyes, and claws that could crush a human with ease.",
            "start_place_name": "Forest",
            "end_place_name": "Forest"
        }
    ],
    "places": [
        {
            "name": "Willow Creek",
            "description": "A quaint town nestled amidst rolling hills and whispering willows."
        },
        {
            "name": "Forest",
            "description": "A shadowy place with rustling undergrowth and whispering trees."
        },
        {
            "name": "Schoolhouse",
            "description": "The only school in the town of Willow Creek."
        },
        {
            "name": "Anya's home",
            "description": "A modest cottage with a creaky wooden door."
        }
    ],
    "things": [
        {
            "name": "Magic backpack",
            "description": "A magical backpack that was handed down to Anya from her grandmother. Its soft, emerald-green fabric shimmered with an ethereal glow, and its leather straps held secrets that only Anya knew.",
            "start_place_name": "Anya's home",
            "end_place_name": "Forest"
        },
        {
            "name": "Shimmering sword",
            "description": "A sword that shimmered in the sunlight and could strike with blinding light.",
            "start_place_name": "Magic backpack",
            "end_place_name": "Forest"
        },
        {
            "name": "Book of ancient spells",
            "description": "A book that contained ancient spells.",
            "start_place_name": "Magic backpack",
            "end_place_name": "Forest"
        },
        {
            "name": "Tiny compass",
            "description": "A compass that always pointed north.",
            "start_place_name": "Magic backpack",
            "end_place_name": "Forest"
        },
        {
            "name": "Magical key",
            "description": "A key that could open any lock.",
            "start_place_name": "Magic backpack",
            "end_place_name": "Forest"
        },
        {
            "name": "Shattered crystals",
            "description": "The remains of the monster after it was defeated by Anya's magic backpack.",
            "start_place_name": "Forest",
            "end_place_name": "Forest"
        }
    ],
    "relationships": [
        {
            "person_1_name": "Anya",
            "person_2_name": "Elise",
            "relationship": "mother-daughter"
        },
        {
            "person_1_name": "Anya",
            "person_2_name": "Edward",
            "relationship": "father-daughter"
        },
        {
            "person_1_name": "Anya",
            "person_2_name": "Samuel",
            "relationship": "best friends"
        }
    ]
}

Es un proceso relativamente simple y a menudo funciona, pero puedes hacerlo más estricto o sólido si defines el esquema con la función de llamada a función de la API.

Usa llamadas a funciones

Si aún no completaste el instructivo Conceptos básicos de la llamada a funciones, asegúrate de hacerlo primero.

Con la función que llama a tu función y sus parámetros se describen a la API como un genai.protos.FunctionDeclaration. En casos básicos, el SDK puede compilar el FunctionDeclaration a partir de la función y sus anotaciones. Por ahora, tienes que definirlos explícitamente.

Define el esquema

Para comenzar, define person como un objeto con campos de string name, description, start_place_name y end_place_name.

person = genai.protos.Schema(
    type = genai.protos.Type.OBJECT,
    properties = {
        'name':  genai.protos.Schema(type=genai.protos.Type.STRING),
        'description':  genai.protos.Schema(type=genai.protos.Type.STRING),
        'start_place_name': genai.protos.Schema(type=genai.protos.Type.STRING),
        'end_place_name': genai.protos.Schema(type=genai.protos.Type.STRING)
    },
    required=['name', 'description', 'start_place_name', 'end_place_name']
)

Luego, define a personas como un ARRAY de objetos person:

people = genai.protos.Schema(
    type=genai.protos.Type.ARRAY,
    items=person
)

Luego, haz lo mismo con cada una de las entidades que intentes extraer:

place = genai.protos.Schema(
    type = genai.protos.Type.OBJECT,
    properties = {
        'name':  genai.protos.Schema(type=genai.protos.Type.STRING),
        'description':  genai.protos.Schema(type=genai.protos.Type.STRING),
    }
)

places = genai.protos.Schema(
    type=genai.protos.Type.ARRAY,
    items=place
)
thing = genai.protos.Schema(
  type = genai.protos.Type.OBJECT,
  properties = {
      'name':  genai.protos.Schema(type=genai.protos.Type.STRING),
      'description':  genai.protos.Schema(type=genai.protos.Type.STRING),
  }
)

things = genai.protos.Schema(
    type=genai.protos.Type.ARRAY,
    items=thing
)
relationship = genai.protos.Schema(
    type = genai.protos.Type.OBJECT,
    properties = {
        'person_1_name':  genai.protos.Schema(type=genai.protos.Type.STRING),
        'person_2_name':  genai.protos.Schema(type=genai.protos.Type.STRING),
        'relationship':  genai.protos.Schema(type=genai.protos.Type.STRING),
    }
)

relationships = genai.protos.Schema(
    type=genai.protos.Type.ARRAY,
    items=relationship
)

Ahora, compila el FunctionDeclaration:

add_to_database = genai.protos.FunctionDeclaration(
    name="add_to_database",
    description=textwrap.dedent("""\
        Adds entities to the database.
        """),
    parameters=genai.protos.Schema(
        type=genai.protos.Type.OBJECT,
        properties = {
            'people': people,
            'places': places,
            'things': things,
            'relationships': relationships
        }
    )
)

Llama a la API

Como viste en Conceptos básicos de la llamada a funciones, ahora puedes pasar este FunctionDeclaration al argumento tools del constructor genai.GenerativeModel (el constructor también aceptaría una representación JSON equivalente de la declaración de la función):

model = model = genai.GenerativeModel(
    model_name='models/gemini-1.5-pro-latest',
    tools = [add_to_database])

Cada vez que llames a la API, el SDK enviará las herramientas junto con tu instrucción, y el modelo debe llamar a la función que definiste:

result = model.generate_content(f"""
Please add the people, places, things, and relationships from this story to the database:

{story}
""",
# Force a function call
tool_config={'function_calling_config':'ANY'})

Ahora no hay texto para analizar. El resultado es una estructura de datos.

'text' in result.candidates[0].content.parts[0]
False
'function_call' in result.candidates[0].content.parts[0]
True
fc = result.candidates[0].content.parts[0].function_call
print(type(fc))
<class 'google.ai.generativelanguage_v1beta.types.content.FunctionCall'>

La clase genai.protos.FunctionCall se basa en los búferes de protocolo de Google. Conviértela en un objeto compatible con JSON más conocido:

print(json.dumps(type(fc).to_dict(fc), indent=4))
{
    "name": "add_to_database",
    "args": {
        "things": [
            {
                "name": "Magical Backpack",
                "description": "Anya's prized possession, the Magical Backpack, is no ordinary satchel. Its soft, emerald-green fabric shimmers with an ethereal glow, and its leather straps have secrets that only Anya knows. Within its capacious interior lay an enchanted world, filled with wonders that would ignite her imagination and change her life forever."
            },
            {
                "name": "Shimmering Sword",
                "description": "Among the wonders in Anya's Magical Backpack, lies a shimmering sword. With a determined gleam in her eye, she retrieved the shimmering sword and charged towards the monster."
            },
            {
                "description": "Residing within the Magical Backpack, the Book of Ancient Spells holds secrets untold.",
                "name": "Book of Ancient Spells"
            },
            {
                "description": "Tucked away in the Magical Backpack is a tiny compass that always points north.",
                "name": "Tiny Compass that Always Points North"
            },
            {
                "description": "Hidden within the Magical Backpack is a magical key that can open any lock.",
                "name": "Magical Key that Can Open Any Lock"
            }
        ],
        "relationships": [
            {
                "relationship": "Mother-Daughter",
                "person_1_name": "Anya",
                "person_2_name": "Elise"
            },
            {
                "person_2_name": "Edward",
                "relationship": "Father-Daughter",
                "person_1_name": "Anya"
            },
            {
                "person_2_name": "Samuel",
                "person_1_name": "Anya",
                "relationship": "Best Friends"
            }
        ],
        "people": [
            {
                "name": "Anya",
                "description": "Anya, the main character of the story, is a young girl with a magical backpack.",
                "start_place_name": "Willow Creek",
                "end_place_name": "Unknown"
            },
            {
                "name": "Elise",
                "description": "Anya's mother, Elise is a kind-hearted woman.",
                "end_place_name": "Unknown",
                "start_place_name": "Willow Creek"
            },
            {
                "start_place_name": "Willow Creek",
                "end_place_name": "Unknown",
                "name": "Edward",
                "description": "Anya's father, Edward is a wise-bearded man."
            },
            {
                "end_place_name": "Unknown",
                "start_place_name": "Willow Creek",
                "description": "Anya's best friend, Samuel is a curious and adventurous boy with a mischievous grin.",
                "name": "Samuel"
            }
        ],
        "places": [
            {
                "description": "The quaint town of Willow Creek is nestled amidst rolling hills and whispering willows.",
                "name": "Willow Creek"
            },
            {
                "description": "The town's only schoolhouse.",
                "name": "Schoolhouse"
            },
            {
                "description": "A shadowy place filled with secrets and dangers, the Forest is home to a terrifying monster.",
                "name": "Forest"
            }
        ]
    }
}

Conclusión

Si bien la API puede manejar problemas de extracción de datos estructurados con entrada y salida de texto pura, es probable que usar la llamada a función sea más confiable, ya que te permite definir un esquema estricto y elimina un paso de análisis potencialmente propenso a errores.