[Kendra] 文字情報がないPDFを検索できるようにする

#Amazon Kendra

#Amazon Bedrock

中村修太

2024.08.26

Introduction

Amazon KendraはAWSが提供する機械学習と自然言語処理をした、フルマネージドな検索サービスです。
Kendraでは従来のものとは違い、文脈に応じた関連性の高い結果を提供します。
さまざまなDataSourceに対応し、PDFやWordなどいろいろな形式のドキュメントを解析可能になっています。

先日たまたま見たPDFが紙をスキャンして作成されたらしく、
テキストが画像化されており、そのままではKendraで検索できませんでした。
なんとかできないものかと試したらとりあえず検索できるようになったので、そのログです。
他にもっとスマートな方法があると思いますが、とりあえず今回試した方法を紹介します。

Setup

AWS CLI : 2.15.56
Node : v22.5.1

Try

今回はここにあるPDFを全文検索できるようにしてみます。
稀代の名車、Z750-RSのパーツカタログです。
このPDFには文字データがなく、そのままではKendraに登録できません。
AcrobatとかのOCR機能を使えば文字認識させることができるみたいですが、
そんなものはないので、PDFをjpgに変換してLLMのOCR機能で画像から文字を抽出し、
メタデータファイルを作成して検索可能にしてみます。

PDFを変換&メタデータファイル作成

こことかを使ってもいいし、
pdf2jpgとかのライブラリで変換してもよいです。
とりあえずPDFを各ページごとにjpgに変換します。

変換したjpgは「hoge_page-0001.jpg」みたいな感じでページごとに作成しました。
そして、jpg郡と元PDFは、適当なS3バケットにアップロードしておきます。
PDFはBucket直下、jpgはimagesディレクトリを作成してそこにアップロードしました。

アップロードしたらBedrockをつかって各jpgから文字情報を抽出してメタデータを作成します。
contentフィールドに抽出したテキストをすべて突っ込むという力技を使用。
なので、ある程度サイズが大きくなると制限にひっかかるので注意。
下記プログラムを実行すると、メタデータのjsonファイルがS3にできてました。

//create_metadata.js
const { S3Client, ListObjectsV2Command, GetObjectCommand, PutObjectCommand } = require("@aws-sdk/client-s3");
const { BedrockRuntimeClient, InvokeModelCommand } = require("@aws-sdk/client-bedrock-runtime");
const sharp = require('sharp');
const path = require('path');

// AWS設定
const region = "<Your Region>";
const s3Client = new S3Client({ region });
const bedrockClient = new BedrockRuntimeClient({ region });

// S3バケットとパスの設定
const bucketName = "<Your S3 Bucket>";
const directoryPath = "images/";
// 原本のPDF名を設定
const originalPdfName = "YourPDF.pdf";

async function processImages() {
    try {
        console.log("処理を開始...");

        // S3バケット内のオブジェクトリストを取得
        const listCommand = new ListObjectsV2Command({
            Bucket: bucketName,
            Prefix: directoryPath
        });
        const { Contents } = await s3Client.send(listCommand);
        
        if (!Contents || Contents.length === 0) {
            console.error("File not found");
            return;
        }

        // 画像ファイルをソートして処理
        const imageFiles = Contents
            .filter(file => file.Key.endsWith('.jpg') || file.Key.endsWith('.png'))
            .sort((a, b) => {
                const pageNumA = parseInt(a.Key.match(/(\d+)\.(jpg|png)$/)[1]);
                const pageNumB = parseInt(b.Key.match(/(\d+)\.(jpg|png)$/)[1]);
                return pageNumA - pageNumB;
            });

        const pdfMetadata = {
            _source_uri: `s3://${bucketName}/${directoryPath}${originalPdfName}`,
            _language_code: "ja",
            title: path.basename(originalPdfName, '.pdf'),
            file_type: "PDF",
            page_count: imageFiles.length,
            last_modified: new Date().toISOString(),
            content: "",
        };

        // 各画像ファイルを処理
        for (let i = 0; i < imageFiles.length; i++) {
            const imageFile = imageFiles[i];
            const pageNumber = i + 1;
            console.log(`ページ ${pageNumber} を処理中: ${imageFile.Key}`);
            const extractedText = await extractTextFromImage(imageFile.Key);
            console.log(`ページ ${pageNumber} のテキスト抽出が完了しました。`);

            pdfMetadata[`page_${pageNumber}_content`] = extractedText;
            pdfMetadata.content += extractedText + "\n\n";
        }

        // メタデータファイルをS3にアップロード
        const metadataFileName = `${originalPdfName}.metadata.json`;
        await uploadMetadataToS3(metadataFileName, JSON.stringify(pdfMetadata, null, 2));
        
        console.log("すべての処理が完了しました。");

    } catch (error) {
        console.error("エラーが発生しました:", error);
    }
}

