transcript_eval/methodClaude.py at main · CAHLR/transcript_eval · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import os
import base64
import csv
from PyPDF2 import PdfReader, PdfWriter
from dotenv import load_dotenv
from anthropic import Anthropic
import random
import time

load_dotenv()
# Claude client
client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))


def split_pdf(input_pdf_path, output_dir, pages_per_split=5):
    os.makedirs(output_dir, exist_ok=True)
    reader = PdfReader(input_pdf_path)
    total_pages = len(reader.pages)
    split_paths = []

    for start in range(0, total_pages, pages_per_split):
        writer = PdfWriter()
        for i in range(start, min(start + pages_per_split, total_pages)):
            writer.add_page(reader.pages[i])
        out_path = os.path.join(output_dir, f"split_{start // pages_per_split + 1}.pdf")
        with open(out_path, "wb") as f:
            writer.write(f)
        split_paths.append(out_path)

    return split_paths


def call_claude_with_pdf(pdf_path, prompt):
    with open(pdf_path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")

    message = client.messages.create(
        model="claude-3-7-sonnet-20250219",
        max_tokens=4096,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "document",
                        "source": {
                            "type": "base64",
                            "media_type": "application/pdf",
                            "data": encoded,
                        },
                    },
                    {"type": "text", "text": prompt},
                ],
            }
        ],
    )

    if message.content:
        return message.content[0].text
    else:
        print("❌ Claude API returned no content.")
        return ""


def extract_tables_from_text(text_response):
    tables = []
    current = []
    for line in text_response.strip().splitlines():
        if line.strip() == "":
            if current:
                tables.append(current)
                current = []
        else:
            current.append([cell.strip() for cell in line.split(",")])
    if current:
        tables.append(current)
    return tables


def save_table_to_csv(table, output_csv_path):
    with open(output_csv_path, "w", newline="") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerows(table)


def method_claude_sdk_pipeline(pdf_path, output_folder, pages_per_split=5):
    os.makedirs(output_folder, exist_ok=True)
    split_dir = os.path.join(output_folder, "splits")
    split_files = split_pdf(pdf_path, split_dir, pages_per_split)

    csv_files = []
    table_index = 1
    for split_pd in split_files:
        print(f"Processing: {split_pd}")
        text = call_claude_with_pdf(
            split_pd, "Extract all tables as CSV. No explanation, just raw CSV."
        )
        tables = extract_tables_from_text(text)
        for table in tables:
            csv_path = os.path.join(output_folder, f"claude_table_{table_index}.csv")
            save_table_to_csv(table, csv_path)
            csv_files.append(csv_path)
            print(f"✔ Saved: {csv_path}")
            table_index += 1
        sleep_time = random.uniform(1.5, 3.5)
        print(f"⏳ Sleeping for {sleep_time:.2f} seconds before next request...")
        time.sleep(sleep_time)