From 02185bef46ccc0e594f62cac226de73005e09748 Mon Sep 17 00:00:00 2001
From: Aaron Roberts <aaronroberts@Aarons-MacBook-Air.local>
Date: Tue, 30 Jun 2026 12:16:16 +0100
Subject: [PATCH] Adding missed files

---
 backend/providers.py                      | 489 ++++++++++++++++++++++
 frontend/src/components/ModelSelector.jsx |  33 ++
 frontend/src/hooks/useModels.js           |  24 ++
 3 files changed, 546 insertions(+)
 create mode 100644 backend/providers.py
 create mode 100644 frontend/src/components/ModelSelector.jsx
 create mode 100644 frontend/src/hooks/useModels.js

diff --git a/backend/providers.py b/backend/providers.py
new file mode 100644
index 0000000..34be178
--- /dev/null
+++ b/backend/providers.py
@@ -0,0 +1,489 @@
+"""
+OCR provider abstraction.
+
+Each provider knows how to turn an image + a semantic OCR request (mode, prompt,
+options) into raw model text. DeepSeek-specific prompt tokens and grounding-box
+parsing live here too so the FastAPI routes stay model-agnostic.
+
+Two providers ship today:
+  - DeepSeekLocalProvider  -> the local HF transformers DeepSeek-OCR model (GPU)
+  - OllamaProvider         -> any vision model served by an external Ollama host
+
+The registry is built from environment variables at startup (see build_registry()).
+"""
+
+import os
+import re
+import base64
+import tempfile
+import shutil
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional
+
+from decouple import config as env_config
+
+# httpx is only needed when an Ollama model is actually used; import lazily so the
+# backend can run DeepSeek-only without the dependency installed.
+try:
+    import httpx
+except Exception:  # pragma: no cover - exercised only when httpx is missing
+    httpx = None
+
+
+# =============================================================================
+# Prompt builders
+# =============================================================================
+def build_prompt(
+    mode: str,
+    user_prompt: str,
+    grounding: bool,
+    find_term: Optional[str],
+    schema: Optional[str],
+    include_caption: bool,
+) -> str:
+    """Build the DeepSeek-OCR prompt (with its special tokens) based on mode."""
+    parts: List[str] = ["<image>"]
+    mode_requires_grounding = mode in {"find_ref", "layout_map", "pii_redact"}
+    if grounding or mode_requires_grounding:
+        parts.append("<|grounding|>")
+
+    parts.append(_instruction_for_mode(mode, user_prompt, find_term, schema, include_caption))
+    return "\n".join(parts)
+
+
+def build_ollama_prompt(
+    mode: str,
+    user_prompt: str,
+    find_term: Optional[str],
+    schema: Optional[str],
+    include_caption: bool,
+) -> str:
+    """Build a plain natural-language prompt for a generic vision model.
+
+    No DeepSeek grounding tokens — Ollama vision models receive the image
+    separately and respond in plain text.
+    """
+    if mode == "plain_ocr":
+        instruction = (
+            "Transcribe all of the text in this image exactly as it appears, "
+            "preserving line breaks and reading order. Output only the transcribed "
+            "text with no commentary."
+        )
+    elif mode == "markdown":
+        instruction = (
+            "Convert this document image to clean GitHub-flavored Markdown, "
+            "preserving headings, lists, and tables. Output only the Markdown."
+        )
+    elif mode == "tables_csv":
+        instruction = (
+            "Extract every table in this image and output CSV only. Use commas with "
+            "minimal quoting. If there are multiple tables, separate them with a line "
+            "containing '---'. Output only the CSV."
+        )
+    elif mode == "tables_md":
+        instruction = (
+            "Extract every table in this image as GitHub-flavored Markdown tables. "
+            "Output only the tables."
+        )
+    elif mode == "kv_json":
+        schema_text = schema.strip() if schema else "{}"
+        instruction = (
+            "Extract the key fields from this image and return strict JSON only "
+            f"(no prose). Use this schema, filling in the values: {schema_text}"
+        )
+    elif mode == "figure_chart":
+        instruction = (
+            "Parse the figure in this image. First extract any numeric series as a "
+            "two-column table (x,y). Then add a line containing '---' followed by a "
+            "two-sentence summary of the chart."
+        )
+    elif mode == "find_ref":
+        key = (find_term or "").strip() or "Total"
+        instruction = (
+            f"Find every occurrence of '{key}' in this image and quote the surrounding "
+            "text for each match. If it does not appear, say so."
+        )
+    elif mode == "layout_map":
+        instruction = (
+            'Identify the layout blocks in this image and return a JSON array of '
+            'objects {"type": one of ["title","paragraph","table","figure"]}. '
+            "Do not include the text content."
+        )
+    elif mode == "pii_redact":
+        instruction = (
+            "Find all emails, phone numbers, postal addresses, and IBANs in this image. "
+            'Return a JSON array of objects {"label", "text"}.'
+        )
+    elif mode == "multilingual":
+        instruction = (
+            "Transcribe all of the text in this image exactly, detecting the language "
+            "automatically and preserving the original script. Output only the text."
+        )
+    elif mode == "describe":
+        instruction = "Describe this image, focusing on the key visible elements."
+    elif mode == "freeform":
+        instruction = user_prompt.strip() if user_prompt else "Transcribe the text in this image."
+    else:
+        instruction = "Transcribe the text in this image."
+
+    if include_caption and mode != "describe":
+        instruction += "\nThen add a one-paragraph description of the image."
+
+    return instruction
+
+
+def _instruction_for_mode(
+    mode: str,
+    user_prompt: str,
+    find_term: Optional[str],
+    schema: Optional[str],
+    include_caption: bool,
+) -> str:
+    """The DeepSeek instruction text (without the <image>/<|grounding|> prefix tokens)."""
+    if mode == "plain_ocr":
+        instruction = "Free OCR."
+    elif mode == "markdown":
+        instruction = "Convert the document to markdown."
+    elif mode == "tables_csv":
+        instruction = (
+            "Extract every table and output CSV only. "
+            "Use commas, minimal quoting. If multiple tables, separate with a line containing '---'."
+        )
+    elif mode == "tables_md":
+        instruction = "Extract every table as GitHub-flavored Markdown tables. Output only the tables."
+    elif mode == "kv_json":
+        schema_text = schema.strip() if schema else "{}"
+        instruction = (
+            "Extract key fields and return strict JSON only. "
+            f"Use this schema (fill the values): {schema_text}"
+        )
+    elif mode == "figure_chart":
+        instruction = (
+            "Parse the figure. First extract any numeric series as a two-column table (x,y). "
+            "Then summarize the chart in 2 sentences. Output the table, then a line '---', then the summary."
+        )
+    elif mode == "find_ref":
+        key = (find_term or "").strip() or "Total"
+        instruction = f"Locate <|ref|>{key}<|/ref|> in the image."
+    elif mode == "layout_map":
+        instruction = (
+            'Return a JSON array of blocks with fields {"type":["title","paragraph","table","figure"],'
+            '"box":[x1,y1,x2,y2]}. Do not include any text content.'
+        )
+    elif mode == "pii_redact":
+        instruction = (
+            'Find all occurrences of emails, phone numbers, postal addresses, and IBANs. '
+            'Return a JSON array of objects {label, text, box:[x1,y1,x2,y2]}.'
+        )
+    elif mode == "multilingual":
+        instruction = "Free OCR. Detect the language automatically and output in the same script."
+    elif mode == "describe":
+        instruction = "Describe this image. Focus on visible key elements."
+    elif mode == "freeform":
+        instruction = user_prompt.strip() if user_prompt else "OCR this image."
+    else:
+        instruction = "OCR this image."
+
+    if include_caption and mode != "describe":
+        instruction = instruction + "\nThen add a one-paragraph description of the image."
+
+    return instruction
+
+
+# =============================================================================
+# Grounding parser (DeepSeek-specific; no-op on plain text)
+# =============================================================================
+DET_BLOCK = re.compile(
+    r"<\|ref\|>(?P<label>.*?)<\|/ref\|>\s*<\|det\|>\s*(?P<coords>\[.*\])\s*<\|/det\|>",
+    re.DOTALL,
+)
+
+
+def clean_grounding_text(text: str) -> str:
+    """Remove grounding tags from text for display, keeping labels."""
+    cleaned = re.sub(
+        r"<\|ref\|>(.*?)<\|/ref\|>\s*<\|det\|>\s*\[.*\]\s*<\|/det\|>",
+        r"\1",
+        text,
+        flags=re.DOTALL,
+    )
+    cleaned = re.sub(r"<\|grounding\|>", "", cleaned)
+    return cleaned.strip()
+
+
+def parse_detections(text: str, image_width: int, image_height: int) -> List[Dict[str, Any]]:
+    """Parse grounding boxes from text and scale 0-999 normalized coords to pixels."""
+    boxes: List[Dict[str, Any]] = []
+    for m in DET_BLOCK.finditer(text or ""):
+        label = m.group("label").strip()
+        coords_str = m.group("coords").strip()
+
+        try:
+            import ast
+
+            parsed = ast.literal_eval(coords_str)
+
+            if (
+                isinstance(parsed, list)
+                and len(parsed) == 4
+                and all(isinstance(n, (int, float)) for n in parsed)
+            ):
+                box_coords = [parsed]
+            elif isinstance(parsed, list):
+                box_coords = parsed
+            else:
+                raise ValueError("Unsupported coords structure")
+
+            for box in box_coords:
+                if isinstance(box, (list, tuple)) and len(box) >= 4:
+                    x1 = int(float(box[0]) / 999 * image_width)
+                    y1 = int(float(box[1]) / 999 * image_height)
+                    x2 = int(float(box[2]) / 999 * image_width)
+                    y2 = int(float(box[3]) / 999 * image_height)
+                    boxes.append({"label": label, "box": [x1, y1, x2, y2]})
+        except Exception as e:
+            print(f"❌ Grounding parse failed: {e}")
+            continue
+
+    return boxes
+
+
+# =============================================================================
+# Providers
+# =============================================================================
+GROUNDING_MODES = {"find_ref", "layout_map", "pii_redact"}
+
+
+class ProviderError(Exception):
+    """Raised when a provider cannot fulfil a request (e.g. backend unreachable)."""
+
+
+class OCRProvider(ABC):
+    """Turns an image + OCR request into raw model text."""
+
+    id: str
+    label: str
+    capabilities: Dict[str, Any]
+
+    @abstractmethod
+    def run(
+        self,
+        image_path: str,
+        *,
+        mode: str,
+        prompt: str,
+        grounding: bool,
+        find_term: Optional[str],
+        schema: Optional[str],
+        include_caption: bool,
+        options: Dict[str, Any],
+    ) -> str:
+        """Return the raw text output of the model for this image/request."""
+
+    def info(self) -> Dict[str, Any]:
+        return {"id": self.id, "label": self.label, "capabilities": self.capabilities}
+
+
+class DeepSeekLocalProvider(OCRProvider):
+    """Local HF transformers DeepSeek-OCR model. Loaded lazily on first use."""
+
+    def __init__(self):
+        self.id = "deepseek-local"
+        self.label = "DeepSeek-OCR (local GPU)"
+        self.capabilities = {"grounding": True, "advanced_settings": True}
+        self._model = None
+        self._tokenizer = None
+
+    @property
+    def loaded(self) -> bool:
+        return self._model is not None and self._tokenizer is not None
+
+    def _ensure_loaded(self):
+        if self.loaded:
+            return
+
+        # Heavy imports kept local so an Ollama-only deployment never needs torch.
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+
+        os.environ.pop("TRANSFORMERS_CACHE", None)
+        model_name = env_config("MODEL_NAME", default="deepseek-ai/DeepSeek-OCR")
+        hf_home = env_config("HF_HOME", default="/models")
+        os.makedirs(hf_home, exist_ok=True)
+
+        print(f"🚀 Loading {model_name}...")
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        model = AutoModel.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            use_safetensors=True,
+            attn_implementation="eager",
+            torch_dtype=torch.bfloat16,
+        ).eval().to("cuda")
+
+        try:
+            if getattr(tokenizer, "pad_token_id", None) is None and getattr(tokenizer, "eos_token_id", None) is not None:
+                tokenizer.pad_token = tokenizer.eos_token
+            if getattr(model.config, "pad_token_id", None) is None and getattr(tokenizer, "pad_token_id", None) is not None:
+                model.config.pad_token_id = tokenizer.pad_token_id
+        except Exception:
+            pass
+
+        self._model = model
+        self._tokenizer = tokenizer
+        print("✅ DeepSeek-OCR loaded and ready!")
+
+    def run(self, image_path, *, mode, prompt, grounding, find_term, schema, include_caption, options):
+        self._ensure_loaded()
+
+        prompt_text = build_prompt(
+            mode=mode,
+            user_prompt=prompt,
+            grounding=grounding,
+            find_term=find_term,
+            schema=schema,
+            include_caption=include_caption,
+        )
+
+        out_dir = tempfile.mkdtemp(prefix="dsocr_")
+        try:
+            res = self._model.infer(
+                self._tokenizer,
+                prompt=prompt_text,
+                image_file=image_path,
+                output_path=out_dir,
+                base_size=int(options.get("base_size", 1024)),
+                image_size=int(options.get("image_size", 640)),
+                crop_mode=bool(options.get("crop_mode", True)),
+                save_results=False,
+                test_compress=bool(options.get("test_compress", False)),
+                eval_mode=True,
+            )
+
+            if isinstance(res, str):
+                text = res.strip()
+            elif isinstance(res, dict) and "text" in res:
+                text = str(res["text"]).strip()
+            elif isinstance(res, (list, tuple)):
+                text = "\n".join(map(str, res)).strip()
+            else:
+                text = ""
+
+            if not text:
+                mmd = os.path.join(out_dir, "result.mmd")
+                if os.path.exists(mmd):
+                    with open(mmd, "r", encoding="utf-8") as fh:
+                        text = fh.read().strip()
+            return text
+        finally:
+            shutil.rmtree(out_dir, ignore_errors=True)
+
+
+class OllamaProvider(OCRProvider):
+    """A single vision model served by an external Ollama host."""
+
+    def __init__(self, tag: str, base_url: str, label: Optional[str] = None):
+        self.tag = tag
+        self.base_url = base_url.rstrip("/")
+        self.id = f"ollama:{tag}"
+        self.label = label or f"{tag} (Ollama)"
+        # Generic vision models don't emit DeepSeek grounding tokens.
+        self.capabilities = {"grounding": False, "advanced_settings": False}
+
+    def run(self, image_path, *, mode, prompt, grounding, find_term, schema, include_caption, options):
+        if httpx is None:
+            raise ProviderError("httpx is not installed; cannot reach Ollama.")
+
+        prompt_text = build_ollama_prompt(
+            mode=mode,
+            user_prompt=prompt,
+            find_term=find_term,
+            schema=schema,
+            include_caption=include_caption,
+        )
+
+        with open(image_path, "rb") as f:
+            img_b64 = base64.b64encode(f.read()).decode("utf-8")
+
+        payload = {
+            "model": self.tag,
+            "prompt": prompt_text,
+            "images": [img_b64],
+            "stream": False,
+        }
+
+        timeout = float(env_config("OLLAMA_TIMEOUT", default=300.0, cast=float))
+        try:
+            resp = httpx.post(f"{self.base_url}/api/generate", json=payload, timeout=timeout)
+            resp.raise_for_status()
+            data = resp.json()
+        except httpx.HTTPStatusError as e:
+            detail = ""
+            try:
+                detail = e.response.json().get("error", "")
+            except Exception:
+                detail = e.response.text[:200]
+            raise ProviderError(f"Ollama returned {e.response.status_code}: {detail}") from e
+        except httpx.HTTPError as e:
+            raise ProviderError(f"Could not reach Ollama at {self.base_url}: {e}") from e
+
+        return (data.get("response") or "").strip()
+
+
+# =============================================================================
+# Registry
+# =============================================================================
+class ModelRegistry:
+    def __init__(self, providers: List[OCRProvider], default_id: str):
+        self._providers: Dict[str, OCRProvider] = {p.id: p for p in providers}
+        # Fall back to the first registered provider if the configured default is gone.
+        self.default_id = default_id if default_id in self._providers else (
+            next(iter(self._providers), None)
+        )
+
+    def get(self, model_id: Optional[str]) -> OCRProvider:
+        chosen = model_id or self.default_id
+        provider = self._providers.get(chosen)
+        if provider is None:
+            raise ProviderError(f"Unknown model '{chosen}'.")
+        return provider
+
+    def list_models(self) -> List[Dict[str, Any]]:
+        out = []
+        for p in self._providers.values():
+            entry = p.info()
+            entry["default"] = (p.id == self.default_id)
+            out.append(entry)
+        return out
+
+
+def build_registry() -> ModelRegistry:
+    """Build the provider registry from environment variables.
+
+    Env:
+      ENABLE_DEEPSEEK_LOCAL  - register the local DeepSeek-OCR model (default: true)
+      OLLAMA_BASE_URL        - Ollama host (default: http://host.docker.internal:11434)
+      OLLAMA_MODELS          - comma-separated tags to surface (e.g. "glm-ocr,llama3.2-vision")
+      DEFAULT_OCR_MODEL      - id to select by default (default: deepseek-local)
+    """
+    providers: List[OCRProvider] = []
+
+    enable_deepseek = env_config("ENABLE_DEEPSEEK_LOCAL", default="true").strip().lower() in {"1", "true", "yes"}
+    if enable_deepseek:
+        providers.append(DeepSeekLocalProvider())
+
+    base_url = env_config("OLLAMA_BASE_URL", default="http://host.docker.internal:11434")
+    raw_tags = env_config("OLLAMA_MODELS", default="")
+    tags = [t.strip() for t in raw_tags.split(",") if t.strip()]
+    for tag in tags:
+        providers.append(OllamaProvider(tag=tag, base_url=base_url))
+
+    default_id = env_config("DEFAULT_OCR_MODEL", default="deepseek-local")
+    if not providers:
+        # Defensive: nothing configured. Register DeepSeek so the app still starts.
+        providers.append(DeepSeekLocalProvider())
+        default_id = "deepseek-local"
+
+    registry = ModelRegistry(providers, default_id)
+    print(f"🧠 OCR models registered: {[p.id for p in providers]} (default: {registry.default_id})")
+    return registry
diff --git a/frontend/src/components/ModelSelector.jsx b/frontend/src/components/ModelSelector.jsx
new file mode 100644
index 0000000..8820596
--- /dev/null
+++ b/frontend/src/components/ModelSelector.jsx
@@ -0,0 +1,33 @@
+import { Cpu } from 'lucide-react'
+
+const SELECT_CLASS =
+  'w-full bg-white/5 border border-white/10 rounded-lg px-3 py-2 text-sm text-gray-200 ' +
+  'focus:outline-none focus:border-purple-500/50 transition-colors'
+
+// Dropdown to pick which OCR model runs the analysis.
+// `models` comes from the useModels() hook; `value` is the selected model id.
+export default function ModelSelector({ models, value, onChange, loading }) {
+  return (
+    <div className="glass p-4 rounded-2xl space-y-3">
+      <div className="flex items-center gap-2">
+        <Cpu className="w-4 h-4 text-purple-400" />
+        <h3 className="text-sm font-semibold text-gray-200">Model</h3>
+      </div>
+
+      <select
+        value={value || ''}
+        onChange={e => onChange(e.target.value)}
+        disabled={loading || models.length === 0}
+        className={SELECT_CLASS}
+      >
+        {loading && <option value="">Loading models…</option>}
+        {!loading && models.length === 0 && <option value="">No models available</option>}
+        {models.map(m => (
+          <option key={m.id} value={m.id}>
+            {m.label}{m.default ? ' (default)' : ''}
+          </option>
+        ))}
+      </select>
+    </div>
+  )
+}
diff --git a/frontend/src/hooks/useModels.js b/frontend/src/hooks/useModels.js
new file mode 100644
index 0000000..a4d5217
--- /dev/null
+++ b/frontend/src/hooks/useModels.js
@@ -0,0 +1,24 @@
+import { useState, useEffect } from 'react'
+
+const API_BASE = import.meta.env.VITE_API_URL || '/api'
+
+// Fetches the OCR models available for selection. Returns { models, loading }.
+// Each model: { id, label, capabilities: { grounding, advanced_settings }, default }
+export function useModels() {
+  const [models, setModels] = useState([])
+  const [loading, setLoading] = useState(true)
+
+  useEffect(() => {
+    let cancelled = false
+    fetch(`${API_BASE}/models`)
+      .then(r => (r.ok ? r.json() : null))
+      .then(data => {
+        if (!cancelled && data?.models) setModels(data.models)
+      })
+      .catch(() => {})
+      .finally(() => { if (!cancelled) setLoading(false) })
+    return () => { cancelled = true }
+  }, [])
+
+  return { models, loading }
+}