extract.py

import argparse
import sys
import time
from pathlib import Path


def extract_text(input_path: Path, output_dir: Path):
    """Normalized extraction of text/markdown files."""
    output_path = output_dir / (input_path.stem + ".md")
    with open(input_path, "r", encoding="utf-8") as f:
        content = f.read()

    if not content.strip():
        content = ""
    elif not content.endswith("\n"):
        content += "\n"

    with open(output_path, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"[extract] {input_path} -> {output_path}", flush=True)


def extract_pdf(input_path: Path, output_dir: Path, model_dict):
    """Convert PDF to markdown using marker-pdf."""
    try:
        from marker.converters.pdf import PdfConverter
        from marker.config.parser import ConfigParser
        from marker.output import text_from_rendered
    except ImportError as e:
        raise ImportError(
            "Missing PDF dependencies. Install OCR extras: uv sync --extra ocr"
        ) from e

    print(f"[extract] {input_path.name} converting ...", flush=True)
    t0 = time.monotonic()

    config_parser = ConfigParser({"output_format": "markdown", "disable_image_extraction": True})
    converter = PdfConverter(
        config=config_parser.generate_config_dict(),
        artifact_dict=model_dict,
        processor_list=config_parser.get_processors(),
        renderer=config_parser.get_renderer(),
    )
    rendered = converter(str(input_path))
    text, _, images = text_from_rendered(rendered)

    elapsed = time.monotonic() - t0
    output_path = output_dir / (input_path.stem + ".md")
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(text)
    print(
        f"[extract] done  {input_path.name} -> {output_path.name}  "
        f"{elapsed:.1f}s  {len(text)} chars",
        flush=True,
    )


def main():
    parser = argparse.ArgumentParser(description="Stage-1 extraction: Source to Markdown")
    parser.add_argument("inputs", nargs="+", help="Input files or directories")
    parser.add_argument("--output-dir", required=True, help="Directory to save extracted markdown")
    # --model and --base-url kept for CLI compatibility but no longer used
    parser.add_argument("--model", default="gemma4-26b-ocr", help=argparse.SUPPRESS)
    parser.add_argument("--base-url", default="http://localhost:11434/v1", help=argparse.SUPPRESS)

    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Collect all files first so we can load marker models once if any PDFs are present
    all_files: list[tuple[Path, str]] = []
    for input_str in args.inputs:
        input_path = Path(input_str)
        if not input_path.exists():
            print(f"[extract] Warning: input path does not exist: {input_path}", flush=True)
            continue

        if input_path.is_dir():
            candidates = sorted(
                list(input_path.glob("*.md"))
                + list(input_path.glob("*.txt"))
                + list(input_path.glob("*.pdf"))
            )
            files = [f for f in candidates if not f.stem.upper().startswith("README")]
        else:
            files = [input_path]

        for f in files:
            suffix = f.suffix.lower()
            if suffix in (".md", ".txt"):
                all_files.append((f, "text"))
            elif suffix == ".pdf":
                all_files.append((f, "pdf"))
            else:
                print(f"[extract] skipping unsupported type: {f}", flush=True)

    pdf_count = sum(1 for _, t in all_files if t == "pdf")
    print(
        f"[extract] {len(all_files)} file(s) to process  ({pdf_count} PDF)",
        flush=True,
    )

    # Load marker models once for all PDFs that actually need processing
    pending_pdfs = sum(
        1 for f, t in all_files
        if t == "pdf" and not (output_dir / (f.stem + ".md")).exists()
    )
    model_dict = None
    if pending_pdfs > 0:
        try:
            from marker.models import create_model_dict
        except ImportError as e:
            print(
                f"[extract] ERROR: marker-pdf not installed. Run: uv sync --extra ocr\n{e}",
                flush=True,
            )
            sys.exit(1)
        print("[extract] loading marker models (first run downloads ~1 GB) ...", flush=True)
        t0 = time.monotonic()
        model_dict = create_model_dict()
        print(f"[extract] models ready  {time.monotonic() - t0:.1f}s", flush=True)

    for f, file_type in all_files:
        output_path = output_dir / (f.stem + ".md")
        if output_path.exists():
            print(f"[extract] skip (exists)  {f.name} -> {output_path.name}", flush=True)
            continue
        try:
            if file_type == "text":
                extract_text(f, output_dir)
            else:
                extract_pdf(f, output_dir, model_dict)
        except Exception as e:
            print(f"[extract] ERROR processing {f}: {e}", flush=True)
            sys.exit(1)


if __name__ == "__main__":
    main()