Pythonの並列処理を使ったコードの例として、Azure OpenAIとLlamaIndexを用いて複数のテキストを並行して要約する方法について書いていきます。LLM APIを逐次実行すると時間がかかると思い、速く回せる方法を探していたら、ついでにいろいろ勉強になったので備忘録的に置いておきます。

1. APIクライアントの設定

Azure OpenAIのAPIキーやエンドポイント情報を使って、llmという変数でクライアントをセットアップしています。

llm = AzureOpenAI(
    model="gpt-4o-mini",
    max_tokens=4096,
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2024-08-01-preview",
    deployment_name="gpt-4o-mini"
)

2. 要約する関数 `asummarize`

asummarizeはテキストを要約する非同期関数で、awaitでAPIレスポンスを待つ仕組み。

async def asummarize(text):
    res = await llm.achat(messages=SUMMARIZE_PROMPT_TEMPLATE.format_messages(context_str=text))
    return res.message.content

3. 並行処理 `parallel_summarize`

parallel_summarize関数で複数のテキストを一気に要約します

async def parallel_summarize(texts: list[str], max_concurrent: int = 4) -> list[str]:
    semaphore = Semaphore(max_concurrent)
    
    async def process_with_progress(text: str, index: int) -> tuple[int, str]:
        async with semaphore:
            summary = await asummarize(text)
            return index, summary
            
    tasks = [asyncio.create_task(process_with_progress(text, i)) for i, text in enumerate(texts)]
    results = [None] * len(texts)
    for completed_task in tqdm(asyncio.as_completed(tasks), total=len(texts), desc="Summarizing"):
        index, summary = await completed_task
        results[index] = summary
    return results

Semaphoreを使った同時実行数の制御と、結果の順番を保つための工夫がポイントです。

Semaphoreによる同時実行数の制御

Semaphoreは、同時に実行できる処理の数を制限する仕組みです。APIへの大量のリクエストを一度に送りすぎると、サーバーの負荷が増加してエラーが発生する可能性があります。そこで、Semaphoreを使って一度に実行するリクエスト数を制限し、システムに過剰な負担がかからないようにしています。

semaphore = Semaphore(max_concurrent)

ここでは、max_concurrentの値で同時実行数を指定します。例えばmax_concurrent=4にすると、同時に4つのリクエストまでしか実行しないようになります。

具体的な流れ:
- async with semaphoreという記述で、セマフォの枠内で処理を行います。
- 4つのリクエストが同時に実行され、1つが終わると次のリクエストがスタートします。

async def process_with_progress(text: str, index: int) -> tuple[int, str]:
    async with semaphore:
        summary = await asummarize(text)
        return index, summary

結果の順番を保つ工夫

asyncio.as_completed(tasks)は、完了した順にタスクを取り出すため、順番が保証されていません。しかし、ここではリストresultsを用意し、タスクごとに対応するインデックスに結果を格納することで、最終的な順番を保っています。

results = [None] * len(texts)
for completed_task in tqdm(asyncio.as_completed(tasks), total=len(texts), desc="Summarizing"):
    index, summary = await completed_task
    results[index] = summary

この工夫により、並行処理で完了順が異なる場合でも、resultsリストには元のテキストの順番に沿った要約が格納されます。

コード全体

import os
import asyncio
from asyncio import Semaphore

from llama_index.llms.azure_openai import AzureOpenAI
from tqdm.auto import tqdm

llm = AzureOpenAI(
    temperature=0,
    model="gpt-4o-mini",
    max_tokens=4096,
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2024-08-01-preview",
    deployment_name="gpt-4o-mini"
)

async def asummarize(text):
    res = await llm.achat(messages=SUMMARIZE_PROMPT_TEMPLATE.format_messages(
        context_str=text
    ))
    return res.message.content


async def parallel_summarize(texts: list[str], max_concurrent: int = 4) -> list[str]:
    semaphore = Semaphore(max_concurrent)

    async def process_with_progress(text: str, index: int) -> tuple[int, str]:
        async with semaphore:
            summary = await asummarize(text)
            return index, summary

    tasks: list[asyncio.Task] = [
        asyncio.create_task(process_with_progress(text, i))
        for i, text in enumerate(texts)
    ]

    results = [None] * len(texts)  # 結果を格納するリストを事前に作成

    for completed_task in tqdm(asyncio.as_completed(tasks), total=len(texts), desc="Summarizing"):
        index, summary = await completed_task
        results[index] = summary

    return results