945 lines
40 KiB
Python
945 lines
40 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import functools
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import traceback
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
import torch
|
|
from flask import Flask, jsonify, request, Response
|
|
import PIL.Image
|
|
|
|
from marker.config.parser import ConfigParser
|
|
from marker.converters.pdf import PdfConverter
|
|
from marker.models import create_model_dict
|
|
from marker.output import text_from_rendered
|
|
from marker.providers.registry import load_extensions
|
|
from marker.settings import settings as marker_settings
|
|
|
|
import request_store
|
|
import deepseek_ocr as dsocr
|
|
|
|
# ── Env configuration ────────────────────────────────────────────────
|
|
API_KEY = os.environ.get("API_KEY", "")
|
|
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://localhost:11435")
|
|
DEESEEK_OCR_MODEL = os.environ.get("DEESEEK_OCR_MODEL", "deepseek-ocr")
|
|
OCR_BACKEND = os.environ.get("OCR_BACKEND", "marker") # "marker" | "deepseek"
|
|
AMD_COMPUTE = os.environ.get("AMD_COMPUTE", "false").lower() in ("true", "1", "yes")
|
|
TORCH_DEVICE = os.environ.get("TORCH_DEVICE", "")
|
|
MODEL_DTYPE = os.environ.get("MODEL_DTYPE", "float32")
|
|
|
|
DEFAULT_LLM_SERVICE = os.environ.get("LLM_SERVICE", "marker.services.ollama.OllamaService")
|
|
DEFAULT_USE_LLM = os.environ.get("USE_LLM", "false").lower() in ("true", "1", "yes")
|
|
DEFAULT_OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "")
|
|
DEFAULT_OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
|
|
DEFAULT_OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "")
|
|
|
|
SUPPORTED_EXTENSIONS = set()
|
|
for provider_type in ("image", "pdf", "epub", "doc", "xls", "ppt"):
|
|
SUPPORTED_EXTENSIONS.update(load_extensions(provider_type))
|
|
SUPPORTED_DISPLAY = sorted({ext.lstrip(".") for ext in SUPPORTED_EXTENSIONS})
|
|
|
|
_marker_dict: Optional[Dict[str, Any]] = None
|
|
_lock = threading.Lock()
|
|
|
|
BASE_URL = os.environ.get("SELF_URL", "http://localhost:8001")
|
|
|
|
|
|
def _configure_env():
|
|
if AMD_COMPUTE and not TORCH_DEVICE:
|
|
os.environ["TORCH_DEVICE"] = "cuda"
|
|
os.environ["TORCH_DEVICE_MODEL"] = "cuda"
|
|
if MODEL_DTYPE == "bfloat16":
|
|
os.environ["MODEL_DTYPE"] = "bfloat16"
|
|
if torch_device_override := os.environ.get("TORCH_DEVICE"):
|
|
os.environ["TORCH_DEVICE"] = torch_device_override
|
|
os.environ["TORCH_DEVICE_MODEL"] = torch_device_override
|
|
|
|
|
|
def get_model_dict() -> Dict[str, Any]:
|
|
global _marker_dict
|
|
if _marker_dict is None:
|
|
with _lock:
|
|
if _marker_dict is None:
|
|
_marker_dict = create_model_dict()
|
|
return _marker_dict
|
|
|
|
|
|
# ── Auth decorator ───────────────────────────────────────────────────
|
|
def require_api_key(f):
|
|
@functools.wraps(f)
|
|
def wrapper(*args, **kwargs):
|
|
if API_KEY:
|
|
header_key = request.headers.get("X-API-Key", "")
|
|
auth_header = request.headers.get("Authorization", "")
|
|
if auth_header.startswith("Bearer "):
|
|
header_key = header_key or auth_header[7:]
|
|
if header_key != API_KEY:
|
|
return jsonify({"success": False, "error": "Invalid or missing API key"}), 403
|
|
return f(*args, **kwargs)
|
|
return wrapper
|
|
|
|
|
|
# ── Option parsing ───────────────────────────────────────────────────
|
|
def parse_options(src: dict) -> dict:
|
|
bool_fields = [
|
|
"paginate_output", "force_ocr", "disable_image_extraction",
|
|
"use_llm", "redo_inline_math", "strip_existing_ocr", "debug",
|
|
"paginate", "skip_cache", "save_checkpoint",
|
|
"disable_image_captions", "fence_synthetic_captions",
|
|
"token_efficient_markdown", "add_block_ids",
|
|
"include_markdown_in_chunks", "word_bboxes",
|
|
]
|
|
str_fields = [
|
|
"page_range", "output_format", "processors", "config_json",
|
|
"converter_cls", "llm_service", "block_correction_prompt",
|
|
"langs", "mode", "file_url", "additional_config",
|
|
"extras", "pipeline_id", "webhook_url", "processing_location",
|
|
"page_schema", "segmentation_schema", "model_override_settings",
|
|
]
|
|
int_fields = ["max_pages", "workflowstepdata_id"]
|
|
opts = {}
|
|
for k in bool_fields:
|
|
v = src.get(k)
|
|
if isinstance(v, str):
|
|
opts[k] = v.lower() in ("true", "1", "yes")
|
|
elif v is not None:
|
|
opts[k] = bool(v)
|
|
for k in str_fields:
|
|
v = src.get(k)
|
|
if v is not None and v != "":
|
|
opts[k] = str(v)
|
|
for k in int_fields:
|
|
v = src.get(k)
|
|
if v is not None:
|
|
opts[k] = int(v)
|
|
return opts
|
|
|
|
|
|
def build_options(**extra: Any) -> Dict[str, Any]:
|
|
opts: Dict[str, Any] = {}
|
|
for k, v in extra.items():
|
|
if v is not None:
|
|
opts[k] = v
|
|
opts.setdefault("output_format", "markdown")
|
|
opts.setdefault("force_ocr", False)
|
|
opts.setdefault("paginate_output", False)
|
|
opts.setdefault("page_range", None)
|
|
opts.setdefault("disable_multiprocessing", True)
|
|
opts.setdefault("disable_image_extraction", False)
|
|
opts.setdefault("output_dir", marker_settings.OUTPUT_DIR)
|
|
opts.setdefault("llm_service", DEFAULT_LLM_SERVICE)
|
|
if DEFAULT_OPENAI_BASE_URL:
|
|
opts.setdefault("openai_base_url", DEFAULT_OPENAI_BASE_URL)
|
|
if DEFAULT_OPENAI_API_KEY:
|
|
opts.setdefault("openai_api_key", DEFAULT_OPENAI_API_KEY)
|
|
if DEFAULT_OPENAI_MODEL:
|
|
opts.setdefault("openai_model", DEFAULT_OPENAI_MODEL)
|
|
if OLLAMA_HOST:
|
|
opts.setdefault("ollama_base_url", OLLAMA_HOST)
|
|
if DEESEEK_OCR_MODEL:
|
|
opts.setdefault("ollama_model", DEESEEK_OCR_MODEL)
|
|
config_parser = ConfigParser(opts)
|
|
return config_parser.generate_config_dict()
|
|
|
|
|
|
# ── Conversion core ──────────────────────────────────────────────────
|
|
def convert_file_bytes(raw_bytes: bytes, filename: str, **opts: Any) -> dict:
|
|
tmp = tempfile.NamedTemporaryFile(delete=False, suffix="." + filename.rsplit(".", 1)[-1] if "." in filename else "")
|
|
try:
|
|
tmp.write(raw_bytes)
|
|
tmp.close()
|
|
filepath = tmp.name
|
|
config_dict = build_options(**opts)
|
|
config_dict["disable_tqdm"] = True
|
|
model_dict = get_model_dict()
|
|
parsed_opts = {k: v for k, v in opts.items() if v is not None}
|
|
for key in ("output_format", "llm_service", "openai_base_url", "openai_api_key", "openai_model"):
|
|
if key not in parsed_opts and key in config_dict:
|
|
parsed_opts[key] = config_dict[key]
|
|
parsed = ConfigParser(parsed_opts)
|
|
converter = PdfConverter(
|
|
config=config_dict,
|
|
artifact_dict=model_dict,
|
|
processor_list=parsed.get_processors(),
|
|
renderer=parsed.get_renderer(),
|
|
llm_service=parsed.get_llm_service(),
|
|
)
|
|
rendered = converter(filepath)
|
|
text, _, images = text_from_rendered(rendered)
|
|
except Exception as exc:
|
|
traceback.print_exc()
|
|
return {"success": False, "error": str(exc)}
|
|
finally:
|
|
tmp.close()
|
|
if os.path.exists(tmp.name):
|
|
os.remove(tmp.name)
|
|
encoded_images: Dict[str, str] = {}
|
|
img_fmt = marker_settings.OUTPUT_IMAGE_FORMAT
|
|
for k, img in images.items():
|
|
buf = io.BytesIO()
|
|
img.save(buf, format=img_fmt)
|
|
encoded_images[k] = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
return {
|
|
"format": opts.get("output_format", "markdown"),
|
|
"output": text,
|
|
"images_b64": encoded_images,
|
|
"metadata": rendered.metadata if hasattr(rendered, 'metadata') else {},
|
|
"success": True,
|
|
}
|
|
|
|
|
|
def convert_with_ocr_backend(raw_bytes: bytes, filename: str, **opts: Any) -> dict:
|
|
if OCR_BACKEND == "deepseek":
|
|
page_range = opts.get("page_range")
|
|
max_pages = opts.get("max_pages")
|
|
ocr_result = dsocr.ocr_pdf(raw_bytes, max_pages=max_pages, page_range=page_range)
|
|
return {
|
|
"format": opts.get("output_format", "markdown"),
|
|
"output": ocr_result.get("text", ""),
|
|
"images_b64": {},
|
|
"metadata": {"page_count": ocr_result.get("page_count", 0), "ocr_backend": "deepseek"},
|
|
"success": ocr_result.get("success", False),
|
|
"page_count": ocr_result.get("page_count", 0),
|
|
}
|
|
return convert_file_bytes(raw_bytes, filename, **opts)
|
|
|
|
|
|
# ── Background worker ────────────────────────────────────────────────
|
|
# The job queue (submit_job, _job_queue, _job_cond) lives in request_store.py
|
|
# so both the Flask app and gunicorn's post_fork hook can access it.
|
|
|
|
|
|
# ── Extraction / Segmentation helpers ────────────────────────────────
|
|
def _apply_extraction_schema(markdown_text: str, schema: dict) -> dict:
|
|
result = {}
|
|
for field_name, field_def in schema.items():
|
|
if isinstance(field_def, dict):
|
|
desc = field_def.get("description", "")
|
|
ftype = field_def.get("type", "string")
|
|
result[field_name] = _extract_field(markdown_text, field_name, desc, ftype)
|
|
else:
|
|
result[field_name] = _extract_field(markdown_text, field_name, str(field_def), "string")
|
|
return result
|
|
|
|
|
|
def _extract_field(text: str, name: str, desc: str, ftype: str) -> Any:
|
|
import re
|
|
patterns = [
|
|
re.compile(rf"(?:{re.escape(name)}\s*:?\s*(.*?))(?:\n\n|\Z)", re.IGNORECASE | re.DOTALL),
|
|
re.compile(rf"(?:{re.escape(desc)}\s*:?\s*(.*?))(?:\n\n|\Z)", re.IGNORECASE | re.DOTALL),
|
|
]
|
|
for pat in patterns:
|
|
m = pat.search(text)
|
|
if m:
|
|
val = m.group(1).strip()
|
|
if ftype == "number":
|
|
try:
|
|
return float(val)
|
|
except ValueError:
|
|
return val
|
|
return val
|
|
return None
|
|
|
|
|
|
def _apply_segmentation_schema(markdown_text: str, schema: dict) -> list[dict]:
|
|
segments = []
|
|
for seg_name, seg_desc in schema.items():
|
|
if isinstance(seg_desc, str):
|
|
desc_text = seg_desc
|
|
elif isinstance(seg_desc, dict):
|
|
desc_text = seg_desc.get("description", seg_name)
|
|
else:
|
|
desc_text = str(seg_desc)
|
|
import re
|
|
pattern = re.compile(rf"(?:{re.escape(desc_text)})(.*?)(?=(?:{'|'.join(re.escape(str(v)) if isinstance(v, str) else re.escape(str(v.get('description', ''))) for v in schema.values())})|\Z)", re.IGNORECASE | re.DOTALL)
|
|
m = pattern.search(markdown_text)
|
|
if m:
|
|
segments.append({"name": seg_name, "description": desc_text, "content": m.group(1).strip()})
|
|
return segments
|
|
|
|
|
|
def _extract_tables(conv_result: dict) -> list[dict]:
|
|
output = conv_result.get("output", "")
|
|
if conv_result.get("format") == "json":
|
|
try:
|
|
data = json.loads(output) if isinstance(output, str) else output
|
|
if isinstance(data, dict):
|
|
return data.get("tables", [])
|
|
if isinstance(data, list):
|
|
return [item for item in data if "table" in str(item).lower() or "rows" in str(item)]
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
import re
|
|
tables = []
|
|
table_pattern = re.compile(r"\|(.+)\|[\s\S]*?(?=\n\n|\Z)", re.MULTILINE)
|
|
for i, m in enumerate(table_pattern.finditer(output)):
|
|
tables.append({"index": i, "table": m.group(0).strip()})
|
|
return tables
|
|
|
|
|
|
# ── File management ──────────────────────────────────────────────────
|
|
def _file_storage_path() -> Path:
|
|
p = Path(os.environ.get("UPLOAD_DIR", "/app/uploads"))
|
|
p.mkdir(parents=True, exist_ok=True)
|
|
return p
|
|
|
|
|
|
def _get_file_metadata(file_id: str) -> dict | None:
|
|
meta_path = _file_storage_path() / file_id / "metadata.json"
|
|
if not meta_path.exists():
|
|
return None
|
|
return json.loads(meta_path.read_text())
|
|
|
|
|
|
def _submit_and_poll(endpoint: str, raw_bytes: bytes, filename: str, form: dict | None = None, opts: dict | None = None) -> Response:
|
|
if opts is None:
|
|
opts = parse_options(form or {})
|
|
request_id = request_store.submit_job(endpoint, raw_bytes, filename, opts)
|
|
check_url = f"{BASE_URL}/api/v1/{endpoint}/{request_id}"
|
|
return jsonify({
|
|
"success": True,
|
|
"request_id": request_id,
|
|
"request_check_url": check_url,
|
|
})
|
|
|
|
|
|
def _poll_result(endpoint: str, request_id: str) -> Response:
|
|
entry = request_store.get_request(request_id)
|
|
if entry is None:
|
|
return jsonify({"success": False, "error": "Request not found"}), 404
|
|
if entry["status"] == "processing":
|
|
return jsonify({"status": "processing"})
|
|
result = request_store.get_result(request_id)
|
|
if result is None:
|
|
return jsonify({"status": "complete", "success": False, "error": "Result not found"}), 500
|
|
fmt = result.get("format", "markdown")
|
|
payload = {
|
|
"status": "complete",
|
|
"success": result.get("success", False),
|
|
"error": result.get("error"),
|
|
"output_format": fmt,
|
|
}
|
|
output_text = result.get("output", "")
|
|
if fmt in ("markdown", "html"):
|
|
payload[fmt] = output_text
|
|
elif fmt == "json":
|
|
try:
|
|
payload["json"] = json.loads(output_text) if output_text else None
|
|
except (json.JSONDecodeError, TypeError):
|
|
payload["json"] = output_text
|
|
elif fmt == "chunks":
|
|
payload["chunks"] = output_text
|
|
payload["images"] = result.get("images_b64", {})
|
|
payload["metadata"] = result.get("metadata", {})
|
|
payload["page_count"] = result.get("page_count")
|
|
return jsonify(payload)
|
|
|
|
|
|
# ── Flask app factory ────────────────────────────────────────────────
|
|
def create_app() -> Flask:
|
|
app = Flask(__name__)
|
|
|
|
# Start cleanup thread for expired results
|
|
request_store.start_cleanup_thread()
|
|
|
|
HTML_DOCS = r"""<html><head><title>marker-api</title>
|
|
<style>body{font-family:sans-serif;max-width:960px;margin:40px auto;padding:0 20px}
|
|
pre{background:#f4f4f4;padding:12px;border-radius:4px;overflow-x:auto}
|
|
table{border-collapse:collapse;width:100%;margin:16px 0}th,td{border:1px solid #ccc;padding:8px;text-align:left}
|
|
th{background:#f0f0f0}
|
|
.alert{padding:12px 16px;border-radius:4px;margin:16px 0}
|
|
.alert-deprecated{background:#fff3cd;border:1px solid #ffc107}
|
|
.alert-info{background:#d1ecf1;border:1px solid #bee5eb}
|
|
code{background:#f4f4f4;padding:2px 6px;border-radius:3px}
|
|
</style></head><body>
|
|
<h1>marker-api</h1>
|
|
<p>Datalab-compatible document conversion API. Convert PDFs, images, documents to Markdown/HTML/JSON.</p>
|
|
|
|
<div class="alert alert-info"><strong>OCR Backend:</strong> <code>{ocr_backend}</code></div>
|
|
|
|
<h2>Datalab-compatible Endpoints</h2>
|
|
<table>
|
|
<tr><th>Method</th><th>Path</th><th>Description</th></tr>
|
|
<tr><td>POST</td><td>/api/v1/convert</td><td>Convert document (async, returns request_id + check_url)</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/convert/<id></td><td>Poll conversion result</td></tr>
|
|
<tr><td>POST</td><td>/api/v1/extract</td><td>Structured extraction with JSON schema</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/extract/<id></td><td>Poll extraction result</td></tr>
|
|
<tr><td>POST</td><td>/api/v1/segment</td><td>Document segmentation with JSON schema</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/segment/<id></td><td>Poll segmentation result</td></tr>
|
|
<tr><td>POST</td><td>/api/v1/ocr</td><td>[DEPRECATED] OCR-only extraction</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/ocr/<id></td><td>Poll OCR result</td></tr>
|
|
<tr><td>POST</td><td>/api/v1/table_rec</td><td>[DEPRECATED] Table recognition</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/table_rec/<id></td><td>Poll table recognition result</td></tr>
|
|
<tr><td>POST</td><td>/api/v1/marker</td><td>[DEPRECATED] Sync conversion (returns directly)</td></tr>
|
|
<tr><td>POST</td><td>/api/v1/create-document</td><td>Create DOCX from markdown</td></tr>
|
|
</table>
|
|
|
|
<h3>File Management</h3>
|
|
<table>
|
|
<tr><th>Method</th><th>Path</th><th>Description</th></tr>
|
|
<tr><td>POST</td><td>/api/v1/files/upload</td><td>Request file upload URL</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/files</td><td>List uploaded files</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/files/<id></td><td>Get file metadata</td></tr>
|
|
<tr><td>GET</td><td>/api/v1/files/<id>/download</td><td>Get file download URL</td></tr>
|
|
<tr><td>DELETE</td><td>/api/v1/files/<id></td><td>Delete file</td></tr>
|
|
</table>
|
|
|
|
<h3>Legacy Endpoints</h3>
|
|
<table>
|
|
<tr><th>Method</th><th>Path</th><th>Description</th></tr>
|
|
<tr><td>POST</td><td>/marker</td><td>Sync conversion (original)</td></tr>
|
|
<tr><td>POST</td><td>/v1/conversions</td><td>Sync conversion (original)</td></tr>
|
|
<tr><td>POST</td><td>/v1/files/convert</td><td>Sync conversion via base64 JSON</td></tr>
|
|
</table>
|
|
|
|
<h2>Authentication</h2>
|
|
<p>All <code>/api/v1/*</code> endpoints accept <code>X-API-Key</code> header (or <code>Authorization: Bearer</code>).</p>
|
|
|
|
<h2>POST /api/v1/convert</h2>
|
|
<pre><code>curl -X POST http://localhost:8001/api/v1/convert \\
|
|
-H "X-API-Key: YOUR_KEY" \\
|
|
-F "file=@document.pdf" \\
|
|
-F "output_format=markdown" \\
|
|
-F "mode=balanced"</code></pre>
|
|
|
|
<h2>Poll for results</h2>
|
|
<pre><code>curl http://localhost:8001/api/v1/convert/REQUEST_ID \\
|
|
-H "X-API-Key: YOUR_KEY"</code></pre>
|
|
|
|
<h3>Parameters</h3>
|
|
<table>
|
|
<tr><th>Parameter</th><th>Type</th><th>Default</th><th>Description</th></tr>
|
|
<tr><td><b>file</b></td><td>file</td><td>required</td><td>Document to convert</td></tr>
|
|
<tr><td>output_format</td><td>string</td><td>markdown</td><td>markdown, html, json, chunks</td></tr>
|
|
<tr><td>mode</td><td>string</td><td>fast</td><td>fast, balanced, accurate</td></tr>
|
|
<tr><td>page_range</td><td>string</td><td>all</td><td>e.g. "0,5-10"</td></tr>
|
|
<tr><td>max_pages</td><td>int</td><td>all</td><td>Max pages to process</td></tr>
|
|
<tr><td>paginate</td><td>bool</td><td>false</td><td>Add page delimiters</td></tr>
|
|
</table>
|
|
|
|
<h3>Environment</h3>
|
|
<table>
|
|
<tr><th>Variable</th><th>Default</th><th>Description</th></tr>
|
|
<tr><td>OCR_BACKEND</td><td>marker</td><td>"marker" (marker-pdf) or "deepseek" (deepseek-ocr via ollama)</td></tr>
|
|
<tr><td>API_KEY</td><td>(none)</td><td>API key for X-API-Key auth (empty = disabled)</td></tr>
|
|
<tr><td>DEEPSEEK_OLLAMA_HOST</td><td>http://localhost:11434</td><td>Ollama host for deepseek-ocr</td></tr>
|
|
<tr><td>DEEPSEEK_OCR_MODEL</td><td>deepseek-ocr:latest</td><td>Model name for deepseek-ocr</td></tr>
|
|
<tr><td>OLLAMA_HOST</td><td>http://localhost:11435</td><td>Ollama host for LLM correction</td></tr>
|
|
<tr><td>DEESEEK_OCR_MODEL</td><td>deepseek-ocr</td><td>LLM correction model name</td></tr>
|
|
<tr><td>TORCH_DEVICE</td><td>auto</td><td>cpu, cuda, rocm</td></tr>
|
|
<tr><td>PORT</td><td>8001</td><td>Listening port</td></tr>
|
|
<tr><td>GUNICORN_WORKERS</td><td>3</td><td>Number of workers</td></tr>
|
|
</table>
|
|
<h3>Supported Formats</h3>
|
|
<p><code>{formats}</code></p>
|
|
</body></html>"""
|
|
|
|
# ── Legacy /info endpoints ───────────────────────────────────────
|
|
@app.route("/")
|
|
def docs():
|
|
body = HTML_DOCS.replace("{ocr_backend}", OCR_BACKEND).replace("{formats}", ", ".join(SUPPORTED_DISPLAY))
|
|
return body
|
|
|
|
@app.route("/health")
|
|
def health():
|
|
try:
|
|
device = TORCH_DEVICE or (marker_settings.TORCH_DEVICE_MODEL if hasattr(marker_settings, 'TORCH_DEVICE_MODEL') else 'auto')
|
|
except Exception:
|
|
device = "unknown"
|
|
return jsonify({
|
|
"status": "ok",
|
|
"ollama": OLLAMA_HOST,
|
|
"ocr_backend": OCR_BACKEND,
|
|
"amd_compute": AMD_COMPUTE,
|
|
"torch_device": device,
|
|
"supported_formats": SUPPORTED_DISPLAY,
|
|
"provider": "flask",
|
|
"default_llm_service": DEFAULT_LLM_SERVICE,
|
|
"default_use_llm": DEFAULT_USE_LLM,
|
|
"openai_base_url": DEFAULT_OPENAI_BASE_URL or None,
|
|
"openai_model": DEFAULT_OPENAI_MODEL or None,
|
|
})
|
|
|
|
# ── Existing original endpoints (unchanged behavior) ────────────
|
|
@app.route("/marker", methods=["POST"])
|
|
def convert_sync():
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
filename = file.filename or "file"
|
|
raw = file.read()
|
|
fmt = request.form.get("output_format", "markdown")
|
|
opts = parse_options(request.form)
|
|
opts["output_format"] = fmt
|
|
result = convert_with_ocr_backend(raw, filename, **opts)
|
|
if not result["success"]:
|
|
return jsonify(result), 500
|
|
if fmt == "markdown":
|
|
return Response(
|
|
result["output"],
|
|
mimetype="text/plain",
|
|
headers={"Content-Disposition": f'attachment; filename="{filename.rsplit(".", 1)[0]}.md"'},
|
|
)
|
|
return jsonify(result)
|
|
if request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" not in data or not data.get("filename"):
|
|
return jsonify({"error": "JSON body must include 'file_b64' and 'filename'"}), 400
|
|
raw = base64.b64decode(data["file_b64"])
|
|
fmt = data.get("output_format", "markdown")
|
|
opts = parse_options(data)
|
|
opts["output_format"] = fmt
|
|
result = convert_with_ocr_backend(raw, data["filename"], **opts)
|
|
return jsonify(result)
|
|
return jsonify({"error": "No file provided. Use multipart/form-data or JSON with 'file_b64'."}), 400
|
|
|
|
@app.route("/v1/conversions", methods=["POST"])
|
|
def convert_async_style():
|
|
if request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" not in data or not data.get("filename"):
|
|
return jsonify({"error": "JSON body must include 'file_b64' and 'filename'"}), 400
|
|
raw = base64.b64decode(data["file_b64"])
|
|
fmt = data.get("output_format", "markdown")
|
|
opts = parse_options(data)
|
|
opts["output_format"] = fmt
|
|
result = convert_with_ocr_backend(raw, data["filename"], **opts)
|
|
if not result["success"]:
|
|
return jsonify(result), 500
|
|
return jsonify({
|
|
"id": str(uuid.uuid4()),
|
|
"filename": data["filename"],
|
|
"format": fmt,
|
|
"output": result["output"],
|
|
"success": True,
|
|
"images_b64": result.get("images_b64", {}),
|
|
"metadata": result.get("metadata", {}),
|
|
})
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
filename = file.filename or "file"
|
|
raw = file.read()
|
|
fmt = request.form.get("output_format", "markdown")
|
|
opts = parse_options(request.form)
|
|
opts["output_format"] = fmt
|
|
result = convert_with_ocr_backend(raw, filename, **opts)
|
|
if not result["success"]:
|
|
return jsonify(result), 500
|
|
return jsonify({
|
|
"id": str(uuid.uuid4()),
|
|
"filename": filename,
|
|
"format": fmt,
|
|
"output": result["output"],
|
|
"success": True,
|
|
"images_b64": result.get("images_b64", {}),
|
|
"metadata": result.get("metadata", {}),
|
|
})
|
|
return jsonify({"error": "No file provided."}), 400
|
|
|
|
@app.route("/v1/files/convert", methods=["POST"])
|
|
def convert_files():
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
filename = file.filename or "file"
|
|
raw = file.read()
|
|
fmt = request.form.get("output_format", "markdown")
|
|
opts = parse_options(request.form)
|
|
opts["output_format"] = fmt
|
|
result = convert_with_ocr_backend(raw, filename, **opts)
|
|
filename_var = filename
|
|
elif request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" not in data or not data.get("filename"):
|
|
return jsonify({"error": "JSON body must include 'file_b64' and 'filename'"}), 400
|
|
raw = base64.b64decode(data["file_b64"])
|
|
fmt = data.get("output_format", "markdown")
|
|
opts = parse_options(data)
|
|
opts["output_format"] = fmt
|
|
result = convert_with_ocr_backend(raw, data["filename"], **opts)
|
|
filename_var = data["filename"]
|
|
else:
|
|
return jsonify({"error": "No file provided."}), 400
|
|
if not result["success"]:
|
|
return jsonify(result), 500
|
|
return jsonify({
|
|
"filename": filename_var,
|
|
"format": fmt,
|
|
"output": result["output"],
|
|
"images_b64": result.get("images_b64", {}),
|
|
"metadata": result.get("metadata", {}),
|
|
})
|
|
|
|
# ── New Datalab-compatible endpoints ─────────────────────────────
|
|
|
|
# POST /api/v1/convert — async submit
|
|
@app.route("/api/v1/convert", methods=["POST"])
|
|
@require_api_key
|
|
def api_convert_submit():
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
raw = file.read()
|
|
filename = file.filename or "file"
|
|
elif request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" in data:
|
|
raw = base64.b64decode(data["file_b64"])
|
|
filename = data.get("filename", "file")
|
|
elif "file_url" in data:
|
|
import requests as req
|
|
resp = req.get(data["file_url"], timeout=120)
|
|
resp.raise_for_status()
|
|
raw = resp.content
|
|
filename = data.get("filename", data["file_url"].rsplit("/", 1)[-1] or "file")
|
|
else:
|
|
return jsonify({"error": "Provide file, file_b64, or file_url"}), 400
|
|
else:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
source = request.form if request.form else (request.get_json() if request.is_json else {})
|
|
opts = parse_options(source)
|
|
opts["output_format"] = source.get("output_format", "markdown")
|
|
return _submit_and_poll("convert", raw, filename, opts=opts)
|
|
|
|
# GET /api/v1/convert/<request_id> — poll
|
|
@app.route("/api/v1/convert/<request_id>", methods=["GET"])
|
|
@require_api_key
|
|
def api_convert_poll(request_id: str):
|
|
return _poll_result("convert", request_id)
|
|
|
|
# POST /api/v1/marker — deprecated (Datalab compat: async submit-and-poll)
|
|
@app.route("/api/v1/marker", methods=["POST"])
|
|
@require_api_key
|
|
def api_marker_deprecated():
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
raw = file.read()
|
|
filename = file.filename or "file"
|
|
elif request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" in data:
|
|
raw = base64.b64decode(data["file_b64"])
|
|
filename = data.get("filename", "file")
|
|
else:
|
|
return jsonify({"error": "Provide file or file_b64"}), 400
|
|
else:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
source = request.form if request.form else (request.get_json() if request.is_json else {})
|
|
opts = parse_options(source)
|
|
opts.setdefault("output_format", "markdown")
|
|
resp = _submit_and_poll("convert", raw, filename, opts=opts)
|
|
if isinstance(resp, Response):
|
|
resp.headers["Warning"] = "299 marker-api \"POST /api/v1/marker is deprecated, use /api/v1/convert\""
|
|
return resp
|
|
|
|
# POST /api/v1/extract — structured extraction
|
|
@app.route("/api/v1/extract", methods=["POST"])
|
|
@require_api_key
|
|
def api_extract_submit():
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
raw = file.read()
|
|
filename = file.filename or "file"
|
|
elif request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" in data:
|
|
raw = base64.b64decode(data["file_b64"])
|
|
filename = data.get("filename", "file")
|
|
else:
|
|
return jsonify({"error": "Provide file or file_b64"}), 400
|
|
else:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
source = request.form if request.form else (request.get_json() if request.is_json else {})
|
|
opts = parse_options(source)
|
|
opts.setdefault("output_format", "markdown")
|
|
return _submit_and_poll("extract", raw, filename, opts=opts)
|
|
|
|
# GET /api/v1/extract/<request_id>
|
|
@app.route("/api/v1/extract/<request_id>", methods=["GET"])
|
|
@require_api_key
|
|
def api_extract_poll(request_id: str):
|
|
return _poll_result("extract", request_id)
|
|
|
|
# POST /api/v1/segment — document segmentation
|
|
@app.route("/api/v1/segment", methods=["POST"])
|
|
@require_api_key
|
|
def api_segment_submit():
|
|
if "file" in request.files:
|
|
file = request.files["file"]
|
|
raw = file.read()
|
|
filename = file.filename or "file"
|
|
elif request.is_json:
|
|
data = request.get_json()
|
|
if "file_b64" in data:
|
|
raw = base64.b64decode(data["file_b64"])
|
|
filename = data.get("filename", "file")
|
|
else:
|
|
return jsonify({"error": "Provide file or file_b64"}), 400
|
|
else:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
source = request.form if request.form else (request.get_json() if request.is_json else {})
|
|
opts = parse_options(source)
|
|
opts.setdefault("output_format", "markdown")
|
|
return _submit_and_poll("segment", raw, filename, opts=opts)
|
|
|
|
# GET /api/v1/segment/<request_id>
|
|
@app.route("/api/v1/segment/<request_id>", methods=["GET"])
|
|
@require_api_key
|
|
def api_segment_poll(request_id: str):
|
|
return _poll_result("segment", request_id)
|
|
|
|
# POST /api/v1/ocr — deprecated OCR-only endpoint
|
|
@app.route("/api/v1/ocr", methods=["POST"])
|
|
@require_api_key
|
|
def api_ocr_submit():
|
|
if "file" not in request.files:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
file = request.files["file"]
|
|
raw = file.read()
|
|
filename = file.filename or "file"
|
|
opts = parse_options(request.form)
|
|
return _submit_and_poll("ocr", raw, filename, opts=opts)
|
|
|
|
# GET /api/v1/ocr/<request_id>
|
|
@app.route("/api/v1/ocr/<request_id>", methods=["GET"])
|
|
@require_api_key
|
|
def api_ocr_poll(request_id: str):
|
|
return _poll_result("ocr", request_id)
|
|
|
|
# POST /api/v1/table_rec — deprecated table recognition
|
|
@app.route("/api/v1/table_rec", methods=["POST"])
|
|
@require_api_key
|
|
def api_table_rec_submit():
|
|
if "file" not in request.files:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
file = request.files["file"]
|
|
raw = file.read()
|
|
filename = file.filename or "file"
|
|
opts = parse_options(request.form)
|
|
return _submit_and_poll("table_rec", raw, filename, opts=opts)
|
|
|
|
# GET /api/v1/table_rec/<request_id>
|
|
@app.route("/api/v1/table_rec/<request_id>", methods=["GET"])
|
|
@require_api_key
|
|
def api_table_rec_poll(request_id: str):
|
|
return _poll_result("table_rec", request_id)
|
|
|
|
# POST /api/v1/create-document — create DOCX from markdown
|
|
@app.route("/api/v1/create-document", methods=["POST"])
|
|
@require_api_key
|
|
def api_create_document():
|
|
data = request.get_json() if request.is_json else {}
|
|
if not data or "markdown" not in data:
|
|
return jsonify({"error": "JSON body must include 'markdown'"}), 400
|
|
md = data["markdown"]
|
|
output_format = data.get("output_format", "docx")
|
|
try:
|
|
import subprocess
|
|
import base64
|
|
with tempfile.NamedTemporaryFile(mode="w", suffix=".md", delete=False) as f:
|
|
f.write(md)
|
|
md_path = f.name
|
|
out_path = md_path.replace(".md", f".{output_format}")
|
|
if output_format == "docx":
|
|
subprocess.run(["pandoc", md_path, "-o", out_path], capture_output=True, timeout=60)
|
|
with open(out_path, "rb") as f:
|
|
b64 = base64.b64encode(f.read()).decode()
|
|
os.unlink(out_path)
|
|
os.unlink(md_path)
|
|
return jsonify({
|
|
"success": True,
|
|
"output_format": "docx",
|
|
"output_base64": b64,
|
|
})
|
|
else:
|
|
return jsonify({"error": f"Unsupported output format: {output_format}"}), 400
|
|
except FileNotFoundError:
|
|
return jsonify({"error": "pandoc not installed, cannot create document"}), 500
|
|
except Exception as exc:
|
|
return jsonify({"success": False, "error": str(exc)}), 500
|
|
|
|
# ── File management endpoints ────────────────────────────────────
|
|
@app.route("/api/v1/files/upload", methods=["POST"])
|
|
@require_api_key
|
|
def api_file_upload():
|
|
data = request.get_json() if request.is_json else {}
|
|
filename = data.get("filename", "file")
|
|
content_type = data.get("content_type", "application/octet-stream")
|
|
import time as tmod
|
|
file_id = str(uuid.uuid4())
|
|
file_dir = _file_storage_path() / file_id
|
|
file_dir.mkdir(parents=True, exist_ok=True)
|
|
metadata = {
|
|
"file_id": file_id,
|
|
"original_filename": filename,
|
|
"filename": filename,
|
|
"content_type": content_type,
|
|
"reference": f"datalab://file-{file_id}",
|
|
"created_at": time.time(),
|
|
"created": tmod.strftime("%Y-%m-%dT%H:%M:%SZ", tmod.gmtime()),
|
|
"file_size": None,
|
|
"upload_status": "pending",
|
|
"error": None,
|
|
}
|
|
(file_dir / "metadata.json").write_text(json.dumps(metadata))
|
|
import secrets
|
|
token = secrets.token_urlsafe(32)
|
|
upload_url = f"{BASE_URL}/api/v1/files/{file_id}/upload/{token}"
|
|
(file_dir / "upload_token.txt").write_text(token)
|
|
return jsonify({
|
|
"file_id": file_id,
|
|
"upload_url": upload_url,
|
|
"reference": f"datalab://file-{file_id}",
|
|
})
|
|
|
|
@app.route("/api/v1/files/<file_id>/upload/<token>", methods=["PUT"])
|
|
def api_file_upload_put(file_id: str, token: str):
|
|
file_dir = _file_storage_path() / file_id
|
|
token_path = file_dir / "upload_token.txt"
|
|
if not token_path.exists() or token_path.read_text().strip() != token:
|
|
return jsonify({"error": "Invalid upload token"}), 403
|
|
(file_dir / "content").write_bytes(request.data)
|
|
metadata_path = file_dir / "metadata.json"
|
|
if metadata_path.exists():
|
|
import time as tmod
|
|
meta = json.loads(metadata_path.read_text())
|
|
meta["uploaded_at"] = tmod.time()
|
|
meta["file_size"] = len(request.data)
|
|
meta["upload_status"] = "completed"
|
|
metadata_path.write_text(json.dumps(meta))
|
|
return jsonify({"success": True})
|
|
|
|
@app.route("/api/v1/files", methods=["GET"])
|
|
@require_api_key
|
|
def api_file_list():
|
|
storage = _file_storage_path()
|
|
limit = int(request.args.get("limit", 50))
|
|
offset = int(request.args.get("offset", 0))
|
|
files = []
|
|
for child in sorted(storage.iterdir(), reverse=True):
|
|
if child.is_dir():
|
|
meta = _get_file_metadata(child.name)
|
|
if meta:
|
|
files.append(meta)
|
|
return jsonify({"files": files[offset:offset + limit], "total": len(files)})
|
|
|
|
@app.route("/api/v1/files/<file_id>", methods=["GET"])
|
|
@require_api_key
|
|
def api_file_get(file_id: str):
|
|
meta = _get_file_metadata(file_id)
|
|
if meta is None:
|
|
return jsonify({"error": "File not found"}), 404
|
|
return jsonify(meta)
|
|
|
|
@app.route("/api/v1/files/<file_id>/confirm", methods=["GET"])
|
|
@require_api_key
|
|
def api_file_confirm(file_id: str):
|
|
meta = _get_file_metadata(file_id)
|
|
if meta is None:
|
|
return jsonify({"error": "File not found"}), 404
|
|
file_dir = _file_storage_path() / file_id
|
|
content_path = file_dir / "content"
|
|
if not content_path.exists():
|
|
return jsonify({"error": "File content not uploaded yet"}), 400
|
|
return jsonify({"success": True, "file_id": file_id})
|
|
|
|
@app.route("/api/v1/files/<file_id>/download", methods=["GET"])
|
|
@require_api_key
|
|
def api_file_download(file_id: str):
|
|
meta = _get_file_metadata(file_id)
|
|
if meta is None:
|
|
return jsonify({"error": "File not found"}), 404
|
|
expires_in = int(request.args.get("expires_in", 3600))
|
|
import secrets
|
|
token = secrets.token_urlsafe(32)
|
|
file_dir = _file_storage_path() / file_id
|
|
(file_dir / "download_token.txt").write_text(token)
|
|
download_url = f"{BASE_URL}/api/v1/files/{file_id}/download/{token}?expires_in={expires_in}"
|
|
return jsonify({"download_url": download_url})
|
|
|
|
@app.route("/api/v1/files/<file_id>/download/<token>", methods=["GET"])
|
|
def api_file_download_token(file_id: str, token: str):
|
|
file_dir = _file_storage_path() / file_id
|
|
token_path = file_dir / "download_token.txt"
|
|
if not token_path.exists() or token_path.read_text().strip() != token:
|
|
return jsonify({"error": "Invalid download token"}), 403
|
|
content_path = file_dir / "content"
|
|
if not content_path.exists():
|
|
return jsonify({"error": "File content not found"}), 404
|
|
meta = _get_file_metadata(file_id) or {}
|
|
return Response(
|
|
content_path.read_bytes(),
|
|
mimetype=meta.get("content_type", "application/octet-stream"),
|
|
headers={"Content-Disposition": f'attachment; filename="{meta.get("filename", "file")}"'},
|
|
)
|
|
|
|
@app.route("/api/v1/files/<file_id>", methods=["DELETE"])
|
|
@require_api_key
|
|
def api_file_delete(file_id: str):
|
|
file_dir = _file_storage_path() / file_id
|
|
if file_dir.exists():
|
|
import shutil
|
|
shutil.rmtree(str(file_dir))
|
|
return jsonify({"success": True})
|
|
|
|
# ── Thumbnails ───────────────────────────────────────────────────
|
|
@app.route("/api/v1/thumbnails/<lookup_key>", methods=["GET"])
|
|
@require_api_key
|
|
def api_thumbnails(lookup_key: str):
|
|
thumb_width = int(request.args.get("thumb_width", 300))
|
|
page_range = request.args.get("page_range")
|
|
entry = request_store.get_request(lookup_key)
|
|
if entry is None:
|
|
return jsonify({"success": False, "error": "Request not found"}), 404
|
|
result = request_store.get_result(lookup_key)
|
|
if result is None:
|
|
return jsonify({"success": False, "error": "No result found"}), 404
|
|
metadata = result.get("metadata", {})
|
|
page_count = metadata.get("page_count", 0) if isinstance(metadata, dict) else 0
|
|
if page_range:
|
|
pages = []
|
|
for part in page_range.split(","):
|
|
part = part.strip()
|
|
if "-" in part:
|
|
a, b = part.split("-", 1)
|
|
pages.extend(range(int(a.strip()), int(b.strip()) + 1))
|
|
else:
|
|
pages.append(int(part))
|
|
else:
|
|
pages = list(range(page_count))
|
|
thumbnails = []
|
|
for p in pages:
|
|
if p < page_count:
|
|
thumbnails.append("") # placeholder — no rendered page images stored
|
|
return jsonify({
|
|
"success": True,
|
|
"thumbnails": thumbnails,
|
|
})
|
|
|
|
return app
|
|
|
|
|
|
app_instance = create_app()
|
|
|
|
if __name__ == "__main__":
|
|
_configure_env()
|
|
_marker_dict = create_model_dict()
|
|
port = int(os.environ.get("PORT", "8001"))
|
|
host = os.environ.get("HOST", "0.0.0.0")
|
|
print("=" * 60)
|
|
print("marker-api starting")
|
|
print(f" OCR_BACKEND = {OCR_BACKEND}")
|
|
print(f" API_KEY = {'set' if API_KEY else '(not set)'}")
|
|
print(f" OLLAMA_HOST = {OLLAMA_HOST}")
|
|
print(f" TORCH_DEVICE = {TORCH_DEVICE or '(auto)'}")
|
|
print(f" LISTENING ON = {host}:{port}")
|
|
print(f" FORMATS = {', '.join(SUPPORTED_DISPLAY)}")
|
|
print("=" * 60)
|
|
app_instance.run(host=host, port=port, debug=(os.environ.get("FLASK_DEBUG", "0") == "1"), threaded=True)
|