Chain Of Density

反復的なテキスト圧縮と評価を行うための chain-of-density 要約手法を、W&B Weave で実装します。

これはインタラクティブなノートブックです。ローカルで実行するか、以下のリンクを使用できます。

• Google Colab で開く
• GitHub でソースを表示

​

Chain of Density による要約

重要な詳細を維持したまま複雑な技術文書を要約するのは、容易ではありません。Chain of Density (CoD) 要約手法は、要約を繰り返し改善して、より簡潔で情報密度の高いものにすることで、この課題に対応します。このガイドでは、Weave を使用してアプリケーションをトラッキングおよび評価しながら、CoD を実装する方法を説明します。

​

Chain of Density要約とは何ですか？

Chain of Density (CoD) は、反復的に要約を洗練しながら、より簡潔で情報密度の高い要約を生成する手法です。具体的には、次のように機能します。

1. 最初の要約から始めます
2. 重要な情報を維持したまま、要約を反復的に洗練して、より簡潔にします
3. 反復のたびに、Entities と技術的な詳細の密度を高めます

このアプローチは、詳細な情報の保持が重要な科学論文や技術文書の要約に特に適しています。

​

なぜ Weave を使用するのか？

このチュートリアルでは、ArXiv 論文向けの Chain of Density 要約パイプラインを、Weave を使用して実装・評価します。学習できる内容は次のとおりです。

1. LLM パイプラインをトラッキングする: Weave を使用して、要約プロセスの入力、出力、途中のステップを自動的にログします。
2. LLM の出力を評価する: Weave の組み込みツールを使用して、要約を同一条件で厳密に評価します。
3. 構成可能なオペレーションを構築する: 要約パイプラインのさまざまな部分で Weave のオペレーションを組み合わせ、再利用します。
4. シームレスに統合する: 既存の Python コードに、最小限のオーバーヘッドで Weave を追加します。

このチュートリアルを終える頃には、モデルサービング、評価、結果のトラッキングに Weave の機能を活用した CoD 要約パイプラインを作成できるようになります。

​

環境を設定する

まず、環境を設定し、必要なライブラリをインポートします。

!pip install -qU anthropic weave pydantic requests PyPDF2 set-env-colab-kaggle-dotenv

AnthropicのAPIキーを取得するには、次の手順に従います。

1. https://www.anthropic.com でアカウントを作成します
2. アカウント設定のAPIセクションにアクセスします
3. 新しいAPIキーを生成します
4. APIキーを .env ファイルに安全に保存します

import io
import os
from datetime import datetime, timezone

import anthropic
import requests
from pydantic import BaseModel
from PyPDF2 import PdfReader
from set_env import set_env

import weave

set_env("WANDB_API_KEY")
set_env("ANTHROPIC_API_KEY")

weave.init("summarization-chain-of-density-cookbook")
anthropic_client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

実験のトラッキングにはWeaveを使用し、テキスト生成にはAnthropicのClaudeモデルを使用します。weave.init(<project name>)を呼び出すと、要約タスク用の新しいWeave projectが設定されます。

​

ArxivPaper モデルを定義する

データを表現するためのシンプルな ArxivPaper クラスを作成します。

# ArxivPaper モデルを定義する
class ArxivPaper(BaseModel):
    entry_id: str
    updated: datetime
    published: datetime
    title: str
    authors: list[str]
    summary: str
    pdf_url: str

# サンプルの ArxivPaper を作成する
arxiv_paper = ArxivPaper(
    entry_id="http://arxiv.org/abs/2406.04744v1",
    updated=datetime(2024, 6, 7, 8, 43, 7, tzinfo=timezone.utc),
    published=datetime(2024, 6, 7, 8, 43, 7, tzinfo=timezone.utc),
    title="CRAG -- Comprehensive RAG Benchmark",
    authors=["Xiao Yang", "Kai Sun", "Hao Xin"],  # 簡略化のため省略
    summary="Retrieval-Augmented Generation (RAG) has recently emerged as a promising solution...",  # 省略
    pdf_url="https://arxiv.org/pdf/2406.04744",
)

このクラスは、要約パイプラインへの入力となる ArXiv 論文のメタデータと内容を保持します。

​

PDF コンテンツを読み込む

