Orquestrando processamento de documentos em larga escala com AWS Step Functions e inferência em lote do Amazon Bedrock

Por Brian Zambrano, Senior Deep Learning Architect e Dan Ford, Senior Applied Scientist na Amazon Web Services.

As organizações frequentemente possuem grandes volumes de documentos contendo informações valiosas que permanecem bloqueadas e não pesquisáveis. Esta solução aborda a necessidade de um pipeline escalável e automatizado de extração de texto e base de conhecimento que transforma coleções de documentos estáticos em repositórios inteligentes e pesquisáveis para aplicações de IA generativa.

As organizações podem automatizar a extração de conteúdo e metadados estruturados para construir bases de conhecimento abrangentes que alimentam soluções de geração aumentada por recuperação (RAG), reduzindo significativamente os custos de processamento manual e o tempo de valorização. A arquitetura não apenas demonstra o processamento automático de 500 artigos de pesquisa, mas também escala para lidar com volumes de documentos empresariais de forma econômica através do modelo de preços de inferência em lote do Amazon Bedrock.

Visão geral

A inferência em lote do Amazon Bedrock é um recurso do Amazon Bedrock que oferece um desconto de 50% nas solicitações de inferência. Embora o Amazon Bedrock agende e execute o trabalho em lote (necessitando de um mínimo de 100 solicitações de inferência) conforme a capacidade fica disponível, a inferência não será em tempo real. Para casos de uso onde você pode acomodar minutos a horas de latência, a inferência em lote do Amazon Bedrock é uma boa opção.

Esta publicação demonstra como construir um pipeline automatizado e Serverless usando AWS Step Functions, Amazon Textract, inferência em lote do Amazon Bedrock e Amazon Bedrock Knowledge Bases para extrair texto, criar metadados e carregá-los em uma base de conhecimento em escala. A solução de exemplo processa 500 artigos de pesquisa em formato PDF da Amazon Science, extrai texto usando Amazon Textract, gera metadados estruturados com inferência em lote do Amazon Bedrock e o modelo Amazon Nova Pro, e carrega a saída final, incluindo filtro do Amazon Bedrock Knowledge Base, em uma Amazon Bedrock Knowledge Base.

Arquitetura

Esta solução usa Step Functions com processamento paralelo de trabalhos do Amazon Textract através de fluxos de trabalho filhos executados por Distributed Map. Você pode usar os controles de simultaneidade oferecidos pelo Distributed Map para processar documentos o mais rápido possível dentro de suas cotas do Amazon Textract. Aumentar a velocidade de processamento requer ajustar sua cota do Amazon Textract e atualizar a configuração do Distributed Map. A inferência em lote do Amazon Bedrock lida com simultaneidade, escalabilidade e limitação. Isso significa que você pode criar o trabalho sem gerenciar essas complexidades.

Nesta implementação de exemplo, a solução processa artigos de pesquisa para extrair metadados como:

Disponibilidade de código e localizações de repositórios
Disponibilidade de conjuntos de dados e métodos de acesso
Tipos de metodologia de pesquisa
Indicadores de reprodutibilidade
Outros atributos de pesquisa relevantes

As partes de alto nível desta solução incluem:

Extrair texto de documentos PDF com Amazon Textract em paralelo, através do Distributed Map do Step Functions.
Analisar texto extraído usando inferência em lote do Amazon Bedrock para extrair metadados estruturados.
Carregar texto extraído e metadados em uma base de conhecimento pesquisável usando Amazon Bedrock Knowledge Bases com Amazon OpenSearch Serverless.

Figura 1. Diagrama completo da arquitetura

Pré-requisitos

Os seguintes pré-requisitos são necessários para concluir esta solução:

Acesso a uma conta AWS através do AWS Management Console e da Interface de Linha de Comando da AWS (AWS CLI). O usuário do AWS Identity and Access Management (IAM) que você usa deve ter permissões para fazer as chamadas de serviço AWS necessárias e gerenciar recursos AWS mencionados neste post. Ao fornecer permissões ao usuário IAM, siga o princípio do menor privilégio.
AWS CLI instalada e configurada. Se você estiver usando credenciais de longo prazo, como chaves de acesso, siga gerenciar chaves de acesso para usuários IAM e proteger chaves de acesso para melhores práticas.
Git instalado.
Python 3.13+ instalado.
Node e npm instalados.
AWS Cloud Development Kit (AWS CDK) instalado.

Executando a solução

A solução completa usa AWS CDK para implementar duas pilhas do AWS CloudFormation:

BedrockKnowledgeBaseStack: Cria a infraestrutura da base de conhecimento
SFNBatchInferenceStack: Implementa o fluxo de trabalho de processamento principal

