improve the OCR pipeline processing and outputs formatting
This commit is contained in:
@@ -1,191 +1,16 @@
|
||||
from __future__ import annotations
|
||||
"""Backward-compatible citation exports.
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
The canonical implementation lives in ``regex_extractor``. Keeping this module as a
|
||||
thin re-export avoids two incompatible Citation models and removes an undeclared
|
||||
``requests`` dependency from the runtime path.
|
||||
"""
|
||||
|
||||
import requests
|
||||
from .regex_extractor import (
|
||||
Citation,
|
||||
GrobidClient,
|
||||
RegexCitationExtractor,
|
||||
extract_citations,
|
||||
format_citations_for_markdown,
|
||||
)
|
||||
|
||||
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)
|
||||
__all__ = ["Citation", "GrobidClient", "RegexCitationExtractor", "extract_citations", "format_citations_for_markdown"]
|
||||
|
||||
@@ -2,8 +2,8 @@ from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
from urllib.parse import urlencode
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
@@ -26,7 +26,7 @@ class RegexCitationExtractor:
|
||||
re.IGNORECASE,
|
||||
)
|
||||
PMID_PATTERN = re.compile(
|
||||
r"\b(?:PMID|PubMed\s*ID|PubMed)\s*[:\-]?\s*(\d{8})\b",
|
||||
r"\b(?:PMID|PubMed\s*ID|PubMed)\s*[:\-]?\s*(\d{1,8})\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
ARXIV_PATTERN = re.compile(
|
||||
@@ -49,7 +49,7 @@ class RegexCitationExtractor:
|
||||
Citation(
|
||||
raw=match.group(),
|
||||
type="doi",
|
||||
identifier=match.group().lower(),
|
||||
identifier=match.group().rstrip(".,;:)]}").lower(),
|
||||
confidence=0.95,
|
||||
context=text[ctx_start:ctx_end].strip(),
|
||||
)
|
||||
@@ -108,13 +108,14 @@ class GrobidClient:
|
||||
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)
|
||||
payload = urlencode({"text": text}).encode("utf-8")
|
||||
request = Request(
|
||||
f"{self.url}/api/processCitationText",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||||
)
|
||||
with urlopen(request, timeout=self.timeout) as response:
|
||||
return self._parse_grobid_xml(response.read().decode("utf-8"))
|
||||
except Exception as e:
|
||||
logger.warning("grobid_request_failed", error=str(e))
|
||||
return []
|
||||
|
||||
Reference in New Issue
Block a user