localRAG/grobid_test.py

"""
Testlauf: PDFs aus data_test/ durch GROBID schicken und das Ergebnis begutachten.

Voraussetzung: GROBID-Container läuft (z. B. auf http://localhost:8070).
Schreibt pro PDF eine TEI-XML-Datei und eine lesbare Textfassung nach grobid_out/
und gibt eine kurze Übersicht aus, was GROBID als Fließtext, Fußnote und
Literaturangabe erkannt hat.

Nur Standardbibliothek, kein zusätzliches Paket nötig.
"""

import sys
import uuid
import urllib.request
import xml.etree.ElementTree as ET
from pathlib import Path

GROBID_URL = "http://localhost:8070"
INPUT_DIR = Path("data_test")
OUTPUT_DIR = Path("grobid_out") / INPUT_DIR.name

TEI = {"tei": "http://www.tei-c.org/ns/1.0"}


def process_pdf(pdf_path: Path) -> bytes:
    """Schickt ein PDF an processFulltextDocument und liefert TEI-XML zurück."""
    boundary = uuid.uuid4().hex
    fields = [
        ("consolidateHeader", "0"),     # keine Online-Abfrage (CrossRef) für Metadaten
        ("consolidateCitations", "0"),
        ("segmentSentences", "0"),
        ("includeRawCitations", "1"),   # Literaturangaben zusätzlich im Originalwortlaut
    ]
    # Seitenkoordinaten an diese Elemente hängen (Format "seite,x,y,b,h;..."),
    # damit sich jedem Absatz eine Seitenzahl zuordnen lässt.
    fields += [("teiCoordinates", el) for el in ("head", "p", "note", "figure", "biblStruct")]
    body = b""
    for name, value in fields:
        body += (f"--{boundary}\r\nContent-Disposition: form-data; name=\"{name}\"\r\n\r\n"
                 f"{value}\r\n").encode()
    body += (f"--{boundary}\r\nContent-Disposition: form-data; name=\"input\"; "
             f"filename=\"{pdf_path.name}\"\r\nContent-Type: application/pdf\r\n\r\n").encode()
    body += pdf_path.read_bytes() + f"\r\n--{boundary}--\r\n".encode()

    req = urllib.request.Request(
        f"{GROBID_URL}/api/processFulltextDocument",
        data=body,
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
    )
    # ganze Monografien brauchen deutlich länger als Zeitschriftenartikel
    with urllib.request.urlopen(req, timeout=1800) as resp:
        return resp.read()


def text_of(el) -> str:
    return " ".join("".join(el.itertext()).split())


def summarize(tei_bytes: bytes) -> tuple[dict, str]:
    """Zählt die erkannten Elemente und baut eine lesbare Textfassung."""
    root = ET.fromstring(tei_bytes)
    title = root.find(".//tei:titleStmt/tei:title", TEI)
    authors = list(dict.fromkeys(  # GROBID listet Autoren gelegentlich doppelt
        " ".join(s.strip() for s in p.itertext() if s.strip())
        for p in root.findall(".//tei:sourceDesc//tei:author/tei:persName", TEI)))
    abstract = root.find(".//tei:profileDesc/tei:abstract", TEI)

    body = root.find(".//tei:text/tei:body", TEI)
    divs = body.findall("tei:div", TEI) if body is not None else []
    paras = body.findall(".//tei:p", TEI) if body is not None else []
    notes = root.findall(".//tei:note[@place='foot']", TEI)
    refs = root.findall(".//tei:back//tei:listBibl/tei:biblStruct", TEI)

    stats = {
        "titel": text_of(title) if title is not None else "",
        "autoren": ", ".join(authors),
        "abschnitte": len(divs),
        "absaetze": len(paras),
        "zeichen_fliesstext": sum(len(text_of(p)) for p in paras),
        "fussnoten": len(notes),
        "literaturangaben": len(refs),
    }

    out = [f"# {stats['titel']}", f"Autoren: {stats['autoren']}", ""]
    if abstract is not None and text_of(abstract):
        out += ["## Abstract", text_of(abstract), ""]
    for div in divs:
        head = div.find("tei:head", TEI)
        if head is not None:
            out += [f"## {text_of(head)}"]
        for p in div.findall("tei:p", TEI):
            out += [text_of(p), ""]
    out += ["", "=" * 70, f"## Fußnoten ({len(notes)})"]
    for n in notes:
        out += [f"[{n.get('n', '?')}] {text_of(n)}"]
    out += ["", "=" * 70, f"## Literaturverzeichnis ({len(refs)})"]
    for r in refs:
        out += [f"- {text_of(r)}"]
    return stats, "\n".join(out)


def main():
    try:
        urllib.request.urlopen(f"{GROBID_URL}/api/isalive", timeout=5)
    except Exception as e:
        sys.exit(f"GROBID unter {GROBID_URL} nicht erreichbar: {e}")

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for pdf in sorted(INPUT_DIR.glob("*.pdf")):
        print(f"\n{pdf.name} ...", flush=True)
        tei = process_pdf(pdf)
        (OUTPUT_DIR / f"{pdf.stem}.tei.xml").write_bytes(tei)
        stats, text = summarize(tei)
        (OUTPUT_DIR / f"{pdf.stem}.txt").write_text(text, encoding="utf-8")
        for k, v in stats.items():
            print(f"  {k:<20} {v}")
    print(f"\nErgebnisse in {OUTPUT_DIR}/")


if __name__ == "__main__":
    main()