0
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@guupys

【コツコツAWS】AWS Lambda - 自分用の論文通知システムを作る -

0
Posted at

生成AIやAgentの研究動向に触れる習慣を作りたく、arXivのAPIとAWSのサービスを使って、１週間ごとにプレプリントのメタデータを収集・通知する構成を実装しました。(ただAWS Step Functionsを使ってみたかった気持ちもあり)

今回の構成

Step Functionsと3つのLambda関数、arXivとOpenAIのAPIを使って、以下の流れでプレプリントのメタデータ収集から要約、通知まで実施。下図がStep Functionsのグラフ。

stepfunctions_graph.png

  1. 直近１週間のプレプリント収集
    arXivのAPIを使って、指定したカテゴリのプレプリントのメタデータを１週間分収集してS3に保存。
  2. 収集したデータからジョブを作成し、Batch APIを実行
    収集したメタデータをGPTで要約させるため、Batchジョブを作成して実行。
  3. Batch完了まで待機 (1日)
    Batch APIは24時間以内に完了するので、Step FunctionsのWait stateで1日待機。
  4. Batch APIの結果からデータ抽出し、保存
    Batch APIから生成結果を抽出してS3に保存。併せて署名付きURLも作成。
  5. Amazon SNSでメール通知
    5にて作成した署名付きURLを指定したメールに送付。

要約生成の部分にBedrockを使用したかったのですが、利用実績の兼ね合いでBatch推論をまだ使えなかったため、OpenAI APIを利用。

Lambda関数を作成した1, 2, 4について、具体的な内容を説明します。
関数の全コードはGithubに載せました。

1. 直近1週間のプレプリント収集

arXivのAPIを使って１週間分のプレプリントのメタデータを収集します。メイン処理が以下の通り。

CollectWeeklyArxiv
def lambda_handler(event, context):
    bucket = os.environ["BUCKET"]
    today = datetime.utcnow()

    all_papers: dict[str, dict] = {}
    for category in event["Categories"]:
        # 1週間分のプレプリントを収集
        request = _make_request(category, today)
        root = _fetch_arxiv_feed(request)
        # categoryごとの抽出結果を統合
        papers = _extract_papers_from_feed(root)
        all_papers.update(papers)

        time.sleep(3)  # API負荷対策
    # csvに変換
    rows = list(all_papers.values())
    csv_text = _to_csv(rows)
    # S3に保存
    key = f"arxiv-papers_{today.strftime('%Y%m%d')}w.csv"
    _upload_csv_to_s3(bucket=bucket, key=key, csv_text=csv_text)

    message = f"Upload {key} ({len(rows)} papers) to {bucket}."
    return {
        "statusCode": 200,
        "bucket": bucket,
        "key": key,
        "body": json.dumps(message),
    }

event["Categories"]はarXivのカテゴリをList形式にしたものです。

関数への入力例
{
    "Categories": ["cs.AI", "cs.CL", "cs.LG", "cs.MA", "econ.EM"]
}

各カテゴリについて、_make_request(category, today)で1週間以内に投稿されたものを収集するクエリを作成し、all_papers.update(papers)の部分で収集した結果を結合する形を取っています。(とりあえず興味のあるカテゴリを拾うためこの形にしました。クエリを作り込めたら要修正)
収集が完了したら、csvファイルでS3に保存して、ファイル名などを次の処理に渡す形としました。(生データも今後何かに活用できれば...)

2. 収集したデータからジョブを作成し、Batch APIを実行

1で保存したデータからOepnAI APIのBatchジョブを作成し、APIを実行します。メイン処理が以下の通り。(openaiのライブラリが必要になるので、レイヤーを追加する必要あり)

MakeArxivBatchInput
def lambda_handler(event, context):
    # S3からCSV取得
    csv_text = _get_s3_object_text(bucket=event["bucket"], key=event["key"])
    # CSVを読み込んでBatch作成, 実行
    reader = csv.DictReader(io.StringIO(csv_text))
    batch = create_batch(reader)

    return {
        "statusCode": 200,
        "batch_id": batch.id,
        "input_file_id": batch.input_file_id,
    }

