AWS Bedrock の Claude で標準入力の文章を15分の1の文字数に要約するPythonスクリプト

summerize-stdin-by-bedrock.py

# AWS Bedrock の Claude で標準入力の文章を15分の1の文字数に要約するPythonスクリプト
# つかいかた;
# pip3 install -u boto3
## デフォの15以外の割合に設定したければ
# export SUMMERIZE_RATE=10
## デフォの anthropic.claude-v2 以外のモデルを使うなら
# export MODEL_NAME="xxxxxxxxxx"
# cat hoge.txt | python3 ./summerize-stdin-by-bedrock.py

import os
import sys
import io
import unicodedata
import json
import boto3

sys.stdin = io.TextIOWrapper(sys.stdin.buffer, encoding='utf-8')
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

# 要約したい文章を stdin で流し込む
text = sys.stdin.read()
# OpenAI Whisper で文字起こした文章は怪しい文字列が含まれてることがある。
# Unicode正規化し、きちんと扱えるようにする。
# https://docs.python.org/ja/3/library/unicodedata.html
text = unicodedata.normalize('NFKC', text)

# 総文字数を summerize_rate で割った量に要約してもらう
summerize_rate = int(os.environ.get('SUMMERIZE_RATE', '15'))
summerize_count = int(len(text)/summerize_rate)

# プロンプトを組み立てる
prompt = "Human: 以下の内容を" + str(summerize_count) + "文字で要約してください\n\n###\n\n" + text + "\n\nAssistant:"

bedrock_runtime = boto3.client("bedrock-runtime", region_name=os.environ.get('AWS_REGION', 'us-east-1'))

body = json.dumps(
    {
        "prompt": prompt,
        "max_tokens_to_sample": 500,
    }
)

resp = bedrock_runtime.invoke_model(
    modelId=os.environ.get('MODEL_NAME', 'anthropic.claude-v2'),
    body=body,
    contentType="application/json",
    accept="application/json",
)
answer = resp["body"].read().decode()

print(json.loads(answer)["completion"])