usomaru技術ぶろぐ

学んだことをちょこちょこメモ( ..)φ

【雑談】ChatGPTでPDF分割とCognitive Search にインデックス作成までのスクリプトをつくった

半分雑談です。
仕事でAzure Cognitive SearchとAzure Open AIを使ってPDFの資料を検索できないか調査をしているのですが、こちらのGitHubを参考にしていました。

github.com

ですが、スクラッチの開始の手順通りに進めてもうまくできず…
会社のサブスクリプションを使っている関係か、権限のエラーなどでできませんでした。
リージョンをEAST USにすればリソースの作成まではできました。  (https://zenn.dev/apeiria_onda/articles/a9b516c501aae1)

今回PDFの分割・インデックスが作成されるまでの実験がしたかったため、Azureのリソースをあらかじめ作成し、この中にあるprepdocs.pyというファイルを直接いじってこのファイル単体で実行してみることにしました。

ここで問題が。
私はpythonでプログラムをかいたことがほとんどありません。
なんとなくノリでわかるようなわからないような、、、
というレベルです。

そこで活躍したのがChatGPTでした。

まず、メソッドごとにコードをコピーし、ChatGPTに貼り付けて、このコードは何をしているのかを聞きます。

def get_document_text(filename):
    offset = 0
    page_map = []
    if args.localpdfparser:
        reader = PdfReader(filename)
        pages = reader.pages
        for page_num, p in enumerate(pages):
            page_text = p.extract_text()
            page_map.append((page_num, offset, page_text))
            offset += len(page_text)
    else:
        if args.verbose: print(f"Extracting text from '{filename}' using Azure Form Recognizer")
        form_recognizer_client = DocumentAnalysisClient(endpoint=f"https://{args.formrecognizerservice}.cognitiveservices.azure.com/", credential=formrecognizer_creds, headers={"x-ms-useragent": "azure-search-chat-demo/1.0.0"})
        with open(filename, "rb") as f:
            poller = form_recognizer_client.begin_analyze_document("prebuilt-layout", document = f)
        form_recognizer_results = poller.result()

        for page_num, page in enumerate(form_recognizer_results.pages):
            tables_on_page = [table for table in form_recognizer_results.tables if table.bounding_regions[0].page_number == page_num + 1]

            # mark all positions of the table spans in the page
            page_offset = page.spans[0].offset
            page_length = page.spans[0].length
            table_chars = [-1]*page_length
            for table_id, table in enumerate(tables_on_page):
                for span in table.spans:
                    # replace all table spans with "table_id" in table_chars array
                    for i in range(span.length):
                        idx = span.offset - page_offset + i
                        if idx >=0 and idx < page_length:
                            table_chars[idx] = table_id

            # build page text by replacing charcters in table spans with table html
            page_text = ""
            added_tables = set()
            for idx, table_id in enumerate(table_chars):
                if table_id == -1:
                    page_text += form_recognizer_results.content[page_offset + idx]
                elif table_id not in added_tables:
                    page_text += table_to_html(tables_on_page[table_id])
                    added_tables.add(table_id)

            page_text += " "
            page_map.append((page_num, offset, page_text))
            offset += len(page_text)

このコードは何をしてますか?

chatGPTからの回答

このコードは、与えられたPDFドキュメントからテキストデータを抽出し、そのテキストデータをページごとに管理するためのPython関数を定義しています。以下はこのコードの主要な機能についての説明です。

  1. 引数と変数の初期化:

    • 関数 get_document_text(filename) は、1つの引数 filename(ファイル名またはパス)を受け取ります。
    • offset 変数はテキストのオフセットを管理し、初期値として0が設定されます。
    • page_map リストは、ページ番号、テキストのオフセット、およびページのテキストをタプルとして保持するために使用されます。このリストに各ページの情報が追加されます。
  2. 条件分岐:

    • コードは args.localpdfparser という条件に基づいて2つの異なるテキスト抽出方法をサポートしています。
  3. 条件に応じたテキスト抽出:

    • args.localpdfparserTrue の場合:

      • PdfReader クラスを使用して指定されたPDFファイルを読み取り、各ページのテキストを抽出します。
      • 各ページのテキスト、ページ番号、およびテキストのオフセットが page_map リストに追加されます。
    • args.localpdfparserFalse の場合:

      • Azure Form Recognizer サービスを使用してPDFドキュメントからテキストデータを抽出します。
      • form_recognizer_client インスタンスが作成され、指定されたAzureサービスに接続します。
      • ドキュメントをバイナリモードで開いて、Form Recognizer クライアントを使用してドキュメントの解析ジョブを開始します。
      • 解析の結果、各ページに含まれるテーブル情報が取得され、テーブルの位置情報がマークされます。
      • ページのテキストは、テーブル情報を含むHTML形式のテキストとして構築され、page_map リストに追加されます。

このコードは、ローカルのPDFパーサーを使用するか、Azure Form Recognizerを使用してPDFドキュメントからテキストデータを抽出し、各ページのテキストデータを page_map リストに格納することを目的としています。また、テーブル情報も適切に処理し、テキストに統合しています。

エラーがでたときも全部質問すればおっけー
その言語独自の関数とかも、これは何ですか?と聞けば返してくれるのでいいですね!

今回はやっていませんが、これを違う言語に書き換えてください。なんてこともできてやっぱすごいと思います。

このChatGPTがまとめている内容を駆使し、自分が必要な作業だけをやってくれるスクリプトをかくことができました。

import os
import argparse
import glob
import html
import io
import re
import time
from pypdf import PdfReader, PdfWriter
from azure.identity import AzureDeveloperCliCredential
from azure.identity import ManagedIdentityCredential
from azure.core.credentials import AzureKeyCredential
from azure.storage.blob import BlobServiceClient
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import *
from azure.search.documents import SearchClient

MAX_SECTION_LENGTH = 1000
SENTENCE_SEARCH_LIMIT = 100
SECTION_OVERLAP = 100

def blob_name_from_file_page(filename, page = 0):
    if os.path.splitext(filename)[1].lower() == ".pdf":
        return os.path.splitext(os.path.basename(filename))[0] + f"-{page}" + ".pdf"
    else:
        return os.path.basename(filename)

def get_document_text(filename):
    offset = 0
    page_map = []
    reader = PdfReader(filename)
    pages = reader.pages
    for page_num, p in enumerate(pages):
        page_text = p.extract_text()
        page_map.append((page_num, offset, page_text))
        offset += len(page_text)

    return page_map

def split_text(page_map):
    SENTENCE_ENDINGS = [".", "!", "?"]
    WORDS_BREAKS = [",", ";", ":", " ", "(", ")", "[", "]", "{", "}", "\t", "\n"]

    def find_page(offset):
        l = len(page_map)
        for i in range(l - 1):
            if offset >= page_map[i][1] and offset < page_map[i + 1][1]:
                return i
        return l - 1

    all_text = "".join(p[2] for p in page_map)
    length = len(all_text)
    start = 0
    end = length
    while start + SECTION_OVERLAP < length:
        last_word = -1
        end = start + MAX_SECTION_LENGTH

        if end > length:
            end = length
        else:
            # Try to find the end of the sentence
            while end < length and (end - start - MAX_SECTION_LENGTH) < SENTENCE_SEARCH_LIMIT and all_text[end] not in SENTENCE_ENDINGS:
                if all_text[end] in WORDS_BREAKS:
                    last_word = end
                end += 1
            if end < length and all_text[end] not in SENTENCE_ENDINGS and last_word > 0:
                end = last_word # Fall back to at least keeping a whole word
        if end < length:
            end += 1

        # Try to find the start of the sentence or at least a whole word boundary
        last_word = -1
        while start > 0 and start > end - MAX_SECTION_LENGTH - 2 * SENTENCE_SEARCH_LIMIT and all_text[start] not in SENTENCE_ENDINGS:
            if all_text[start] in WORDS_BREAKS:
                last_word = start
            start -= 1
        if all_text[start] not in SENTENCE_ENDINGS and last_word > 0:
            start = last_word
        if start > 0:
            start += 1

        section_text = all_text[start:end]
        yield (section_text, find_page(start))

        last_table_start = section_text.rfind("<table")
        if (last_table_start > 2 * SENTENCE_SEARCH_LIMIT and last_table_start > section_text.rfind("</table")):
            start = min(end - SECTION_OVERLAP, start + last_table_start)
        else:
            start = end - SECTION_OVERLAP
        
    if start + SECTION_OVERLAP < end:
        yield (all_text[start:end], find_page(start))

# セクションごとに分割する
def create_sections(filename, page_map):
    for i, (section, pagenum) in enumerate(split_text(page_map)):
        yield {
            "id": re.sub("[^0-9a-zA-Z_-]","_",f"{filename}-{i}"),
            "content": section,
            "sourcepage": blob_name_from_file_page(filename, pagenum),
            "sourcefile": filename
        }

#blobにアップロード
def upload_blobs(filename):
    storage_creds = default_creds if "ストレージアカウントのキー" is None else "ストレージアカウントのキー"
    blob_service = BlobServiceClient(account_url=f"https://{ストレージアカウント名}.blob.core.windows.net", credential=storage_creds)
    blob_container = blob_service.get_container_client("コンテナ名")
    if not blob_container.exists():
        blob_container.create_container()

    # if file is PDF split into pages and upload each page as a separate blob
    if os.path.splitext(filename)[1].lower() == ".pdf":
        reader = PdfReader(filename)
        pages = reader.pages
        for i in range(len(pages)):
            blob_name = blob_name_from_file_page(filename, i)
            f = io.BytesIO()
            writer = PdfWriter()
            writer.add_page(pages[i])
            writer.write(f)
            f.seek(0)
            blob_container.upload_blob(blob_name, f, overwrite=True)
    else:
        blob_name = blob_name_from_file_page(filename)
        with open(filename,"rb") as data:
            blob_container.upload_blob(blob_name, data, overwrite=True)

def create_search_index():
    search_creds = default_creds if "Cognitive Search account key" is None else AzureKeyCredential("Cognitive Search account key")
    index_client = SearchIndexClient(endpoint=f"https://{Cognitive Searchのリソース名}.search.windows.net/",
                                     credential=search_creds)
    index = SearchIndex(
    name="searchindextest",
    fields=[
        SimpleField(name="id", type="Edm.String", key=True),
        SearchableField(name="content", type="Edm.String", analyzer_name="ja.microsoft"),
        SimpleField(name="category", type="Edm.String", filterable=True, facetable=True),
        SimpleField(name="sourcepage", type="Edm.String", filterable=True, facetable=True),
        SimpleField(name="sourcefile", type="Edm.String", filterable=True, facetable=True)
    ],
    index_client.create_index(index)


# Cognitive Searchにインデックス作る
def index_sections(filename, sections):
    print('index_sections')
    search_creds = default_creds if "Cognitive Search account key" is None else AzureKeyCredential("Cognitive Search account key")
    search_client = SearchClient(endpoint=f"https://{Cognitive Searchのリソース名}.search.windows.net/",
                                    index_name="インデックス名",
                                    credential=search_creds)
    i = 0
    batch = []
    for s in sections:
        batch.append(s)
        i += 1
        if i % 1000 == 0:
            results = search_client.upload_documents(documents=batch)
            succeeded = sum([1 for r in results if r.succeeded])
            print(succeeded)
            batch = []

    if len(batch) > 0:
        results = search_client.upload_documents(documents=batch)
        succeeded = sum([1 for r in results if r.succeeded])
        print(succeeded)

# ここからが呼び出されたときの処理
# pdf分割したのをAzure Blobにあげてくれて、それをもとにindex作ってくれる処理
filename = r"pdfのパス"
upload_blobs(filename)
page_map = get_document_text(filename)
sections = create_sections(os.path.basename(filename), page_map)
create_search_index()
index_sections(os.path.basename(filename), sections)

※自分の検証用のため、ファイルパスもキーとかも直書き
PDFを分割・ストレージにアップロード・インデックスを作成までをしています。

普段かき慣れてない言語でも解読は簡単にできてしまえるのはすごくいいですね!
今後コードの保守とかで解読が必要な時に役立てることができるかもと思いました。

すごい今さらな内容ですが、やっぱすごいと思ったので今後もよい副操縦士(Microsoft風)として活躍していただきたいところです。