1で作成したデータはそこまで大きいサイズにならないので、Lambda関数のメモリ上に読み込み。
Batchのジョブ作成では、以下のようなプロンプトを使って、プレプリントから要約の生成をする他に、手法やキーワードの抽出も行いました。

プロンプト作成関数
def _build_prompt(title: str, abstract: str) -> str:
    """プロンプト作成"""
    return f"""
You are a research assistant.

Please process the following paper.

# TITLE
{title}

# ABSTRACT
{abstract}

# TASK
1. Translate the abstract into Japanese.
2. Extract structured information.

Return JSON only in this format:

{{
  "japanese_summary": "...",
  "key_points": [
    "..."
  ],
  "method": "...",
  "contribution": "...",
  "keywords": ["..."]
}}
"""

json形式での出力をプロンプトで指示しているため、Batchのリクエストの方でもschemaのパラメータを使って指定した項目をjsonで出力するようにしています。

Batch APIを実行したら完了まで待ちとなります。後の抽出時に必要となるので、ジョブのidを次のWait Stateに渡して1日待機するようにしました。

4. Batch APIの結果からデータ抽出し、保存

1日待って完了したであろうBatch APIから結果を抽出します。
(1日待たずとも定期的に状態を確認して、完了したら抽出に進む実装をしたり、失敗した時の処理の実装をしたりする方が良いのですが、ご容赦ください)

ExtractArxivBatchOutput
def lambda_handler(event, context):
    batch_id = event["batch_id"]
    bucket = os.environ["BUCKET"]
    # API設定
    os.environ["OPENAI_API_KEY"] = _get_apikey()
    client = OpenAI()
    # batchの状態確認
    batch = client.batches.retrieve(batch_id)
    # 完了以外ならエラーで終了
    output_file_id = getattr(batch, "output_file_id", None)
    if batch.status != "completed" or not output_file_id:
        return {
            "statusCode": 400,
            "batch_id": batch_id,
            "status": batch.status,
            "output_file_id": output_file_id,
            "error_file_id": getattr(batch, "error_file_id", None),
            "message": "Batch not completed or output_file_id is missing.",
        }
    # データ抽出
    content = client.files.content(output_file_id)
    jsonl_bytes = content.read()
    rows = _parse_batch_output_jsonl(jsonl_bytes)
    # CSV化
    csv_text = _rows_to_csv_text(rows)
    # S3へアップロード
    key = f"{OUTPUT_CSV_KEY_PREFIX}{datetime.utcnow().strftime('%Y%m%d')}{OUTPUT_CSV_KEY_SUFFIX}"
    _upload_text_to_s3(bucket=bucket, key=key, text=csv_text)
    # 署名付きURL発行
    presigned_url = _generate_presigned_get_url(
        bucket=bucket,
        key=key,
        expires_in=PRESIGNED_URL_EXPIRES_IN,
    )

    message = f"Upload {key} ({len(rows)} papers) to {bucket}."
    return {
        "statusCode": 200,
        "bucket": bucket,
        "key": key,
        "body": json.dumps(message),
        "presigned_url": presigned_url,
        "expires_in": PRESIGNED_URL_EXPIRES_IN,
    }

抽出した結果はcsvファイルでS3に保存し、署名付きURLを発行します。(わざわざS3まで取りに行くのも面倒なので、メールで通知してダウンロードできる形にしたかった)
署名付きURLを次のSNSのステップに渡せば、SNSが指定したメールにURLを載せて送信するので、これで完了です。

まとめ

今回は個人的にやりたかったarXivのプレプリント収集とその通知を行うシステムを、Step Functionsで実装しました。上記の内容で構築したステートマシンを、Amazon EventBridgeで定期実行するスケジュールを作成すれば、定期的な通知システムの完成です。

初めてStep Functionsを触りましたが、今回のようなシンプルなものであれば、特にドキュメントを見ることもなく直感的に作ることができました。(Wait Stateのおかげで、Batch APIの状態を見に行くような実装も不要になってよかったです)

とりあえず初案として作っただけで利便性に課題も残っているので、使いながらLambda関数の修正や、他サービスとの連携を試したいと思います。

0
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?