Primeiro, clone o repositório GitHub em seu ambiente de desenvolvimento local e instale os requisitos:

git clone https://github.com/aws-samples/sample-step-functions-batch-inference.git .

cd sample-step-functions-batch-inference

npm install

Em seguida, implante a solução usando AWS CDK:

cdk deploy --all

Após implantar as pilhas do cdk, carregue suas fontes de dados (arquivos PDF) no bucket de entrada do Amazon S3 criado pelo AWS CDK. Neste exemplo, carreguei 500 artigos da Amazon Science. O nome do bucket de entrada está incluído nas saídas do AWS CDK:

Saídas:

SFNBatchInference.BatchInputBucketName = sfnbatchinference-batchinputbucket11aaa222-nrjki8tewwww

Extração de texto paralela

O processo começa quando você carrega um arquivo manifest.json no bucket de entrada. O arquivo de manifesto lista os arquivos para processamento, que já existem no bucket de entrada. Os nomes de arquivo listados em manifest.json definem o que constitui uma única execução de trabalho de processamento. Para criar outra execução, você criaria um manifest.json diferente e o carregaria no mesmo bucket S3.

[
  {
    "filename": "flexecontrol-flexible-and-efficient-multimodal-control-for-text-to-image-generation.pdf"
  },
  {
    "filename": "adaptive-global-local-context-fusion-for-multi-turn-spoken-language-understanding.pdf"
  }
]

A definição do AWS CDK para o bucket de entrada inclui notificações do Amazon EventBridge e cria uma regra que aciona o fluxo de trabalho do Step Functions sempre que um arquivo manifest.json é carregado.

private createS3Buckets() {
    const batchBucket = new s3.Bucket(this, "BatchInputBucket", {
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    })
    batchBucket.enableEventBridgeNotification()

    new cdk.CfnOutput(this, "BatchInputBucketName", {
      value: batchBucket.bucketName,
      description: "Name of input bucket to send PDF documents that Textract will read.",
    })

    const manifestFileCreatedRule = new eventBridge.Rule(this, "ManifestFileCreatedRule", {
      eventPattern: {
        source: ["aws.s3"],
        detailType: ["Object Created"],
        detail: {
          bucket: {
            name: [batchBucket.bucketName],
          },
          object: {
            key: ["manifest.json"],
          },
        },
      },
    })

    return { batchBucket, manifestFileCreatedRule }
  }

O primeiro passo no fluxo de trabalho do Step Functions é uma execução do Distributed Map que realiza as seguintes ações para cada PDF no arquivo de manifesto:

Inicia um trabalho do Amazon Textract, fornecendo um tópico do Amazon Simple Notification Service (Amazon SNS) para notificação de conclusão.
Grava o token de tarefa do Step Functions no Amazon DynamoDB, pausando o fluxo de trabalho filho individual.
Processa a mensagem do Amazon SNS quando o trabalho do Amazon Textract é concluído, acionando uma função do AWS Lambda.
Usa uma função Lambda para recuperar o token de tarefa do DynamoDB usando o JobId do Amazon Textract.
Busca os resultados brutos do Amazon Textract, organiza o texto para legibilidade e grava os resultados em um bucket S3

Um componente chave desta arquitetura é o padrão de callback que o Amazon Textract suporta usando a opção NotificationChannel, conforme mostrado na figura anterior. A definição do AWS CDK do estado do Step Functions que inicia o trabalho do Amazon Textract é mostrada a seguir.

const startTextractStep = new tasks.CallAwsService(this, "StartTextractJob", {
  service: "textract",
  action: "startDocumentAnalysis",
  resultPath: "$.textractOutput",
  parameters: {
    DocumentLocation: {
      S3Object: {
        Bucket: sourceBucket.bucketName,
        Name: sfn.JsonPath.stringAt("$.filename"),
      },
    },
    FeatureTypes: ["LAYOUT"],
    NotificationChannel: {
      RoleArn: textractRoleArn,
      SnsTopicArn: snsTopicArn,
    },
  },
  iamResources: ["*"],
})

A função Lambda que lida com tokens de tarefa extrai o JobId do Amazon Textract da mensagem do Amazon SNS, busca o TaskToken do DynamoDB e retoma o fluxo de trabalho do Step Functions enviando o TaskToken:

from aws_lambda_powertools.utilities.data_classes import SNSEvent, event_source

@event_source(data_class=SNSEvent)
def handle_textract_task_complete(event, context):
    # Multiple records can be delivered in a single event
    for record in event.records:
        sns_message = json.loads(record.sns.message)
        textract_job_id = sns_message["JobId"]

        # Get both task token and original file from DynamoDB
        ddb_item = _get_item_from_ddb(textract_job_id)

        # Send both the job ID and original file name in the response
        _send_task_success(
            ddb_item["TaskToken"],
            {
                "TextractJobId": textract_job_id,
                "OriginalFile": ddb_item["OriginalFile"],
            },
        )
        # Delete the task token from DynamoDB after use
        _delete_item_from_ddb(textract_job_id)

