Automated PDF summarization of arXiv papers with Claude 3.5 Sonnet and W&B Weave

import weave
weave.init("arxiv-paper-summarization-anthropic")

import anthropic
import os

client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

from datetime import datetime, timezone
from pydantic import BaseModel, Field
from typing import List, Optional

class Author(BaseModel):
    full_name: str

class Link(BaseModel):
    href: str
    title: Optional[str] = None
    rel: Optional[str] = None
    content_type: Optional[str] = None

class ArxivPaper(BaseModel):
    entry_id: str
    updated: datetime
    published: datetime
    title: str
    authors: List[Author]
    summary: str
    comment: Optional[str] = None
    journal_ref: Optional[str] = None
    doi: Optional[str] = None
    primary_category: str
    categories: List[str]
    links: List[Link]
    pdf_url: Optional[str] = None

    def __getitem__(self, key):
        return getattr(self, key)

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=[Author(full_name="Xiao Yang"), Author(full_name="Kai Sun"), Author(full_name="Hao Xin")],
    summary="The long summary from the paper",
    doi="10.48550/arXiv.2406.04744",
    primary_category="cs.CL",
    categories=["cs.CL"],
    links=[
        Link(href="<https://arxiv.org/abs/2406.04744>", title="Abstract", rel="alternate"),
        Link(href="<https://arxiv.org/pdf/2406.04744>", title="pdf", rel="related")
    ]
    pdf_url="https://arxiv.org/pdf/2406.04744"
)

import requests
import io
import PyPDF2

def load_pdf(arxiv_result):
    pdf_url = arxiv_result["pdf_url"]
    response = requests.get(pdf_url)
    pdf_file = io.BytesIO(response.content)
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    return pdf_reader

from pdf2image import convert_from_bytes
from PIL import Image
import base64

def convert_vector_graphic_page_to_image(pdf_page: PyPDF2.PageObject, scale_factor: float = 0.5) -> Optional[str]:
    def get_object(obj):
        return obj.get_object() if isinstance(obj, PyPDF2.generic.IndirectObject) else obj

    resources = get_object(pdf_page.get('/Resources', {}))
    xobject = get_object(resources.get('/XObject', {}))

    for obj in xobject.values():
        obj = get_object(obj)
        if isinstance(obj, dict) and obj.get('/Subtype') == '/Form':  # Indicates a vector graphic
            pdf_bytes = io.BytesIO()
            PyPDF2.PdfWriter().add_page(pdf_page).write(pdf_bytes)
            pdf_bytes.seek(0)

            images = convert_from_bytes(pdf_bytes.getvalue(), fmt='png')
            if images:
                image = images[0]
                new_size = tuple(int(dim * scale_factor) for dim in image.size)
                image = image.resize(new_size, Image.LANCZOS)
                img_byte_arr = io.BytesIO()
                image.save(img_byte_arr, format='PNG')
                img_str = base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
                data_url = f"data:image/png;base64,{img_str}"
                return data_url

    return None

@weave.op()
def process_figure_image(data_url: str, model: str = "claude-3-5-sonnet-20240620") -> str:
    img_str = data_url.split(",")[1]
    prompt = """Analyze this image as if it's a figure from a scientific research paper. Provide a detailed technical description addressing:
    1. Type of figure (e.g., graph, diagram, flowchart, experimental setup)
    2. Key components or variables represented
    3. Relationships or trends depicted
    4. Quantitative information (if present)
    5. Methodology or process illustrated (if applicable)
    6. Potential implications or conclusions that can be drawn
    7. Any limitations or assumptions evident in the figure
    Focus on technical accuracy and relevance to scientific research. Avoid general descriptions and concentrate on the specific scientific content presented."""

    response = client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{
            "role": "user",
            "content": [
                {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": img_str}},
                {"type": "text", "text": prompt}
            ]
        }]
    )
    return response.content[0].text