async function extractTextFromImage(imageKey) {
    console.log(`    画像 ${imageKey} からテキストを抽出中...`);
    // S3から画像を取得
    const getCommand = new GetObjectCommand({
        Bucket: bucketName,
        Key: imageKey
    });
    const { Body } = await s3Client.send(getCommand);
    const imageBuffer = await streamToBuffer(Body);

    // 画像をBase64エンコード
    const base64Image = await sharp(imageBuffer).toBuffer().then(buffer => buffer.toString('base64'));

    // Bedrock APIを呼び出してテキスト抽出
    console.log(`    Bedrock APIを呼び出してテキスト抽出中...`);
    const params = {
        modelId: "<Model IDを指定。今回はclaude3.5 sonnetを使用>",
        contentType: "application/json",
        accept: "application/json",
        body: JSON.stringify({
            anthropic_version: "bedrock-2023-05-31",
            max_tokens: 1000,
            messages: [
                {
                    role: "user",
                    content: [
                        {
                            type: "image",
                            source: {
                                type: "base64",
                                media_type: "image/jpeg",
                                data: base64Image
                            }
                        },
                        {
                            type: "text",
                            text: "この画像に含まれる日本語や英語のテキストを全て抽出し、読み取った順番通りに列挙してください。レイアウトや改行は無視して、テキストのみを抽出してください。"
                        }
                    ]
                }
            ]
        })
    };

    const command = new InvokeModelCommand(params);
    const response = await bedrockClient.send(command);
    const result = JSON.parse(new TextDecoder().decode(response.body));

    console.log(`    テキスト抽出が完了しました。`);
    // 抽出されたテキストを返す
    return result.content[0].text;
}

async function uploadMetadataToS3(fileName, content) {
    const uploadParams = {
        Bucket: bucketName,
        Key: `${directoryPath}${fileName}`,
        Body: content,
        ContentType: "application/json"
    };

    const command = new PutObjectCommand(uploadParams);
    await s3Client.send(command);
    console.log(`メタデータファイルがアップロードされました: ${fileName}`);
}

function streamToBuffer(stream) {
    return new Promise((resolve, reject) => {
        const chunks = [];
        stream.on('data', (chunk) => chunks.push(chunk));
        stream.on('error', reject);
        stream.on('end', () => resolve(Buffer.concat(chunks)));
    });
}

console.log("スクリプトを開始します...");
processImages().then(() => console.log("スクリプトが完了しました。"));

Kendraの準備

AWSコンソールでもいいし↓みたいなプログラムでもいいので、Kendra Indexを作成します。
IAM Roleとか必要なので、コンソールのほうが楽かもしれないです。

const { KendraClient, CreateIndexCommand } = require("@aws-sdk/client-kendra");

const region = "＜your region＞";
const kendra = new KendraClient({ region });

async function createKendraIndex() {
  const params = {
    Name: 'Z2-Manual-Index',
    Edition: 'DEVELOPER_EDITION',
    RoleArn: '<Kendra作成用Role>'
  };

  try {
    const command = new CreateIndexCommand(params);
    const result = await kendra.send(command);
    console.log('Index created:', result.Id);
    return result.Id;
  } catch (error) {
    console.error('Error creating index:', error);
  }
}

createKendraIndex();

↑のプログラムを実行したら、ステータスがactiveになるまで待ちます。
下記コマンドでステータスがわかります。

% aws kendra describe-index --id <Kendra Index ID> --region <Your Region> --query "Status"

Indexができたら次はDataSourceの作成です。
S3において自動でsyncできれば楽だったのですが、文字情報がないPDFはsyncできませんでした。
なので、カスタムデータソースを使います。

//create_datasource.js
const { KendraClient, CreateDataSourceCommand } = require("@aws-sdk/client-kendra");

const region = "＜Your Region＞";
const kendra = new KendraClient({ region });
const kendra_name = 'Z2-Manual-Index';
const kendra_index = '<Kendra Index ID>';

async function createDataSource(indexId) {
    const params = {
        IndexId: indexId,
        Name: 'Z2-Manual-DataSource',
        Type: 'CUSTOM',
        DataSourceConfiguration: {
            CustomDataSourceConfiguration: {
                ConnectionConfiguration: {
                    //実際のdata sourceの接続情報に応じて変更
                    AuthenticationType: "NO_AUTH",
                },
            }
        },
        LanguageCode: 'ja',
        Description: 'Custom data source for Z2 Manual'
    };

    try {
        const command = new CreateDataSourceCommand(params);
        const result = await kendra.send(command);
        console.log('Data source created:', result.Id);
        return result.Id;
    } catch (error) {
        console.error('Error creating data source:', error);
        if (error.message) {
            console.error('Error message:', error.message);
        }
        if (error.$metadata) {
            console.error('Error metadata:', error.$metadata);
        }
    }
}