def _send_task_success(task_token: str, output: None | dict = None) -> None:
    """Sends task success to Step Functions with the provided output"""
    sfn = boto3.client("stepfunctions")
    sfn.send_task_success(taskToken=task_token, output=json.dumps(output or {}))

O Distributed Map executa até 10 fluxos de trabalho filhos simultaneamente, controlado pela configuração maxConcurrency. Embora o Step Functions suporte a execução de até 10.000 execuções de fluxo de trabalho filho, a simultaneidade prática para esta solução é restrita pelas cotas do Amazon Textract. A API startDocumentAnalysis tem uma cota padrão de 10 solicitações por segundo (RPS), o que significa que você deve considerar esse limite ao escalar suas cargas de trabalho de processamento de documentos e potencialmente solicitar aumentos de cota para requisitos de throughput mais altos.

const distributedMap = new sfn.DistributedMap(this, "DistributedMap", {
  mapExecutionType: sfn.StateMachineType.STANDARD,
  maxConcurrency: 10,
  itemReader: new sfn.S3JsonItemReader({
    bucket: sourceBucket,
    key: "manifest.json",
  }),
  resultPath: "$.files",
}

Executando inferência em lote do Amazon Bedrock

Quando todos os trabalhos do Amazon Textract terminam, o estado Distributed Map cria um arquivo de entrada de inferência em lote do Amazon Bedrock, inicia o trabalho de inferência do Amazon Bedrock e aguarda sua conclusão.

Uma função Lambda coleta resultados de texto do Amazon S3 e cria um arquivo de entrada de inferência em lote do Amazon Bedrock com prompts personalizados.
O fluxo de trabalho inicia o trabalho de inferência em lote do Amazon Bedrock chamando createModelInvocationJob e enviando o arquivo de entrada de inferência em lote como entrada.
O fluxo de trabalho pausa e armazena o token de tarefa no DynamoDB.
Uma regra do EventBridge corresponde a eventos de inferência em lote do Amazon Bedrock concluídos e, após a conclusão do trabalho, aciona uma função Lambda. A função Lambda recupera o token de tarefa e retoma o fluxo de trabalho, conforme mostrado na figura a seguir.

Uma entrada de inferência em lote é um único arquivo jsonl com múltiplas entradas como o exemplo a seguir. O prompt em cada solicitação de inferência instrui o modelo de linguagem grande (LLM) a analisar o artigo e extrair metadados. Leia o template de prompt completo no repositório GitHub.

{
  "recordId": "c1b8a3b2086141f963",
  "modelInput": {
    "messages": [
      {
        "role": "user",
        "content": [
          {
            "text": "Analyze the following research paper transcript and extract metadata about code and dataset availability. Extract the following metadata from this research paper transcript:\n\n1. **has_code**: Does the paper mention or link to source code? (true/false) ...... Return only valid JSON matching the schema above. Do not include any text outside of the JSON structure."
          }
        ]
      }
    ],
    "inferenceConfig": { "maxTokens": 4096 }
  }
}

Populando a Amazon Bedrock Knowledge Base

Após a conclusão da inferência em lote, o fluxo de trabalho faz o seguinte:

Extrai resultados de inferência e cria arquivos de metadados com base nos resultados de inferência do Amazon Bedrock (exemplo de metadados mostrado na figura a seguir).
Inicia um trabalho de ingestão da Amazon Bedrock Knowledge Base.
Monitora o status do trabalho de ingestão usando estados Wait e Choice do Step Functions.
Envia uma notificação de conclusão através do Amazon SNS.

O seguinte mostra o formato de metadados de exemplo:

{
  "metadataAttributes": {
    "has_code": true,
    "has_dataset": false,
    "code_availability": "publicly_available",
    "dataset_availability": "not_available",
    "research_type": "methodology",
    "is_reproducible": true,
    "code_repository_url": "https://github.com/amazon-science/PIXELS"
  }
}

Testando a base de conhecimento

Após a conclusão bem-sucedida do fluxo de trabalho, você pode testar a base de conhecimento para verificar se os documentos e metadados foram devidamente ingeridos e são pesquisáveis. Existem dois métodos práticos para testar uma Amazon Bedrock Knowledge Base:

Usando o Console
Usando o AWS SDK para executar uma consulta

Testando através do Console

O Console fornece uma interface intuitiva para testar suas consultas de base de conhecimento com filtros de metadados:

Navegue até o console do Amazon Bedrock.
No painel de navegação à esquerda, escolha Knowledge Bases na seção Build.
Escolha a base de conhecimento criada pela implantação do AWS CDK (o nome será exibido na saída da pilha do AWS CDK).
Escolha o botão Test no canto superior direito.
Na interface de teste, escolha seu modelo de fundação (FM) preferido (como Amazon Nova Pro).
Expanda a coluna Configurations e navegue até a seção Filters.
Configure filtros com base nos metadados extraídos, conforme mostrado na figura a seguir.

Digite uma consulta em linguagem natural relacionada aos seus documentos, por exemplo: “Pesquisas recentes sobre geração aumentada por recuperação?”

O console exibe a resposta gerada junto com atribuições de origem mostrando quais documentos foram recuperados e usados para formular a resposta, filtrados pelos atributos de metadados especificados, conforme mostrado na figura a seguir.

Testando via API

Para testes programáticos e integração em aplicações, use o AWS SDK com filtragem de metadados. O seguinte é um exemplo em Python usando boto3:

model_arn = "arn:aws:bedrock:us-east-1::foundation-model/amazon.nova-pro-v1:0"

# Query for papers with publicly available code
response = bedrock_agent_runtime.retrieve_and_generate(
    input={'text': "What recent research has been done on RAG?"},
    retrieveAndGenerateConfiguration={
        'type': 'KNOWLEDGE_BASE',
        'knowledgeBaseConfiguration': {
            'knowledgeBaseId': knowledge_base_id,
            'modelArn': model_arn,
            'retrievalConfiguration': {
                'vectorSearchConfiguration': {
                    'numberOfResults': 5,
                    'filter': {"equals": {"key": "has_code", "value": True}},
                }
            },
        },
    },
)

# Display results
print(f"Response: {response['output']['text']}\n")
print("Source Documents:")

for citation in response.get('citations', []):
    for reference in citation.get('retrievedReferences', []):
        metadata = reference.get('metadata', {})
        print(f" Document: {reference['location']['s3Location']['uri']}\n")

O seguinte é a saída do script de teste:

Response: Recent research on Retrieval-Augmented Generation (RAG) has focused on enhancing the system's ability to dynamically retrieve and utilize relevant information from a Vector Database (VDB) to improve decision-making and performance. Key innovations include:

1. **Dynamic Retrieval and Utilization**: The system is designed to query the VDB for contextually relevant past experiences, which significantly improves decision quality and accelerates performance by leveraging a growing repository of relevant experiences.

2. **Teacher-Student Instructional Tuning**: A novel mechanism where a Teacher agent refines a Student agent's core policy through direct interaction. The Teacher generates a modified SYSTEM prompt based on the Student's actions, creating a meta-learning loop that enhances the Student's reasoning policy over time.

Conclusão

Esta solução demonstra como combinar múltiplos serviços de IA e Serverless da AWS para construir um pipeline escalável de processamento de documentos. As organizações podem usar AWS Step Functions para orquestração, Amazon Textract para processamento de documentos, inferência em lote do Amazon Bedrock para análise inteligente de conteúdo e Amazon Bedrock Knowledge Bases para armazenamento pesquisável. Por sua vez, elas podem automatizar a extração de insights de grandes coleções de documentos enquanto otimizam custos.

Seguindo esta solução, você pode construir uma base sólida para pipelines de processamento de documentos em escala de produção que mantêm a flexibilidade para se adaptar aos seus requisitos específicos, garantindo confiabilidade, escalabilidade e excelência operacional. Siga este link para saber mais sobre arquiteturas Serverless.

Este conteúdo foi traduzido do post original do blog, que pode ser encontrado aqui.

Autores

	Brian Zambrano é Senior Deep Learning Architect na Amazon Web Services.
	Dan Ford é Senior Applied Scientist na Amazon Web Services.

Tradutores

Daniel Abib é Arquiteto de Soluções Sênior e Especialista em Amazon Bedrock na AWS, com mais de 25 anos trabalhando com gerenciamento de projetos, arquiteturas de soluções escaláveis, desenvolvimento de sistemas e CI/CD, microsserviços, arquitetura Serverless & Containers e especialização em Machine Learning. Ele trabalha apoiando Startups, ajudando-os em sua jornada para a nuvem.
https://www.linkedin.com/in/danielabib/

Nicolas Tarzia é Senior Technical Account Manager na AWS, com mais de 13 anos de experiência, com ampla experiência em arquitetura cloud, engenharia e design de software. Atualmente está habilitando empresas do ramo de ISV (Independent Software Vendors) simplificando a operação na nuvem e otimizando os custos em cloud. Sua área de interesse são tecnologias serverless.
https://www.linkedin.com/in/nicolastarzia

O blog da AWS