論文全体の内容を扱うため、PDF からテキストを読み込んで抽出する関数を追加します。

@weave.op()
def load_pdf(pdf_url: str) -> str:
    # PDFをダウンロードする
    response = requests.get(pdf_url)
    pdf_file = io.BytesIO(response.content)

    # PDFを読み込む
    pdf_reader = PdfReader(pdf_file)

    # 全ページからテキストを抽出する
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()

    return text

​

Chain of Density の要約を実装する

それでは、Weaveのオペレーションを使って、CoD要約の中核ロジックを実装します。

# Chain of Density Summarization
@weave.op()
def summarize_current_summary(
    document: str,
    instruction: str,
    current_summary: str = "",
    iteration: int = 1,
    model: str = "claude-3-sonnet-20240229",
):
    prompt = f"""
    Document: {document}
    Current summary: {current_summary}
    Instruction to focus on: {instruction}
    Iteration: {iteration}

    Generate an increasingly concise, entity-dense, and highly technical summary from the provided document that specifically addresses the given instruction.
    """
    response = anthropic_client.messages.create(
        model=model, max_tokens=4096, messages=[{"role": "user", "content": prompt}]
    )
    return response.content[0].text

@weave.op()
def iterative_density_summarization(
    document: str,
    instruction: str,
    current_summary: str,
    density_iterations: int,
    model: str = "claude-3-sonnet-20240229",
):
    iteration_summaries = []
    for iteration in range(1, density_iterations + 1):
        current_summary = summarize_current_summary(
            document, instruction, current_summary, iteration, model
        )
        iteration_summaries.append(current_summary)
    return current_summary, iteration_summaries

@weave.op()
def final_summary(
    instruction: str, current_summary: str, model: str = "claude-3-sonnet-20240229"
):
    prompt = f"""
    Given this summary: {current_summary}
    And this instruction to focus on: {instruction}
    Create an extremely dense, final summary that captures all key technical information in the most concise form possible, while specifically addressing the given instruction.
    """
    return (
        anthropic_client.messages.create(
            model=model, max_tokens=4096, messages=[{"role": "user", "content": prompt}]
        )
        .content[0]
        .text
    )

@weave.op()
def chain_of_density_summarization(
    document: str,
    instruction: str,
    current_summary: str = "",
    model: str = "claude-3-sonnet-20240229",
    density_iterations: int = 2,
):
    current_summary, iteration_summaries = iterative_density_summarization(
        document, instruction, current_summary, density_iterations, model
    )
    final_summary_text = final_summary(instruction, current_summary, model)
    return {
        "final_summary": final_summary_text,
        "accumulated_summary": current_summary,
        "iteration_summaries": iteration_summaries,
    }

各関数の役割は次のとおりです。

• summarize_current_summary: 現在の状態に基づいて、1 回分の要約イテレーションを生成します。
• iterative_density_summarization: summarize_current_summary を複数回呼び出して、CoD 手法を適用します。
• chain_of_density_summarization: 要約プロセス全体を制御し、結果を返します。

@weave.op() デコレータを使用することで、Weave がこれらの関数の入力、出力、実行をトラッキングできるようにしています。

​

Weave Model を作成する

それでは、要約パイプラインを Weave Model としてラップします。

# Weave Model
class ArxivChainOfDensityPipeline(weave.Model):
    model: str = "claude-3-sonnet-20240229"
    density_iterations: int = 3

    @weave.op()
    def predict(self, paper: ArxivPaper, instruction: str) -> dict:
        text = load_pdf(paper.pdf_url)
        result = chain_of_density_summarization(
            text,
            instruction,
            model=self.model,
            density_iterations=self.density_iterations,
        )
        return result

この ArxivChainOfDensityPipeline クラスは、要約ロジックを Weave Model としてカプセル化しており、次のような主な利点があります。

1. 自動的な実験管理: Weave は、モデルの各 run について inputs、outputs、parameters を取得します。
2. バージョン管理: モデルの属性やコードへの変更は自動的にバージョン管理されるため、要約パイプラインが時間とともにどのように変化したかを明確に追跡できます。
3. 再現性: バージョン管理とトラッキングにより、要約パイプラインの過去の結果や設定を簡単に再現できます。
4. ハイパーパラメーター管理: モデルの属性 (model や density_iterations など) は明確に定義され、異なる Runs 間でトラッキングされるため、実験を進めやすくなります。
5. Weave エコシステムとのインテグレーション: weave.Model を使用すると、評価やサービング機能など、他の Weave ツールとシームレスに連携できます。

