132 lines
3.6 KiB
Python
132 lines
3.6 KiB
Python
"""GROBID integration.
|
|
|
|
Extracts structured reference lists and full-text TEI XML from PDFs by
|
|
calling a self-hosted GROBID HTTP server. No DB access, no embedding,
|
|
no network fetching of papers happens here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import xml.etree.ElementTree as ET
|
|
|
|
import httpx
|
|
|
|
from codex.config import Settings
|
|
|
|
_TEI_NS = "{http://www.tei-c.org/ns/1.0}"
|
|
|
|
|
|
def _text(element: ET.Element | None) -> str:
|
|
"""Return element text or empty string if element is None."""
|
|
if element is None:
|
|
return ""
|
|
return (element.text or "").strip()
|
|
|
|
|
|
def extract_references(
|
|
pdf_path: str,
|
|
grobid_url: str | None = None,
|
|
) -> list[dict[str, str]]:
|
|
"""Extract a structured reference list from a PDF via GROBID.
|
|
|
|
Parameters
|
|
----------
|
|
pdf_path:
|
|
Path to the PDF file on disk.
|
|
grobid_url:
|
|
Base URL of the GROBID server. Defaults to ``Settings().grobid_url``.
|
|
|
|
Returns
|
|
-------
|
|
list[dict[str, str]]
|
|
One dict per reference. All dicts contain the keys
|
|
``title``, ``authors``, ``year``, ``doi``, ``arxiv_id``
|
|
(missing values are empty strings).
|
|
"""
|
|
if grobid_url is None:
|
|
grobid_url = Settings().grobid_url
|
|
|
|
with open(pdf_path, "rb") as fh, httpx.Client(timeout=60.0) as client:
|
|
response = client.post(
|
|
f"{grobid_url}/api/processReferences",
|
|
files={"input": (pdf_path, fh, "application/pdf")},
|
|
)
|
|
response.raise_for_status()
|
|
|
|
root = ET.fromstring(response.text)
|
|
results: list[dict[str, str]] = []
|
|
|
|
for bib in root.iter(f"{_TEI_NS}biblStruct"):
|
|
# Title
|
|
title_el = bib.find(f".//{_TEI_NS}title[@level='a']")
|
|
title = _text(title_el)
|
|
|
|
# Authors: collect all persName elements
|
|
authors_parts: list[str] = []
|
|
for person in bib.iter(f"{_TEI_NS}persName"):
|
|
forename_el = person.find(f"{_TEI_NS}forename")
|
|
surname_el = person.find(f"{_TEI_NS}surname")
|
|
forename = _text(forename_el)
|
|
surname = _text(surname_el)
|
|
full = " ".join(p for p in (forename, surname) if p)
|
|
if full:
|
|
authors_parts.append(full)
|
|
authors = "; ".join(authors_parts)
|
|
|
|
# Year
|
|
date_el = bib.find(f".//{_TEI_NS}date[@type='published']")
|
|
year = ""
|
|
if date_el is not None:
|
|
when = date_el.get("when", "")
|
|
year = when[:4] if when else ""
|
|
|
|
# DOI
|
|
doi_el = bib.find(f".//{_TEI_NS}idno[@type='DOI']")
|
|
doi = _text(doi_el)
|
|
|
|
# arXiv ID
|
|
arxiv_el = bib.find(f".//{_TEI_NS}idno[@type='arxiv']")
|
|
arxiv_id = _text(arxiv_el)
|
|
|
|
results.append(
|
|
{
|
|
"title": title,
|
|
"authors": authors,
|
|
"year": year,
|
|
"doi": doi,
|
|
"arxiv_id": arxiv_id,
|
|
}
|
|
)
|
|
|
|
return results
|
|
|
|
|
|
def extract_structure(
|
|
pdf_path: str,
|
|
grobid_url: str | None = None,
|
|
) -> str:
|
|
"""Extract full-text TEI XML from a PDF via GROBID.
|
|
|
|
Parameters
|
|
----------
|
|
pdf_path:
|
|
Path to the PDF file on disk.
|
|
grobid_url:
|
|
Base URL of the GROBID server. Defaults to ``Settings().grobid_url``.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Raw TEI XML response text from GROBID.
|
|
"""
|
|
if grobid_url is None:
|
|
grobid_url = Settings().grobid_url
|
|
|
|
with open(pdf_path, "rb") as fh, httpx.Client(timeout=120.0) as client:
|
|
response = client.post(
|
|
f"{grobid_url}/api/processFulltextDocument",
|
|
files={"input": (pdf_path, fh, "application/pdf")},
|
|
)
|
|
response.raise_for_status()
|
|
return response.text
|