Feat/UI deepseek

.gitignore

@@ -47,3 +47,8 @@ node_modules/
 tmp/
 demos/demo_tiles/
 assets/*.mp4
+
+# Research knowledge base (local only — papers are large & copyrighted, indexes are generated)
+papers/
+knowledge_index/
+e2e_test/index/

CLAUDE.md

@@ -1,3 +1,7 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
 # PixelRAG
 
 Visual RAG: render documents (web pages, PDFs) as screenshot tiles and retrieve over the images
@@ -6,16 +10,87 @@ for how it runs in production.
 
 ## Layout
 - `render/` — `pixelshot` capture (Playwright/CDP, PDF)
-- `embed/`, `index/` — tiles → vectors → FAISS index
-- `serve/` — FAISS search API (`pixelrag serve`)
+- `embed/` — tiles → chunks → vectors → index (the `chunk`, `embed`, `build-index` stages)
+- `index/` — `pixelrag index` orchestrator + pluggable document sources
+- `serve/` — FAISS/Qdrant search API (`pixelrag serve`)
 - `web/` — Next.js frontend (on Vercel) + `agent-server.mjs` (the chat agent backend)
 - `train/` — **separate uv project** (LoRA finetune); install from inside `train/`, not the root
+- `eval/` — **separate uv project** reproducing paper Table 1 (reader + grader over HTTP)
 - `deploy/` — systemd units, CD workflow, blue-green scripts
 
+## Commands
+
+Python (uv only — `uv add`, never `uv pip install`; work in `.venv`; commit `uv.lock`):
+
+```bash
+# Install dev deps and run the full Python suite (render tests need Chrome too)
+uv sync --extra dev
+uv run pixelshot install-chrome        # one-time; CI also runs `npx -y playwright install-deps chromium`
+uv run pytest tests/ -v
+uv run pytest tests/test_cli.py::test_<name> -v   # single test
+
+# Lint / format (ruff; config in pyproject.toml excludes train/, demos/, docs/, skill/)
+uvx ruff check .
+uvx ruff format --check .
+
+# Pipeline: build an index from a pixelrag.yaml config, then serve it
+pixelrag index build -c pixelrag.yaml
+pixelrag serve --index-dir ./index --port 30001
+```
+
+Frontend (`cd web`):
+
+```bash
+npm ci
+npm test            # node --test lib/*.test.mjs
+npm run build
+npm run typecheck   # tsc --noEmit
+```
+
+## Architecture
+
+**One distribution, many packages.** `pixelrag` is the single PyPI package (`pyproject.toml`
+bundles `src/pixelrag` plus `render/embed/index/serve` under `src/`). Two entry points: the
+standalone `pixelshot` (`pixelrag_render.render:main`) and the `pixelrag` umbrella
+(`pixelrag.cli:main`). Heavy ML stages are opt-in extras (`embed`, `index`, `serve`, `qdrant`, `pdf`).
+
+**CLI dispatch is lazy.** `src/pixelrag/cli.py` maps `pixelrag <stage>` to a stage module and
+imports it on demand, printing an install hint if the extra is missing. Stage 0 (render) is the
+separate `pixelshot` command, kept torch-free in core.
+
+**`pixelrag index` orchestrates by subprocess.** `index/src/pixelrag_index/pipelines.py:build`
+walks `source → ingest (render tiles) → chunk → embed → build-index`, shelling out to
+`python -m pixelrag_embed.chunk / .embed_cpu / .embed / .index` rather than importing them.
+Device routing happens here: `device: cpu|mps|auto` → `embed_cpu.py`, `cuda` → `embed.py`
+(vLLM/sglang).
+
+**Tile data model.** Each document gets a numeric position index. Tiles live in
+`<output>/tiles/<idx>.png.tiles/` containing `tile_XXXX.jpg`, a `tiles.json` manifest
+(`article_id`, `source`, `page_height`, `tiles`, `complete`), and — after chunking —
+`chunks.json` + `chunk_XXXX_YY.png` (≤1024px strips). `embed` reads `article_id` from the
+manifest, never from the directory name. `<output>/articles.json` maps position index →
+title/url/department (department = first subdirectory under a local source root, used for
+filtered search). Manifests are written atomically (`*.tmp` + `os.replace`) and `_needs_render`
+uses the `source` field to detect stale tile dirs when the source set changes between runs.
+
+**Config.** `pixelrag.yaml` is deep-merged over `DEFAULT_CONFIG` in
+`index/src/pixelrag_index/config.py` (`ingest`, `embed` sections merge key-wise). Sources are
+pluggable via the `SOURCES` registry in `index/src/pixelrag_index/sources/__init__.py`
+(`local`, `web`, `kiwix`, `pdf`); each yields `Document(id, url, path, metadata)`.
+
+**Serve.** `serve/src/pixelrag_serve/api.py` is a FastAPI app serving FAISS (default) or Qdrant
+backends, with text and image queries. It also holds `render_ondemand.py` (render retrieved
+pages at query time from a kiwix ZIM) and `zim_server*.py`. `web/agent-server.mjs` calls
+`/search`; its `/health` is liveness and `/ready` is readiness (see `deploy/README.md`).
+
+**Training and eval are separate uv projects.** `train/` LoRA-fine-tunes
+`Qwen/Qwen3-VL-Embedding-2B` (`train_contrastors.py` is the primary script); it has its own
+`train/CLAUDE.md` and `train/README.md` — read those before touching it. `eval/` is a pure HTTP
+client (reader + grader) that reproduces paper numbers; also self-contained.
+
 ## Conventions
-- Python: `uv` only (`uv add`, never `uv pip install`); work in `.venv`; commit `uv.lock`.
-- One distribution (`pixelrag`) with extras; the CLIs are `pixelshot` and `pixelrag <stage>`.
 - `main` is branch-protected — land changes via PRs.
+- Releasing is automated: `uv version X.Y.Z`, tag `vX.Y.Z`, GitHub Release → `.github/workflows/release.yml` publishes `pixelrag` to PyPI.
 
 ## Operational context (host-specific, kept out of this public repo)
 Deploy/runtime details that are specific to the live host are **not** committed here. They are

PixelRAG客户端.bat

@@ -0,0 +1,3 @@
+@echo off
+cd /d "%~dp0"
+start "" ".venv\Scripts\pythonw.exe" pixelrag_client.py

e2e_test/pixelrag.yaml

@@ -0,0 +1,9 @@
+source:
+  type: local
+  path: D:/ccproject/PixelRAG/e2e_test/test.pdf
+
+embed:
+  model: Qwen/Qwen3-VL-Embedding-2B
+  device: cpu
+
+output: D:/ccproject/PixelRAG/e2e_test/index

embed/src/pixelrag_embed/embed_cpu.py

@@ -142,18 +142,33 @@ def scan_chunks(shard_dir: str) -> list[dict]:
     return items
 
 
+def _save_checkpoint(path: Path, embeddings: np.ndarray, done: int) -> None:
+    """Atomically persist a resume checkpoint holding the first `done` rows."""
+    tmp = path.parent / (path.name + ".tmp.npz")
+    np.savez(tmp, embeddings=embeddings[:done], done=np.int64(done))
+    os.replace(tmp, path)
+
+
 def embed_items(
     items: list[dict],
     model_name: str,
     device: str = "cpu",
     instruction: str = "",
-) -> np.ndarray:
-    """Embed image items using transformers on the given device."""
+    checkpoint_dir: str | None = None,
+    checkpoint_every: int = 10,
+) -> tuple[np.ndarray, Path | None]:
+    """Embed image items using transformers on the given device.
+
+    When ``checkpoint_dir`` is set, progress is written to
+    ``<checkpoint_dir>/embed_checkpoint.npz`` every ``checkpoint_every``
+    chunks and resumed from there on restart, so an interrupted CPU embed
+    continues instead of restarting. Returns ``(embeddings, checkpoint_path)``.
+    """
     import torch
     from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
 
     device = _resolve_device(device)
-    dtype = torch.float32 if device == "cpu" else torch.float16
+    dtype = torch.bfloat16 if device == "cpu" else torch.float16
 
     logger.info("Loading model %s on %s (%s)...", model_name, device, dtype)
     processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
@@ -170,9 +185,36 @@ def embed_items(
     dim = model.config.text_config.hidden_size
     embeddings = np.zeros((len(items), dim), dtype=np.float16)
 
+    checkpoint_path = None
+    start = 0
+    if checkpoint_dir:
+        checkpoint_path = Path(checkpoint_dir) / "embed_checkpoint.npz"
+        if checkpoint_path.exists():
+            try:
+                with np.load(checkpoint_path) as ck:
+                    partial = ck["embeddings"]
+                    done = int(ck["done"])
+                if partial.shape == (done, dim) and 0 < done <= len(items):
+                    embeddings[:done] = partial
+                    start = done
+                    logger.info(
+                        "Resuming from checkpoint: %d/%d chunks already embedded",
+                        done,
+                        len(items),
+                    )
+                else:
+                    logger.warning(
+                        "Ignoring unusable checkpoint (shape=%s, done=%d)",
+                        partial.shape,
+                        done,
+                    )
+            except (OSError, ValueError, KeyError):
+                logger.warning("Ignoring unreadable checkpoint at %s", checkpoint_path)
+
     prefix = f"Instruct: {instruction}\n" if instruction else ""
 
-    for i, item in enumerate(tqdm(items, desc="Embedding")):
+    for i in tqdm(range(start, len(items)), desc="Embedding", initial=start):
+        item = items[i]
         img = Image.open(item["path"]).convert("RGB")
         img = _clamp_width(img)
 
@@ -202,15 +244,18 @@ def embed_items(
             last_idx = seq_lens - 1
             pooled = last_hidden[0, last_idx[0]]
             pooled = pooled / pooled.norm()
-            embeddings[i] = pooled.cpu().numpy().astype(np.float16)
+            embeddings[i] = pooled.cpu().float().numpy().astype(np.float16)
+
+        if checkpoint_path and (i + 1) % checkpoint_every == 0:
+            _save_checkpoint(checkpoint_path, embeddings, i + 1)
 
         # tqdm covers interactive runs but never reaches log files; when
         # stderr is not a TTY (nohup/CI/redirects) emit a periodic record so
         # long embeds stay observable after the fact.
         if (i + 1) % 100 == 0 and not sys.stderr.isatty():
             logger.info("Embedded %d/%d", i + 1, len(items))
 
-    return embeddings
+    return embeddings, checkpoint_path
 
 
 def main():
@@ -230,6 +275,15 @@ def main():
         "--instruction", default="", help="Instruction prefix for queries"
     )
     parser.add_argument("--limit", type=int, default=None, help="Max chunks to embed")
+    parser.add_argument(
+        "--checkpoint-every",
+        type=int,
+        default=1,
+        help="Save a resume checkpoint every N chunks (default: 1). Embedding "
+        "keeps all progress in memory; only the checkpoint survives a kill, so "
+        "this bounds what an OOM kill costs. The write is a few MB at most, so "
+        "saving after every chunk is cheap.",
+    )
     args = parser.parse_args()
 
     os.makedirs(args.output_dir, exist_ok=True)
@@ -244,8 +298,20 @@ def main():
         items = items[: args.limit]
     else:
         logger.info("Found %d chunks to embed", len(items))
-    embeddings = embed_items(
-        items, args.model, device=args.device, instruction=args.instruction
+
+    checkpoint_every = max(1, args.checkpoint_every)
+    logger.info(
+        "Checkpoint every %d chunk(s) -> %s (an OOM kill costs at most that much)",
+        checkpoint_every,
+        Path(args.output_dir) / "embed_checkpoint.npz",
+    )
+    embeddings, checkpoint_path = embed_items(
+        items,
+        args.model,
+        device=args.device,
+        instruction=args.instruction,
+        checkpoint_dir=args.output_dir,
+        checkpoint_every=checkpoint_every,
     )
 
     output_path = Path(args.output_dir) / "shard_000.npz"
@@ -259,6 +325,9 @@ def main():
         tile_heights=np.array([it["height"] for it in items], dtype=np.int32),
     )
 
+    if checkpoint_path and checkpoint_path.exists():
+        checkpoint_path.unlink()
+
     logger.info("Saved %d embeddings to %s", len(items), output_path)
 
 

index/src/pixelrag_index/pipelines.py

@@ -213,7 +213,7 @@ def _repl(m: re.Match) -> str:
                 if not _needs_render(tiles_dir, idx, doc):
                     continue
                 src_path = Path(doc.path)
-                content = src_path.read_text(errors="replace")
+                content = src_path.read_text(encoding="utf-8", errors="replace")
                 ext = (doc.metadata or {}).get("extension", src_path.suffix.lower())
                 if ext == ".md":
                     body = md_lib.markdown(
@@ -224,7 +224,7 @@ def _repl(m: re.Match) -> str:
                 body = _resolve_relative_paths(body, str(src_path.parent))
                 html_content = f"{_HTML_TEMPLATE}{body}</body></html>"
                 html_path = Path(tmp_dir) / f"{idx}.html"
-                html_path.write_text(html_content)
+                html_path.write_text(html_content, encoding="utf-8")
                 text_urls.append(f"file://{html_path.resolve()}")
                 text_stems.append(str(idx))
 
@@ -349,7 +349,7 @@ def _repl(m: re.Match) -> str:
         article_entries.append(
             {"title": title, "url": url, "department": _department_of(a, source_root)}
         )
-    with open(articles_path, "w") as f:
+    with open(articles_path, "w", encoding="utf-8") as f:
         json.dump(article_entries, f)
     logger.info(
         "  Saved %d article mappings to %s", len(article_entries), articles_path

pixelrag.yaml

@@ -0,0 +1,9 @@
+source:
+  type: local
+  path: D:/ccproject/PixelRAG/papers
+
+embed:
+  model: Qwen/Qwen3-VL-Embedding-2B
+  device: cpu
+
+output: D:/ccproject/PixelRAG/knowledge_index