​

評価メトリクスを実装する

要約の品質を評価するために、シンプルな評価メトリクスを実装します：

import json

@weave.op()
def evaluate_summary(
    summary: str, instruction: str, model: str = "claude-3-sonnet-20240229"
) -> dict:
    prompt = f"""
    Summary: {summary}
    Instruction: {instruction}

    Evaluate the summary based on the following criteria:
    1. Relevance (1-5): How well does the summary address the given instruction?
    2. Conciseness (1-5): How concise is the summary while retaining key information?
    3. Technical Accuracy (1-5): How accurately does the summary convey technical details?

    Your response MUST be in the following JSON format:
    {{
        "relevance": {{
            "score": <int>,
            "explanation": "<string>"
        }},
        "conciseness": {{
            "score": <int>,
            "explanation": "<string>"
        }},
        "technical_accuracy": {{
            "score": <int>,
            "explanation": "<string>"
        }}
    }}

    Ensure that the scores are integers between 1 and 5, and that the explanations are concise.
    """
    response = anthropic_client.messages.create(
        model=model, max_tokens=1000, messages=[{"role": "user", "content": prompt}]
    )
    print(response.content[0].text)

    eval_dict = json.loads(response.content[0].text)

    return {
        "relevance": eval_dict["relevance"]["score"],
        "conciseness": eval_dict["conciseness"]["score"],
        "technical_accuracy": eval_dict["technical_accuracy"]["score"],
        "average_score": sum(eval_dict[k]["score"] for k in eval_dict) / 3,
        "evaluation_text": response.content[0].text,
    }

これらの評価関数では、関連性、簡潔さ、技術的な正確性の観点から生成された要約の品質を評価するために、Claude モデルを使用します。

​

Weave データセットを作成して評価を実行する

パイプラインを評価するために、Weave データセットを作成し、評価を実行します。

# Weave Dataset を作成する
dataset = weave.Dataset(
    name="arxiv_papers",
    rows=[
        {
            "paper": arxiv_paper,
            "instruction": "What was the approach to experimenting with different data mixtures?",
        },
    ],
)

weave.publish(dataset)

今回の評価では、LLM-as-a-judge アプローチを使用します。この手法では、あるモデルやシステムが生成した出力の品質を、別の言語モデルを使って評価します。LLM の理解力と推論能力を活用することで、従来のメトリクスでは十分に捉えきれないタスクでも、よりきめ細かな評価が可能になります。

# スコアラー関数を定義する
@weave.op()
def quality_scorer(instruction: str, output: dict) -> dict:
    result = evaluate_summary(output["final_summary"], instruction)
    return result
python
# 評価を実行する
evaluation = weave.Evaluation(dataset=dataset, scorers=[quality_scorer])
arxiv_chain_of_density_pipeline = ArxivChainOfDensityPipeline()
results = await evaluation.evaluate(arxiv_chain_of_density_pipeline)

このコードは、ArXiv のサンプル論文を含むデータセットを作成し、品質スコアラーを定義して、要約パイプラインを評価します。

​

結論

この例では、Weave を使って ArXiv 論文向けの Chain of Density の要約パイプラインを実装する方法を紹介しました。具体的には、以下の方法を示しました。

1. 要約プロセスの各ステップに対応する Weave オペレーションを作成する
2. トラッキングと評価を容易にするために、パイプラインを Weave Model でラップする
3. Weave オペレーションを使用してカスタム評価メトリクスを実装する
4. データセットを作成し、パイプラインの評価を実施する

Weave のシームレスな統合により、要約プロセス全体を通じて入力、出力、中間ステップをトラッキングできるため、LLM アプリケーションのデバッグ、最適化、評価が容易になります。 この例は、より大規模なデータセットに対応させたり、より高度な評価メトリクスを実装したり、ほかの LLM ワークフローと統合したりするように拡張できます。

W&B で full report を表示