Chain Of Density

W&B Weave で Chain Of Density を使う方法を学ぶ

これはインタラクティブなノートブックです。ローカルで実行するか、以下のリンクから開くことができます:

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

​

Chain of Density を用いた要約

重要な詳細を保持しながら複雑な技術文書を要約することは、難しい作業です。Chain of Density (CoD) 要約手法は、要約を反復的に洗練させて、より簡潔かつ情報密度の高いものにすることで、この課題に対する有効な解決策となります。このガイドでは、Weave を使ってアプリケーションを追跡・評価しながら CoD を実装する方法を示します。

​

Chain of Density 要約とは？

Chain of Density (CoD) は、反復的な要約手法であり、要約を段階的に洗練していくことで、より簡潔かつ情報密度の高い要約を生成する手法です。仕組みは次のとおりです:

1. まず初期要約を作成する
2. 要約を反復的に洗練し、重要な情報を保持しつつ、より簡潔にしていく
3. 各反復ごとに、Entities や技術的な詳細の密度を高めていく

このアプローチは、詳細な情報の保持が重要となる科学論文や技術文書を要約する際に特に有用です。

​

なぜ Weave を使うのか？

このチュートリアルでは、Weave を使って ArXiv 論文向けの Chain of Density 要約パイプラインを実装し、評価します。次のことを学びます:

1. LLM パイプラインを追跡する: Weave を使って、要約プロセスの入力、出力、および中間ステップを自動でログに記録します。
2. LLM の出力を評価する: Weave の組み込みツールを使って、要約結果を厳密かつ公平に評価します。
3. 合成可能なオペレーションを構築する: 要約パイプラインのさまざまな部分で再利用できる Weave のオペレーション（operation）を組み合わせて構築します。
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 プロジェクトが初期化されます。

​

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 モデルを作成する

では、要約パイプラインを Weave モデルとして定義してみましょう。

# Weave モデル
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 のモデルとしてカプセル化しており、次のような主な利点があります。

1. 実験の自動トラッキング: Weave は、モデルの各 run に対して入力・出力・パラメータを記録します。
2. バージョニング: モデルの属性やコードの変更は自動的にバージョン管理され、要約パイプラインが時間とともにどのように進化してきたかの明確な履歴が作成されます。
3. 再現性: バージョニングとトラッキングにより、要約パイプラインの過去の任意の結果や設定を簡単に再現できます。
4. ハイパーパラメータ管理: モデル属性（model や density_iterations など）が明確に定義され、run 間で追跡されるため、実験がしやすくなります。
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 Dataset を作成して評価を実行する

パイプラインを評価するために、Weave Dataset を作成し、評価を実行します。

# 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 operation を作成する方法
2. パイプラインを Weave モデルでラップして、追跡と評価を容易にする方法
3. Weave operation を使ってカスタム評価メトリクスを実装する方法
4. データセットを作成し、パイプラインの評価を実行する方法

Weave のシームレスなインテグレーションにより、要約プロセス全体を通して入力、出力、および中間ステップを追跡できるため、LLM アプリケーションのデバッグ、最適化、および評価が容易になります。 この例を拡張して、より大きなデータセットを扱ったり、より高度な評価メトリクスを実装したり、他の LLM ワークフローとインテグレーションしたりすることもできます。

W&B 上の完全なレポートを表示