createDataSource(kendra_index);

実行するとDataSourceが作成されます。

% node create_datasource.js
Data source created: <Data Source ID>

次にドキュメント（PDF）を追加します。
プログラム実行前にfile_type、page_count、last_modifiedのフィールドを
Kendraに作成しておきましょう。（とりあえず手動で作成した）
その後、下記プログラム（かなり簡易版）を実行するとドキュメント追加と同期が実行されます。

//sync.js
const { KendraClient, BatchPutDocumentCommand, StartDataSourceSyncJobCommand } = require("@aws-sdk/client-kendra");
const { S3Client, GetObjectCommand } = require("@aws-sdk/client-s3");
const fs = require('fs').promises;

const region = "ap-northeast-1";
const kendra = new KendraClient({ region });
const s3 = new S3Client({ region });

const kendra_index = '<Kendra Index ID>';
const dataSourceId = '<DataSource ID>';
const bucketName = '<S3 Bucket Name>';
const pdfKey = 'Z2.pdf';
const metadataKey = 'Z2.pdf.metadata.json';

async function getS3Object(bucket, key) {
    const command = new GetObjectCommand({ Bucket: bucket, Key: key });
    const response = await s3.send(command);
    return response.Body;
}

async function streamToBuffer(stream) {
    const chunks = [];
    for await (const chunk of stream) {
        chunks.push(chunk);
    }
    return Buffer.concat(chunks);
}

async function addDocumentToIndex(indexId) {
    try {
        const metadataStream = await getS3Object(bucketName, metadataKey);
        const metadataBuffer = await streamToBuffer(metadataStream);
        const metadata = JSON.parse(metadataBuffer.toString());

        const params = {
            IndexId: indexId,
            Documents: [
                {
                    Id: 'Z2-Manual',
                    Title: metadata.title,
                    ContentType: 'PLAIN_TEXT',
                    Blob: Buffer.from(metadata.contents, 'utf-8'),
                    Attributes: [
                        {
                            Key: '_language_code',
                            Value: { StringValue: metadata._language_code }
                        },
                        {
                            Key: 'file_type',
                            Value: { StringValue: metadata.file_type }
                        },
                        {
                            Key: 'page_count',
                            Value: { LongValue: metadata.page_count }
                        },
                        {
                            Key: 'last_modified',
                            Value: { DateValue: new Date(metadata.last_modified) }
                        },
                        {
                            Key: '_source_uri',
                            Value: { StringValue: metadata._source_uri }
                        }
                    ]
                }
            ]
        };

        const command = new BatchPutDocumentCommand(params);
        const result = await kendra.send(command);
        console.log('Document added:', result);

        if (result.FailedDocuments && result.FailedDocuments.length > 0) {
            console.error('Some documents failed to be added:', result.FailedDocuments);
            return false;
        }

        return true;
    } catch (error) {
        console.error('Error adding document:', error);
        return false;
    }
}

async function syncDataSource(indexId, dataSourceId) {
    const params = {
        IndexId: indexId,
        Id: dataSourceId
    };

    try {
        const command = new StartDataSourceSyncJobCommand(params);
        const result = await kendra.send(command);
        console.log('Sync job started:', result.ExecutionId);
    } catch (error) {
        console.error('Error starting sync job:', error);
    }
}

async function main() {
    const documentAdded = await addDocumentToIndex(kendra_index);
    if (documentAdded) {
        console.log('Document Add done.');
    } else {
        console.log('Skipping sync due to document addition failure');
    }
}

main();

上記プログラムを実行すると、ドキュメント追加とsyncが実行されます。

% node sync.js
Document added: {
  '$metadata': {
    httpStatusCode: 200,
    requestId: 'xxxxxxxxxxx',
    extendedRequestId: undefined,
    cfId: undefined,
    attempts: 1,
    totalRetryDelay: 0
  },
  FailedDocuments: []
}
Sync job started: xxxxxxxxxx

なお、syncを中断したいときは下記コマンドを実行しましょう。

% aws kendra stop-data-source-sync-job --index-id <Kendra Index ID> --id <DataSource ID>

進捗を確認したいときは下記コマンドを実行します。

% aws kendra list-data-source-sync-jobs \
    --index-id <Kendra Index ID> \
    --id <DataSource ID>

{
    "History": [
        {
            "ExecutionId": "xxxxxxx",
            "StartTime": "2024-08-21T10:23:26.774000+09:00",
            "Status": "SYNCING",
            "Metrics": {
                ・・・・・・
            }
        }
    ]
}