OCR-to-RAG pipeline for life science screenshots
PaddleOCR with preprocessing, scispaCy NER, figure/table detection, citation extraction, and chunked Markdown output with frontmatter. Includes watch mode and notebook reprocessing.
This commit is contained in:
7
src/ocr_pipeline/citations/__init__.py
Normal file
7
src/ocr_pipeline/citations/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
from .regex_extractor import Citation, extract_citations, format_citations_for_markdown
|
||||
|
||||
__all__ = [
|
||||
"extract_citations",
|
||||
"format_citations_for_markdown",
|
||||
"Citation",
|
||||
]
|
||||
BIN
src/ocr_pipeline/citations/__pycache__/__init__.cpython-313.pyc
Normal file
BIN
src/ocr_pipeline/citations/__pycache__/__init__.cpython-313.pyc
Normal file
Binary file not shown.
BIN
src/ocr_pipeline/citations/__pycache__/extractor.cpython-313.pyc
Normal file
BIN
src/ocr_pipeline/citations/__pycache__/extractor.cpython-313.pyc
Normal file
Binary file not shown.
Binary file not shown.
191
src/ocr_pipeline/citations/extractor.py
Normal file
191
src/ocr_pipeline/citations/extractor.py
Normal file
@@ -0,0 +1,191 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Citation:
|
||||
type: str
|
||||
value: str
|
||||
start: int
|
||||
end: int
|
||||
confidence: float
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class RegexCitationExtractor:
|
||||
DOI_PATTERN = re.compile(r"\b(?:10\.\d{4,9}/[-._;()/:A-Z0-9]+)\b", re.IGNORECASE)
|
||||
PMID_PATTERN = re.compile(r"\b(?:PMID|PubMed|pubmed)[:/\s]?(\d{1,8})\b", re.IGNORECASE)
|
||||
ARXIV_PATTERN = re.compile(
|
||||
r"\b(?:arXiv|arxiv)[:/\s]?(\d{4}\.\d{4,5}(?:v\d+)?)\b", re.IGNORECASE
|
||||
)
|
||||
PMC_PATTERN = re.compile(r"\b(?:PMC|pmc)[:/\s]?(\d{1,8})\b", re.IGNORECASE)
|
||||
ISBN_PATTERN = re.compile(
|
||||
r"\b(?:ISBN[-:\s]?)?(?:97[89][-\s]?)?\d{1,5}[-\s]?\d{1,7}[-\s]?\d{1,7}[-\s]?[\dX]\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def extract(cls, text: str) -> list[Citation]:
|
||||
citations = []
|
||||
|
||||
for match in cls.DOI_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="DOI",
|
||||
value=match.group(),
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.95,
|
||||
metadata={"url": f"https://doi.org/{match.group()}"},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.PMID_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="PMID",
|
||||
value=f"PMID:{match.group(1)}",
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.9,
|
||||
metadata={"url": f"https://pubmed.ncbi.nlm.nih.gov/{match.group(1)}/"},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.ARXIV_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="arXiv",
|
||||
value=f"arXiv:{match.group(1)}",
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.9,
|
||||
metadata={"url": f"https://arxiv.org/abs/{match.group(1)}"},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.PMC_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="PMC",
|
||||
value=f"PMC:{match.group(1)}",
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.85,
|
||||
metadata={
|
||||
"url": f"https://www.ncbi.nlm.nih.gov/pmc/articles/PMC{match.group(1)}/"
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.ISBN_PATTERN.finditer(text):
|
||||
isbn = re.sub(r"[^\dX]", "", match.group(), flags=re.IGNORECASE)
|
||||
if len(isbn) in (10, 13):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="ISBN",
|
||||
value=isbn,
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.8,
|
||||
metadata={"url": f"https://isbnsearch.org/isbn/{isbn}"},
|
||||
)
|
||||
)
|
||||
|
||||
return citations
|
||||
|
||||
|
||||
class GrobidClient:
|
||||
def __init__(self):
|
||||
self.url = settings.citations.grobid_url
|
||||
self.timeout = settings.citations.grobid_timeout
|
||||
|
||||
def extract_references(self, text: str) -> list[Citation]:
|
||||
if not settings.citations.grobid_enabled:
|
||||
return []
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{self.url}/api/processCitationList",
|
||||
data={"citations": text},
|
||||
headers={"Accept": "application/xml"},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
logger.warning("grobid_request_failed", status=response.status_code)
|
||||
return []
|
||||
|
||||
return self._parse_grobid_xml(response.text)
|
||||
except Exception as e:
|
||||
logger.warning("grobid_extraction_failed", error=str(e))
|
||||
return []
|
||||
|
||||
def _parse_grobid_xml(self, xml_text: str) -> list[Citation]:
|
||||
citations = []
|
||||
try:
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
root = ET.fromstring(xml_text)
|
||||
for cit in root.findall(".//{*}citation"):
|
||||
citation_text = ET.tostring(cit, encoding="unicode")
|
||||
doi_match = self.DOI_PATTERN.search(citation_text)
|
||||
if doi_match:
|
||||
citations.append(
|
||||
Citation(
|
||||
type="DOI",
|
||||
value=doi_match.group(),
|
||||
start=0,
|
||||
end=len(citation_text),
|
||||
confidence=0.99,
|
||||
metadata={"source": "grobid", "full_citation": citation_text},
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return citations
|
||||
|
||||
|
||||
def extract_citations(text: str) -> list[Citation]:
|
||||
citations = []
|
||||
|
||||
if settings.citations.regex_enabled:
|
||||
citations.extend(RegexCitationExtractor.extract(text))
|
||||
|
||||
if settings.citations.grobid_enabled:
|
||||
grobid = GrobidClient()
|
||||
citations.extend(grobid.extract_references(text))
|
||||
|
||||
seen = set()
|
||||
unique = []
|
||||
for cit in citations:
|
||||
key = (cit.type, cit.value)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(cit)
|
||||
|
||||
return unique
|
||||
|
||||
|
||||
def format_citations_for_markdown(citations: list[Citation]) -> str:
|
||||
if not citations:
|
||||
return "No citations found."
|
||||
|
||||
lines = ["## Citations", ""]
|
||||
for cit in citations:
|
||||
url = cit.metadata.get("url", "") if cit.metadata else ""
|
||||
if url:
|
||||
lines.append(f"- [{cit.type}: {cit.value}]({url})")
|
||||
else:
|
||||
lines.append(f"- {cit.type}: {cit.value}")
|
||||
|
||||
return "\n".join(lines)
|
||||
196
src/ocr_pipeline/citations/regex_extractor.py
Normal file
196
src/ocr_pipeline/citations/regex_extractor.py
Normal file
@@ -0,0 +1,196 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Citation:
|
||||
raw: str
|
||||
type: str
|
||||
identifier: str
|
||||
confidence: float
|
||||
context: str = ""
|
||||
|
||||
|
||||
class RegexCitationExtractor:
|
||||
DOI_PATTERN = re.compile(
|
||||
r"\b(?:10\.\d{4,9}/[-._;()/:A-Z0-9]+)\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
PMID_PATTERN = re.compile(
|
||||
r"\b(?:PMID|PubMed\s*ID|PubMed)\s*[:\-]?\s*(\d{8})\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
ARXIV_PATTERN = re.compile(
|
||||
r"\b(?:arXiv|arxiv)\s*[:\-]?\s*(\d{4}\.\d{4,5}(?:v\d+)?)\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
ISBN_PATTERN = re.compile(
|
||||
r"\b(?:ISBN[-\s]*(?:13|10)?[-\s]*)?(?:97[89][-\s]*)?\d{1,5}[-\s]*\d{1,7}[-\s]*\d{1,7}[-\s]*[\dX]\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def extract(cls, text: str) -> list[Citation]:
|
||||
citations = []
|
||||
|
||||
for match in cls.DOI_PATTERN.finditer(text):
|
||||
ctx_start = max(0, match.start() - 80)
|
||||
ctx_end = min(len(text), match.end() + 80)
|
||||
citations.append(
|
||||
Citation(
|
||||
raw=match.group(),
|
||||
type="doi",
|
||||
identifier=match.group().lower(),
|
||||
confidence=0.95,
|
||||
context=text[ctx_start:ctx_end].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.PMID_PATTERN.finditer(text):
|
||||
ctx_start = max(0, match.start() - 80)
|
||||
ctx_end = min(len(text), match.end() + 80)
|
||||
citations.append(
|
||||
Citation(
|
||||
raw=match.group(),
|
||||
type="pmid",
|
||||
identifier=f"PMID:{match.group(1)}",
|
||||
confidence=0.9,
|
||||
context=text[ctx_start:ctx_end].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.ARXIV_PATTERN.finditer(text):
|
||||
ctx_start = max(0, match.start() - 80)
|
||||
ctx_end = min(len(text), match.end() + 80)
|
||||
citations.append(
|
||||
Citation(
|
||||
raw=match.group(),
|
||||
type="arxiv",
|
||||
identifier=f"arXiv:{match.group(1)}",
|
||||
confidence=0.9,
|
||||
context=text[ctx_start:ctx_end].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.ISBN_PATTERN.finditer(text):
|
||||
ctx_start = max(0, match.start() - 80)
|
||||
ctx_end = min(len(text), match.end() + 80)
|
||||
isbn = re.sub(r"[-\s]", "", match.group())
|
||||
citations.append(
|
||||
Citation(
|
||||
raw=match.group(),
|
||||
type="isbn",
|
||||
identifier=f"ISBN:{isbn}",
|
||||
confidence=0.7,
|
||||
context=text[ctx_start:ctx_end].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
return citations
|
||||
|
||||
|
||||
class GrobidClient:
|
||||
def __init__(self):
|
||||
self.url = settings.citations.grobid_url
|
||||
self.timeout = settings.citations.grobid_timeout
|
||||
|
||||
def process_text(self, text: str) -> list[Citation]:
|
||||
if not settings.citations.grobid_enabled:
|
||||
return []
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=self.timeout) as client:
|
||||
response = client.post(
|
||||
f"{self.url}/api/processCitationText",
|
||||
data={"text": text},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return self._parse_grobid_xml(response.text)
|
||||
except Exception as e:
|
||||
logger.warning("grobid_request_failed", error=str(e))
|
||||
return []
|
||||
|
||||
def _parse_grobid_xml(self, xml_text: str) -> list[Citation]:
|
||||
try:
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
root = ET.fromstring(xml_text)
|
||||
citations = []
|
||||
|
||||
ns = {"tei": "http://www.tei-c.org/ns/1.0"}
|
||||
for cit in root.findall(".//tei:biblStruct", ns):
|
||||
doi_elem = cit.find(".//tei:idno[@type='DOI']", ns)
|
||||
pmid_elem = cit.find(".//tei:idno[@type='PMID']", ns)
|
||||
|
||||
if doi_elem is not None and doi_elem.text:
|
||||
citations.append(
|
||||
Citation(
|
||||
raw=doi_elem.text,
|
||||
type="doi",
|
||||
identifier=doi_elem.text.lower(),
|
||||
confidence=0.99,
|
||||
)
|
||||
)
|
||||
if pmid_elem is not None and pmid_elem.text:
|
||||
citations.append(
|
||||
Citation(
|
||||
raw=pmid_elem.text,
|
||||
type="pmid",
|
||||
identifier=f"PMID:{pmid_elem.text}",
|
||||
confidence=0.99,
|
||||
)
|
||||
)
|
||||
|
||||
return citations
|
||||
except Exception as e:
|
||||
logger.warning("grobid_xml_parse_failed", error=str(e))
|
||||
return []
|
||||
|
||||
|
||||
def extract_citations(text: str) -> list[Citation]:
|
||||
citations = []
|
||||
|
||||
if settings.citations.regex_enabled:
|
||||
citations.extend(RegexCitationExtractor.extract(text))
|
||||
|
||||
if settings.citations.grobid_enabled:
|
||||
client = GrobidClient()
|
||||
citations.extend(client.process_text(text))
|
||||
|
||||
seen = set()
|
||||
unique = []
|
||||
for cit in citations:
|
||||
key = (cit.type, cit.identifier)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(cit)
|
||||
|
||||
return unique
|
||||
|
||||
|
||||
def format_citations_for_markdown(citations: list[Citation]) -> str:
|
||||
if not citations:
|
||||
return "No citations found."
|
||||
|
||||
lines = [
|
||||
"## Citations",
|
||||
"",
|
||||
"| Type | Identifier | Context |",
|
||||
"|------|------------|---------|",
|
||||
]
|
||||
for cit in citations:
|
||||
context = (
|
||||
cit.context[:80].replace("|", "\\|") + "..." if len(cit.context) > 80 else cit.context
|
||||
)
|
||||
lines.append(f"| {cit.type.upper()} | `{cit.identifier}` | {context} |")
|
||||
|
||||
return "\n".join(lines)
|
||||
Reference in New Issue
Block a user