@weave.op()
def process_vector_image_pdf(data_url: str, model: str = "claude-3-5-sonnet-20240620") -> str:
    img_str = data_url.split(",")[1]
    prompt = """This image is a full page from a scientific paper PDF, converted to PNG format. It may contain one or more vector graphic figures or charts. Your task is to:
    1. Identify and focus solely on the vector graphic figures or charts within the page.
    2. For each identified figure or chart, provide a detailed technical analysis addressing:
       a. Type of figure (e.g., graph, diagram, flowchart)
       b. Key components or variables represented
       c. Relationships or trends depicted
       d. Quantitative information (if present)
       e. Methodology or process illustrated (if applicable)
       f. Potential implications or conclusions that can be drawn
    3. Ignore any text or other elements on the page that are not part of the vector graphic figures.
    4. If multiple figures are present, analyze each separately and clearly indicate which figure you are describing.
    Focus on providing accurate, technical descriptions of the vector graphic content only."""

    response = client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{
            "role": "user",
            "content": [
                {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": img_str}},
                {"type": "text", "text": prompt}
            ]
        }]
    )
    return response.content[0].text

import filetype

@weave.op()
def extract_images(paper, model="claude-3-5-sonnet-20240620"):
    """Extract text and images from PDF content."""
    pdf_reader = load_pdf(paper)
    all_images = []

    aspect_ratio_sizes = {
        (1, 1): (1092, 1092),
        (3, 4): (951, 1268),
        (2, 3): (896, 1344),
        (9, 16): (819, 1456),
        (1, 2): (784, 1568)
    }

    def get_closest_aspect_ratio(width, height):
        img_ratio = width / height
        return min(aspect_ratio_sizes.keys(), key=lambda x: abs(x[0]/x[1] - img_ratio))

    for page in pdf_reader.pages:
        images = []

        for image in page.images:
            img_data = image.data
            kind = filetype.guess(img_data)
            if kind is None:
                print("Cannot guess file type!")
                continue

            # Resize image if necessary
            img = Image.open(io.BytesIO(img_data))
            closest_ratio = get_closest_aspect_ratio(img.width, img.height)
            new_size = aspect_ratio_sizes[closest_ratio]
            
            if img.width != new_size[0] or img.height != new_size[1]:
                img = img.resize(new_size, Image.LANCZOS)
                
                # Convert resized image back to bytes
                img_byte_arr = io.BytesIO()
                img.save(img_byte_arr, format=img.format if img.format else 'PNG')
                img_data = img_byte_arr.getvalue()


            img_str = base64.b64encode(img_data).decode("utf-8")
            data_url = f"data:{kind.mime};base64,{img_str}"
            try:
                images.append(
                    {"image": data_url, "description": process_figure_image(data_url, model=model)}
                )
            except Exception as e:
                print(f"Error processing image: {e}")
                images.append({"image": data_url, "description": ""})

        vector_graphics_image_data_url = convert_vector_graphic_page_to_image(page)
        if vector_graphics_image_data_url:
            images.append({"image": vector_graphics_image_data_url, "description": process_vector_image_pdf(vector_graphics_image_data_url, model=model)})
        all_images.append(images)

    return all_images

@weave.op()
def replace_images_with_descriptions(paper, images):
    pdf_reader = load_pdf(paper)
    text = ""
    for page_num, page in enumerate(pdf_reader.pages):
        text += page.extract_text() + "\n\n"
        if images[page_num] and len(images[page_num]) > 0:
            text += f"\n\n[Image Descriptions for page {page_num+1}]\n"
            for image_num, image in enumerate(images[page_num]):
                text += f"\n[Image {image_num+1}]: {image['description']}\n"
            text += "[END OF IMAGE DESCRIPTIONS]\n"

    return text

