rd-tablebench/docparse.py at master · aryn-ai/rd-tablebench · 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
import os
import json
import csv
from pathlib import Path
from dotenv import load_dotenv


from aryn_sdk.partition import partition_file_async_submit

load_dotenv()
ARYN_API_KEY = os.environ["ARYN_TEST_API_KEY"]

PROVIDER = "detr-5.0_deformable"
INPUT_DIR = Path("data/rd-tablebench/pdfs")
OUTPUT_ROOT = Path("partitioned")
DOCPARSE_URL = "https://test-api.aryn.ai/v1/document/partition"

TASKS_DIR = Path("tasks")

LANGUAGE_MAP = {
    "en": "english",
    "ru": "russian",
    "de": "german",
    "es": "spanish",
    "fr": "french",
    "ja": "japanese",
    "zh": "chinese",
    "it": "italian",
    "pt": "portuguese",
    "uk": "ukrainian",
    "nl": "dutch",
    "cs": "czech",
    "pl": "polish",
    "sv": "swedish",
    "ro": "romanian",
}


def load_score_map():
    with open("score_map.json") as f:
        return json.load(f)


def load_tasks(tasks_file_path):
    try:
        with open(tasks_file_path, "r") as f:
            reader = csv.reader(f)
            next(reader)  # Skip header
            return {row[0] for row in reader}
    except FileNotFoundError:
        tasks_file_path.parent.mkdir(parents=True, exist_ok=True)
        with open(tasks_file_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["file_path", "task_id"])
        return set()


def is_partitioned(output_path):
    if not output_path.exists():
        return False

    with open(output_path, "r") as f:
        data = json.load(f)
        return "elements" in data


def process_pdf(pdf_file, pdf_path, output_path, language, tasks_file):
    with open(pdf_path, "rb") as f:
        response = partition_file_async_submit(
            f,
            extract_table_structure=True,
            ocr_language=language,
            docparse_url=DOCPARSE_URL,
            aryn_api_key=ARYN_API_KEY,
        )

        with open(tasks_file, "a", newline="") as task_file:
            writer = csv.writer(task_file)
            writer.writerow([pdf_file, response["task_id"]])


def main():
    output_dir = OUTPUT_ROOT / PROVIDER
    output_dir.mkdir(parents=True, exist_ok=True)

    score_map = load_score_map()
    tasks_file = Path(TASKS_DIR / f"tasks_{PROVIDER}.csv")
    sent_files = load_tasks(tasks_file)

    for pdf_path in INPUT_DIR.glob("*.pdf"):
        output_path = output_dir / f"{pdf_path.name}.json"
        print(f"Processing {pdf_path.name}")

        if output_path.exists() and is_partitioned(output_path):
            print(f"Skipping {pdf_path.name} because it already exists")
            continue
        if pdf_path.name in sent_files:
            print(f"Skipping {pdf_path.name} because it has already been sent")
            continue

        language = LANGUAGE_MAP[score_map[pdf_path.name]["language"]]
        process_pdf(pdf_path.name, pdf_path, output_path, language, tasks_file)


if __name__ == "__main__":
    main()