Create marker/markdown_extractor.py

marker/markdown_extractor.py

@@ -0,0 +1,36 @@
+from marker.convert import convert_single_pdf
+from marker.models import load_all_models
+import tempfile
+from indexify_extractor_sdk import Content, Extractor, Feature
+
+from pydantic import BaseModel
+from typing import Optional, Literal, List, Union
+
+class MarkdownExtractorConfig(BaseModel):
+    max_pages: Optional[int] = None
+    langs: Optional[str] = None
+    batch_multiplier: Optional[int] = 2
+
+class MarkdownExtractor(Extractor):
+    name = "tensorlake/marker"
+    description = "Markdown Extractor for PDFs"
+    system_dependencies = []
+    input_mime_types = ["application/pdf"]
+
+    def __init__(self):
+        super(MarkdownExtractor, self).__init__()
+        self.model_lst = load_all_models()
+
+    def extract(self, content: Content, params: MarkdownExtractorConfig) -> List[Union[Feature, Content]]:
+        contents = []
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as inputtmpfile:
+            inputtmpfile.write(content.data)
+            inputtmpfile.flush()
+
+            full_text, images, out_meta = convert_single_pdf(inputtmpfile.name, self.model_lst, max_pages=params.max_pages, langs=params.langs, batch_multiplier=params.batch_multiplier)
+            
+            feature = Feature.metadata(value=out_meta, name="text")
+            contents.append(Content.from_text(full_text, features=[feature]))
+
+        return contents