@weave.op()
def summarize_current_summary(document, instruction, current_summary="", iteration=1, model="claude-3-5-sonnet-20240620"):
    # Define the maximum number of tokens for the model's response
    max_tokens = 4096
    
    # Construct the prompt for the LLM
    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 using the below approach:
    1. Carefully read the current summary and the instruction.
    2. Identify 1-3 new, important technical entities or ideas from the original text that:
       - Are directly relevant to the instruction
       - Are not yet present in the current summary
       - Add significant, specific information to the summary
       - Are preferably 5 words or fewer
       - May include methodologies, algorithms, metrics, or key findings
       - Ensure to include this in the output before the summary
    3. Write a new summary that:
       - Incorporates the newly identified entities/ideas
       - Retains all crucial information from the current summary
       - Increases overall information density
       - Remains focused on addressing the instruction
       - Utilizes the response window of {max_tokens} tokens
    
    Guidelines:
    - Prioritize technical accuracy and specificity over general readability
    - Use precise terminology, domain-specific jargon, and include quantitative details where relevant
    - Ensure all information is directly related to the instruction
    - Make every word count: rewrite to improve density and make space for new technical entities
    - Employ fusion, compression, and removal of less informative phrases to increase density
    - Never drop entities or technical details from the current summary that are relevant to the instruction
    - Maintain coherence while maximizing information density
    
    Your goal is to create a summary that is noticeably denser, more technical, and more informative than the previous one, utilizing the response window of {max_tokens} tokens while staying laser-focused on the instruction. The summary should be suitable for an expert audience in the field.
    """
    
    # Make the API call to the LLM
    response = anthropic_client.messages.create(
        model=model,
        max_tokens=max_tokens,
        messages=[{"role": "user", "content": prompt}]
    )
    
    # Return the generated summary
    return response.content[0].text

@weave.op()
def iterative_density_summarization(document, instruction, current_summary, density_iterations, model):
    # Initialize a list to store summaries from each iteration
    iteration_summaries = []

    # Iterate through the specified number of density iterations
    for iteration in range(1, density_iterations + 1):
        # Generate a new summary based on the current summary and document
        current_summary = summarize_current_summary(document, instruction, current_summary, iteration, model)
        
        # Add the new summary to the list of iteration summaries
        iteration_summaries.append(current_summary)
        
        # Print the current iteration and summary for monitoring
        print(f"Iteration {iteration}:\n{current_summary}\n")

    # Return the final summary and the list of all iteration summaries
    return current_summary, iteration_summaries

@weave.op()
def final_summary(instruction, current_summary, model):
    # Construct the prompt for the final summary generation
    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. Follow these guidelines:
        1. Aim to reduce length by 30-40% while retaining all critical technical content relevant to the instruction.
        2. Prioritize highly specific methodologies, algorithms, metrics, and findings that directly address the instruction.
        3. Preserve precise quantitative data, including statistical significance and error margins where applicable and relevant to the instruction.
        4. Maintain the use of domain-specific terminology and technical jargon pertinent to the instruction.
        5. Ensure that all key entities and concepts from the original summary that relate to the instruction are represented.
        6. Use compact phrasing and remove any remaining non-essential information that doesn't directly contribute to addressing the instruction.
        7. If relevant to the instruction, include brief mentions of limitations, assumptions, or conflicting viewpoints.
        8. Optimize for information density while maintaining coherence for an expert audience, always keeping the focus on the given instruction.
        The final summary should be a highly concentrated, technical distillation of the research that specifically addresses the given instruction, suitable for specialists in the field."""
    )

    # Make the API call to the LLM for the final summary
    response = anthropic_client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{"role": "user", "content": prompt}]
    )

    # Return the generated final summary
    return response.content[0].text

@weave.op()
def chain_of_density_summarization(document, instruction, current_summary="", model="claude-3-5-sonnet-20240620", density_iterations=2):
    # Perform iterative density summarization
    current_summary, iteration_summaries = iterative_density_summarization(
        document, instruction, current_summary, density_iterations, model
    )
    
    # Generate the final, highly condensed summary
    final_summary_text = final_summary(instruction, current_summary, model)
    
    # Print the final summary for monitoring
    print(f"Final Summary:\n{final_summary_text}\n")
    
    # Return a dictionary containing all generated summaries
    return {
        "final_summary": final_summary_text,
        "accumulated_summary": current_summary,
        "iteration_summaries": iteration_summaries,
    }


