File size: 2,203 Bytes
e831c40
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
schema.py — Paper2Lab data contracts.

Field-agnostic contracts for PDF extraction across ML, NLP, CV, biomedical,
physics, education, social-science, economics, and interdisciplinary papers.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any, Dict, List, Optional


KNOWN_ROLES = frozenset({
    "front_matter",
    "abstract",
    "keywords",
    "introduction",
    "related_work",
    "background",
    "theory",
    "methodology",
    "experiments",
    "results",
    "discussion",
    "limitations",
    "future_work",
    "conclusion",
    "references",
    "appendix",
    "boilerplate",
    "other",
})


@dataclass
class Section:
    title: str
    text: str
    level: int = 1
    page_start: Optional[int] = None
    page_end: Optional[int] = None
    role: str = "other"
    word_count: int = 0

    def __post_init__(self) -> None:
        if self.role not in KNOWN_ROLES:
            self.role = "other"
        if not self.word_count:
            self.word_count = len((self.text or "").split())


@dataclass
class Caption:
    label: str
    caption: str
    page_number: Optional[int] = None


@dataclass
class Table:
    page_number: Optional[int]
    table_index: int
    data: Any
    engine: str = "unknown"
    caption: Optional[str] = None


@dataclass
class DocumentExtraction:
    source_pdf: str
    title: Optional[str]
    abstract: Optional[str]
    text: str
    clean_text: str
    raw_text: str = ""
    num_pages: Optional[int] = None
    sections: List[Section] = field(default_factory=list)
    all_sections: List[Section] = field(default_factory=list)
    references: List[str] = field(default_factory=list)
    references_text: str = ""
    appendix_text: str = ""
    boilerplate_text: str = ""
    captions: List[Caption] = field(default_factory=list)
    tables: List[Table] = field(default_factory=list)
    pages: List[Dict[str, Any]] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    quality: Dict[str, Any] = field(default_factory=dict)
    extraction_engine: str = "unknown"

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)