feat: Create a backend to transform PubMed XML files to DoclingDocument

Signed-off-by: lucas-morin <[email protected]>

docling/datamodel/base_models.py

@@ -31,6 +31,7 @@ class InputFormat(str, Enum):
     DOCX = "docx"
     PPTX = "pptx"
     HTML = "html"
+    XML_PUBMED = "xml_pubmed"
     IMAGE = "image"
     PDF = "pdf"
     ASCIIDOC = "asciidoc"
@@ -52,6 +53,7 @@ class OutputFormat(str, Enum):
     InputFormat.PDF: ["pdf"],
     InputFormat.MD: ["md"],
     InputFormat.HTML: ["html", "htm", "xhtml"],
+    InputFormat.XML_PUBMED: ["xml"],
     InputFormat.IMAGE: ["jpg", "jpeg", "png", "tif", "tiff", "bmp"],
     InputFormat.ASCIIDOC: ["adoc", "asciidoc", "asc"],
     InputFormat.XLSX: ["xlsx"],
@@ -68,6 +70,7 @@ class OutputFormat(str, Enum):
         "application/vnd.openxmlformats-officedocument.presentationml.presentation",
     ],
     InputFormat.HTML: ["text/html", "application/xhtml+xml"],
+    InputFormat.XML_PUBMED: ["application/xml"],
     InputFormat.IMAGE: [
         "image/png",
         "image/jpeg",

docling/datamodel/document.py

@@ -264,6 +264,11 @@ def _guess_format(self, obj: Union[Path, DocumentStream]):
                 with obj.open("rb") as f:
                     content = f.read(1024)  # Read first 1KB
 
+            # Detect PubMed XML documents
+            xml_doctype = re.search(r"<!DOCTYPE [^>]+>", content.decode("utf-8"))
+            if xml_doctype and ("/NLM//DTD JATS" in xml_doctype.group()):
+                return InputFormat.XML_PUBMED
+
         elif isinstance(obj, DocumentStream):
             content = obj.stream.read(8192)
             obj.stream.seek(0)

docling/document_converter.py

@@ -15,6 +15,7 @@
 from docling.backend.msexcel_backend import MsExcelDocumentBackend
 from docling.backend.mspowerpoint_backend import MsPowerpointDocumentBackend
 from docling.backend.msword_backend import MsWordDocumentBackend
+from docling.backend.pubmed_backend import PubMedDocumentBackend
 from docling.datamodel.base_models import (
     ConversionStatus,
     DoclingComponentType,
@@ -82,6 +83,11 @@ class HTMLFormatOption(FormatOption):
     backend: Type[AbstractDocumentBackend] = HTMLDocumentBackend
 
 
+class XMLPubMedFormatOption(FormatOption):
+    pipeline_cls: Type = SimplePipeline
+    backend: Type[AbstractDocumentBackend] = PubMedDocumentBackend
+
+
 class PdfFormatOption(FormatOption):
     pipeline_cls: Type = StandardPdfPipeline
     backend: Type[AbstractDocumentBackend] = DoclingParseV2DocumentBackend
@@ -118,6 +124,9 @@ def _get_default_option(format: InputFormat) -> FormatOption:
         InputFormat.PDF: FormatOption(
             pipeline_cls=StandardPdfPipeline, backend=DoclingParseV2DocumentBackend
         ),
+        InputFormat.XML_PUBMED: FormatOption(
+            pipeline_cls=SimplePipeline, backend=PubMedDocumentBackend
+        ),
     }
     if (options := format_to_default_options.get(format)) is not None:
         return options
@@ -162,7 +171,6 @@ def convert(
         max_num_pages: int = sys.maxsize,
         max_file_size: int = sys.maxsize,
     ) -> ConversionResult:
-
         all_res = self.convert_all(
             source=[source],
             raises_on_error=raises_on_error,