class ArxivChainOfDensityPipeline(weave.Model):
    model: str = "claude-3-5-sonnet-20240620"    density_iterations: int = 3    def __init__(self, model: str = "claude-3-5-sonnet-20240620", density_iterations: int = 3):
        super().__init__()
        self.model = model
        self.density_iterations = density_iterations
    @weave.op()
    def predict(self, paper: ArxivPaper, instruction: str) -> dict:
        extracted_images = extract_images(paper)
        cleaned_text = replace_images_with_descriptions(paper, extracted_images)
        result = chain_of_density_summarization(cleaned_text, instruction, model=self.model, density_iterations=self.density_iterations)
        return result

eval_papers = [arxiv_paper3]
eval_instructions = [
    "Summarize the key methodologies and novel contributions of this research, focusing on their potential impact in the field.",
]
eval_data = list(product(eval_papers, eval_instructions))
dataset = weave.Dataset(name="we-paper-reading-eval-data", rows=[{"paper": arxiv_paper, "instruction": instruction, "summary": arxiv_paper.summary} for arxiv_paper, instruction in eval_data])
weave.publish(dataset)

@weave.op()
def score_summary(summary, summary_type, instruction, model):
    openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    # Construct a detailed prompt for the GPT model to evaluate the summary    prompt = f"""Evaluate the quality of the following {summary_type} based on how well it addresses the given instruction. Use the scoring rules below to calculate three numerical scores between 0 and 10.Instruction: {instruction}{summary_type}:{summary}Scoring Rules:1. Relevance (0-5): [Detailed scoring criteria for relevance]2. Technical Quality (0-5): [Detailed scoring criteria for technical quality]3. Conciseness (0-5): [Detailed scoring criteria for conciseness]Provide your evaluation in the following JSON format:{{    "relevance": {{        "score": <float>    }},    "technical_quality": {{        "score": <float>    }},    "conciseness": {{        "score": <float>    }}}}Ensure your response is ONLY valid JSON. Do not include any other text outside the JSON object.Ensure you have the keys: relevance, technical_quality, conciseness, each containing only a score.Ensure each score is a float between 0 and 10, using the scoring rules provided above."""    # Make an API call to the GPT model for evaluation    response = openai_client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"}
    )
    # Parse and return the JSON response    return json.loads(response.choices[0].message.content)

@weave.op()
def calculate_long_tail_stats(scores):
    if not scores:
        return None    aspects = ['relevance', 'technical_quality', 'conciseness']
    stats = {}
    for aspect in aspects:
        try:
            # Handle different input formats (list of lists or list of dicts)            if isinstance(scores[0], list):
                flattened_scores = [score[aspect]['score'] for sublist in scores for score in sublist]
            elif isinstance(scores[0], dict):
                flattened_scores = [score[aspect]['score'] for score in scores]
            else:
                print(f"Unexpected format for scores: {scores}")
                return None            # Calculate statistics for each aspect            stats[aspect] = {
                "mean": np.mean(flattened_scores),
                "tail_ratio": np.mean(sorted(flattened_scores)[-max(1, int(len(flattened_scores)*0.05)):]) / np.mean(flattened_scores),
            }
        except Exception as e:
            print(f"Error calculating stats for {aspect}: {str(e)}")
            stats[aspect] = None    return stats

@weave.op()
def analyze_iteration_impact(scores):
    if len(scores) < 2:
        return {aspect: {"diminishing_returns_point": 0, "cumulative_improvement": 0} for aspect in ['relevance', 'technical_quality', 'conciseness']}
    aspects = ['relevance', 'technical_quality', 'conciseness']
    results = {}
    for aspect in aspects:
        aspect_scores = [s[aspect]['score'] for s in scores]
        improvements = [aspect_scores[i+1] - aspect_scores[i] for i in range(len(aspect_scores)-1)]
        results[aspect] = {
            "diminishing_returns_point": next((i for i, imp in enumerate(improvements) if imp <= 0), len(improvements)),
            "cumulative_improvement": sum(improvements),
        }
    return results