valcyon 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- valcyon-0.1.0/.gitignore +10 -0
- valcyon-0.1.0/LICENSE +0 -0
- valcyon-0.1.0/PKG-INFO +24 -0
- valcyon-0.1.0/README.md +7 -0
- valcyon-0.1.0/pyproject.toml +32 -0
- valcyon-0.1.0/src/valcyon/__init__.py +1 -0
- valcyon-0.1.0/src/valcyon/adapters/__init__.py +0 -0
- valcyon-0.1.0/src/valcyon/adapters/base.py +7 -0
- valcyon-0.1.0/src/valcyon/adapters/errors.py +2 -0
- valcyon-0.1.0/src/valcyon/adapters/mineru.py +76 -0
- valcyon-0.1.0/src/valcyon/adapters/mineru_cloud.py +138 -0
- valcyon-0.1.0/src/valcyon/chunking.py +135 -0
- valcyon-0.1.0/src/valcyon/cli.py +80 -0
- valcyon-0.1.0/src/valcyon/config.py +27 -0
- valcyon-0.1.0/src/valcyon/embedders/__init__.py +0 -0
- valcyon-0.1.0/src/valcyon/embedders/base.py +14 -0
- valcyon-0.1.0/src/valcyon/embedders/bge_m3.py +24 -0
- valcyon-0.1.0/src/valcyon/embedders/local.py +0 -0
- valcyon-0.1.0/src/valcyon/eval/__init__.py +0 -0
- valcyon-0.1.0/src/valcyon/eval/metrics.py +40 -0
- valcyon-0.1.0/src/valcyon/eval/runner.py +0 -0
- valcyon-0.1.0/src/valcyon/ids.py +20 -0
- valcyon-0.1.0/src/valcyon/models.py +49 -0
- valcyon-0.1.0/src/valcyon/pipeline.py +46 -0
- valcyon-0.1.0/src/valcyon/rerankers/__init__.py +18 -0
- valcyon-0.1.0/src/valcyon/rerankers/bge.py +23 -0
- valcyon-0.1.0/src/valcyon/retrieve/__init__.py +0 -0
- valcyon-0.1.0/src/valcyon/retrieve/dense.py +53 -0
- valcyon-0.1.0/src/valcyon/retrieve/fusion.py +0 -0
- valcyon-0.1.0/src/valcyon/retrieve/hybrid.py +29 -0
- valcyon-0.1.0/src/valcyon/retrieve/keyword.py +41 -0
- valcyon-0.1.0/src/valcyon/retrieve/rerank.py +23 -0
- valcyon-0.1.0/src/valcyon/retrieve/search.py +0 -0
- valcyon-0.1.0/src/valcyon/retrieve/tree_ops.py +0 -0
- valcyon-0.1.0/src/valcyon/store/__init__.py +3 -0
- valcyon-0.1.0/src/valcyon/store/keyword.py +41 -0
- valcyon-0.1.0/src/valcyon/store/sqlite.py +181 -0
- valcyon-0.1.0/src/valcyon/store/vectors.py +0 -0
- valcyon-0.1.0/src/valcyon/tree/__init__.py +0 -0
- valcyon-0.1.0/src/valcyon/tree/builder.py +52 -0
- valcyon-0.1.0/src/valcyon/tree/levels.py +60 -0
- valcyon-0.1.0/src/valcyon/tree/promote.py +61 -0
- valcyon-0.1.0/src/valcyon/tree/restructure.py +89 -0
- valcyon-0.1.0/src/valcyon/tree/validate.py +97 -0
valcyon-0.1.0/.gitignore
ADDED
valcyon-0.1.0/LICENSE
ADDED
|
File without changes
|
valcyon-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: valcyon
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Accurate, traceable retrieval for PDFs.
|
|
5
|
+
Author: Osama AKESKOUS
|
|
6
|
+
License: MIT
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Requires-Python: >=3.10
|
|
9
|
+
Requires-Dist: numpy
|
|
10
|
+
Provides-Extra: cloud
|
|
11
|
+
Requires-Dist: requests; extra == 'cloud'
|
|
12
|
+
Provides-Extra: dev
|
|
13
|
+
Requires-Dist: pytest; extra == 'dev'
|
|
14
|
+
Provides-Extra: local
|
|
15
|
+
Requires-Dist: sentence-transformers; extra == 'local'
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# Valcyon
|
|
19
|
+
|
|
20
|
+
<p align="center">
|
|
21
|
+
<img src="assets/kingfisher-logo.png" alt="pdftree logo" width="220">
|
|
22
|
+
</p>
|
|
23
|
+
|
|
24
|
+
Deterministic hierarchical PDF extraction and retrieval toolkit. This project is in early development.
|
valcyon-0.1.0/README.md
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "valcyon"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Accurate, traceable retrieval for PDFs."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [{ name = "Osama AKESKOUS" }]
|
|
13
|
+
dependencies = ["numpy"]
|
|
14
|
+
|
|
15
|
+
[project.optional-dependencies]
|
|
16
|
+
local = ["sentence-transformers"]
|
|
17
|
+
cloud = ["requests"]
|
|
18
|
+
dev = ["pytest"]
|
|
19
|
+
|
|
20
|
+
[project.scripts]
|
|
21
|
+
valcyon = "valcyon.cli:main"
|
|
22
|
+
|
|
23
|
+
[tool.hatch.build.targets.wheel]
|
|
24
|
+
packages = ["src/valcyon"]
|
|
25
|
+
|
|
26
|
+
[tool.hatch.build.targets.sdist]
|
|
27
|
+
include = [
|
|
28
|
+
"/src/valcyon",
|
|
29
|
+
"/README.md",
|
|
30
|
+
"/LICENSE",
|
|
31
|
+
"/pyproject.toml",
|
|
32
|
+
]
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.1.0"
|
|
File without changes
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
|
|
5
|
+
from valcyon.models import Element
|
|
6
|
+
|
|
7
|
+
# block type in middle.json -> element type in Valcyon
|
|
8
|
+
_TYPE_MAP = {
|
|
9
|
+
"title": "heading",
|
|
10
|
+
"text": "text",
|
|
11
|
+
"list": "text",
|
|
12
|
+
"index": "text",
|
|
13
|
+
"code": "text",
|
|
14
|
+
"table": "table", # to verify on a real file
|
|
15
|
+
"image": "figure", # to verify on a real file
|
|
16
|
+
"interline_equation": "equation", # to verify on a real file
|
|
17
|
+
}
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _span_text(span: dict) -> str:
|
|
21
|
+
"""Text of one span. Tables/equations may use other keys than 'content'."""
|
|
22
|
+
return span.get("content") or span.get("html") or span.get("latex") or ""
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def _block_text(block: dict) -> str:
|
|
26
|
+
"""Text of a block: its lines joined by spaces, or its sub-blocks if nested."""
|
|
27
|
+
if "lines" in block:
|
|
28
|
+
lines = []
|
|
29
|
+
for line in block["lines"]:
|
|
30
|
+
parts = [_span_text(s) for s in line.get("spans", [])]
|
|
31
|
+
lines.append("".join(parts).strip())
|
|
32
|
+
return " ".join(x for x in lines if x)
|
|
33
|
+
# image/table blocks nest their content in sub-blocks (caption, body...)
|
|
34
|
+
sub = [_block_text(b) for b in block.get("blocks", [])]
|
|
35
|
+
return "\n".join(x for x in sub if x)
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
def _bbox(block: dict):
|
|
39
|
+
b = block.get("bbox")
|
|
40
|
+
if b and len(b) == 4:
|
|
41
|
+
return tuple(float(x) for x in b)
|
|
42
|
+
return None
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class MinerUParser:
|
|
46
|
+
"""Reads MinerU's middle.json.
|
|
47
|
+
|
|
48
|
+
Assumed shape:
|
|
49
|
+
{"pdf_info": [ {"page_idx": 0, "preproc_blocks": [ block, ... ]}, ... ]}
|
|
50
|
+
block = {"type", "bbox", "index", "level" (titles), "lines": [...]}
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
def parse(self, source: str) -> list[Element]:
|
|
54
|
+
with open(source, "r", encoding="utf-8") as f:
|
|
55
|
+
data = json.load(f)
|
|
56
|
+
|
|
57
|
+
elements: list[Element] = []
|
|
58
|
+
for page_pos, page in enumerate(data.get("pdf_info", [])):
|
|
59
|
+
page_no = int(page.get("page_idx", page_pos)) + 1 # 1-based
|
|
60
|
+
blocks = page.get("preproc_blocks", [])
|
|
61
|
+
# explicit sort key -> same file always gives same order
|
|
62
|
+
ordered = sorted(enumerate(blocks), key=lambda t: (t[1].get("index", t[0]), t[0]))
|
|
63
|
+
|
|
64
|
+
for _, block in ordered:
|
|
65
|
+
kind = _TYPE_MAP.get(block.get("type"))
|
|
66
|
+
if kind is None:
|
|
67
|
+
continue # unknown block type: ignore for now
|
|
68
|
+
text = _block_text(block)
|
|
69
|
+
if not text:
|
|
70
|
+
continue
|
|
71
|
+
level = int(block.get("level", 1)) if kind == "heading" else None
|
|
72
|
+
elements.append(Element(
|
|
73
|
+
type=kind, text=text, page=page_no, order=len(elements),
|
|
74
|
+
bbox=_bbox(block), raw_level=level,
|
|
75
|
+
))
|
|
76
|
+
return elements
|
|
@@ -0,0 +1,138 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import io
|
|
4
|
+
import os
|
|
5
|
+
import time
|
|
6
|
+
import zipfile
|
|
7
|
+
from pathlib import Path, PurePosixPath
|
|
8
|
+
from valcyon.config import get_setting
|
|
9
|
+
|
|
10
|
+
from valcyon.adapters.errors import MinerUError
|
|
11
|
+
|
|
12
|
+
API = "https://mineru.net/api/v4"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class CloudMinerURunner:
|
|
16
|
+
"""Sends a PDF to the MinerU cloud API and saves its middle.json.
|
|
17
|
+
|
|
18
|
+
Flow (from MinerU's docs):
|
|
19
|
+
1. POST /file-urls/batch -> batch_id + presigned upload URL
|
|
20
|
+
2. PUT <upload url> -> the PDF bytes (no Content-Type header)
|
|
21
|
+
3. GET /extract-results/batch/{batch_id} until state == "done"
|
|
22
|
+
4. download full_zip_url -> the cloud names middle.json "layout.json"
|
|
23
|
+
"""
|
|
24
|
+
|
|
25
|
+
def __init__(self, token: str | None = None, model_version: str = "pipeline",
|
|
26
|
+
language: str = "ch", is_ocr: bool = False,
|
|
27
|
+
timeout: float = 600, interval: float = 3,
|
|
28
|
+
base_url: str = API, session=None, sleep=time.sleep):
|
|
29
|
+
self.token = token or get_setting("MINERU_TOKEN")
|
|
30
|
+
if not self.token:
|
|
31
|
+
raise MinerUError("No token. Set the MINERU_TOKEN environment variable.")
|
|
32
|
+
self.model_version = model_version
|
|
33
|
+
self.language = language
|
|
34
|
+
self.is_ocr = is_ocr
|
|
35
|
+
self.timeout = timeout
|
|
36
|
+
self.interval = interval
|
|
37
|
+
self.base_url = base_url.rstrip("/")
|
|
38
|
+
self.sleep = sleep
|
|
39
|
+
self._session = session
|
|
40
|
+
|
|
41
|
+
# -- helpers ---------------------------------------------------------
|
|
42
|
+
def _http(self):
|
|
43
|
+
if self._session is None:
|
|
44
|
+
try:
|
|
45
|
+
import requests
|
|
46
|
+
except ImportError:
|
|
47
|
+
raise MinerUError("Install the cloud extra: pip install valcyon[cloud]")
|
|
48
|
+
self._session = requests.Session()
|
|
49
|
+
return self._session
|
|
50
|
+
|
|
51
|
+
@staticmethod
|
|
52
|
+
def _data(resp) -> dict:
|
|
53
|
+
"""Check an API response and return its 'data' part."""
|
|
54
|
+
if resp.status_code != 200:
|
|
55
|
+
raise MinerUError(f"HTTP {resp.status_code} from MinerU")
|
|
56
|
+
body = resp.json()
|
|
57
|
+
if body.get("code") != 0:
|
|
58
|
+
raise MinerUError(f"MinerU error {body.get('code')}: {body.get('msg')}")
|
|
59
|
+
return body.get("data") or {}
|
|
60
|
+
|
|
61
|
+
@staticmethod
|
|
62
|
+
def _pick(result, file_name: str) -> dict | None:
|
|
63
|
+
"""extract_result is a list in the docs; accept a dict too."""
|
|
64
|
+
if isinstance(result, dict):
|
|
65
|
+
return result
|
|
66
|
+
for row in result or []:
|
|
67
|
+
if row.get("file_name") == file_name:
|
|
68
|
+
return row
|
|
69
|
+
return (result or [None])[0]
|
|
70
|
+
|
|
71
|
+
@staticmethod
|
|
72
|
+
def _find_middle(names: list[str]) -> str | None:
|
|
73
|
+
found = [n for n in names
|
|
74
|
+
if PurePosixPath(n).name in ("layout.json",)
|
|
75
|
+
or PurePosixPath(n).name.endswith("middle.json")]
|
|
76
|
+
return sorted(found)[0] if found else None
|
|
77
|
+
|
|
78
|
+
# -- main ------------------------------------------------------------
|
|
79
|
+
def run(self, pdf_path: str, out_dir: str) -> str:
|
|
80
|
+
pdf = Path(pdf_path)
|
|
81
|
+
if not pdf.is_file():
|
|
82
|
+
raise MinerUError(f"PDF not found: {pdf}")
|
|
83
|
+
out = Path(out_dir)
|
|
84
|
+
out.mkdir(parents=True, exist_ok=True)
|
|
85
|
+
http = self._http()
|
|
86
|
+
auth = {"Authorization": f"Bearer {self.token}"}
|
|
87
|
+
|
|
88
|
+
# 1. ask for an upload URL
|
|
89
|
+
body = {
|
|
90
|
+
"files": [{"name": pdf.name, "is_ocr": self.is_ocr}],
|
|
91
|
+
"model_version": self.model_version,
|
|
92
|
+
"language": self.language,
|
|
93
|
+
"enable_table": True,
|
|
94
|
+
"enable_formula": True,
|
|
95
|
+
}
|
|
96
|
+
data = self._data(http.post(f"{self.base_url}/file-urls/batch",
|
|
97
|
+
headers={**auth, "Content-Type": "application/json"},
|
|
98
|
+
json=body, timeout=60))
|
|
99
|
+
batch_id = data.get("batch_id")
|
|
100
|
+
urls = data.get("file_urls") or []
|
|
101
|
+
if not batch_id or not urls:
|
|
102
|
+
raise MinerUError("MinerU did not return a batch_id and an upload URL")
|
|
103
|
+
|
|
104
|
+
# 2. upload the PDF (deliberately no headers: the URL is signed)
|
|
105
|
+
with open(pdf, "rb") as f:
|
|
106
|
+
put = http.put(urls[0], data=f, timeout=300)
|
|
107
|
+
if put.status_code not in (200, 201):
|
|
108
|
+
raise MinerUError(f"Upload failed: HTTP {put.status_code}")
|
|
109
|
+
|
|
110
|
+
# 3. poll until done
|
|
111
|
+
deadline = time.monotonic() + self.timeout
|
|
112
|
+
while True:
|
|
113
|
+
data = self._data(http.get(f"{self.base_url}/extract-results/batch/{batch_id}",
|
|
114
|
+
headers=auth, timeout=60))
|
|
115
|
+
row = self._pick(data.get("extract_result"), pdf.name) or {}
|
|
116
|
+
state = row.get("state")
|
|
117
|
+
if state == "done":
|
|
118
|
+
zip_url = row.get("full_zip_url")
|
|
119
|
+
if not zip_url:
|
|
120
|
+
raise MinerUError("Task is done but has no full_zip_url")
|
|
121
|
+
break
|
|
122
|
+
if state == "failed":
|
|
123
|
+
raise MinerUError(f"MinerU failed: {row.get('err_msg') or 'unknown reason'}")
|
|
124
|
+
if time.monotonic() >= deadline:
|
|
125
|
+
raise MinerUError(f"Timed out (state={state}). batch_id={batch_id}")
|
|
126
|
+
self.sleep(self.interval)
|
|
127
|
+
|
|
128
|
+
# 4. download the ZIP and keep only the middle.json inside it
|
|
129
|
+
resp = http.get(zip_url, timeout=300)
|
|
130
|
+
if resp.status_code != 200:
|
|
131
|
+
raise MinerUError(f"Download failed: HTTP {resp.status_code}")
|
|
132
|
+
with zipfile.ZipFile(io.BytesIO(resp.content)) as z:
|
|
133
|
+
member = self._find_middle(z.namelist())
|
|
134
|
+
if member is None:
|
|
135
|
+
raise MinerUError(f"No layout.json / middle.json in ZIP: {z.namelist()}")
|
|
136
|
+
target = out / f"{pdf.stem}_middle.json"
|
|
137
|
+
target.write_bytes(z.read(member))
|
|
138
|
+
return str(target)
|
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
from dataclasses import dataclass
|
|
5
|
+
from typing import Callable
|
|
6
|
+
|
|
7
|
+
from valcyon.ids import sha
|
|
8
|
+
from valcyon.models import BBox, Node
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def approx_tokens(text: str) -> int:
|
|
12
|
+
"""Cheap token estimate: whitespace-separated words."""
|
|
13
|
+
return len(text.split())
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass
|
|
17
|
+
class Chunk:
|
|
18
|
+
id: str
|
|
19
|
+
doc_id: str
|
|
20
|
+
section_id: str
|
|
21
|
+
text: str
|
|
22
|
+
breadcrumb: str # "Expérience > Assistant projet – Stage"
|
|
23
|
+
path: list[str]
|
|
24
|
+
page_start: int
|
|
25
|
+
page_end: int
|
|
26
|
+
bboxes: list[tuple[int, BBox]] # (page, bbox) of each source block
|
|
27
|
+
chunk_type: str # text | table | figure | equation
|
|
28
|
+
order: int # position inside its section
|
|
29
|
+
token_count: int
|
|
30
|
+
|
|
31
|
+
@property
|
|
32
|
+
def embed_text(self) -> str:
|
|
33
|
+
"""What we will embed: the breadcrumb gives the chunk its context."""
|
|
34
|
+
return f"{self.breadcrumb}\n\n{self.text}" if self.breadcrumb else self.text
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def _split_words(sentence: str, max_tokens: int, count: Callable[[str], int]) -> list[str]:
|
|
38
|
+
"""Last resort for a single sentence longer than max_tokens."""
|
|
39
|
+
out: list[str] = []
|
|
40
|
+
cur: list[str] = []
|
|
41
|
+
for w in sentence.split():
|
|
42
|
+
if cur and count(" ".join(cur + [w])) > max_tokens:
|
|
43
|
+
out.append(" ".join(cur))
|
|
44
|
+
cur = []
|
|
45
|
+
cur.append(w)
|
|
46
|
+
if cur:
|
|
47
|
+
out.append(" ".join(cur))
|
|
48
|
+
return out
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def _split_text(text: str, max_tokens: int, overlap: int,
|
|
52
|
+
count: Callable[[str], int]) -> list[str]:
|
|
53
|
+
"""Split on sentence boundaries; each piece repeats up to `overlap` tokens of the last."""
|
|
54
|
+
pieces: list[str] = []
|
|
55
|
+
for s in re.split(r"(?<=[.!?])\s+", text.strip()):
|
|
56
|
+
if not s:
|
|
57
|
+
continue
|
|
58
|
+
pieces.extend([s] if count(s) <= max_tokens else _split_words(s, max_tokens, count))
|
|
59
|
+
|
|
60
|
+
chunks: list[str] = []
|
|
61
|
+
cur: list[str] = []
|
|
62
|
+
for p in pieces:
|
|
63
|
+
if cur and count(" ".join(cur + [p])) > max_tokens:
|
|
64
|
+
chunks.append(" ".join(cur))
|
|
65
|
+
carry: list[str] = []
|
|
66
|
+
total = 0
|
|
67
|
+
for q in reversed(cur): # keep trailing pieces as overlap
|
|
68
|
+
c = count(q)
|
|
69
|
+
if total + c > overlap:
|
|
70
|
+
break
|
|
71
|
+
carry.insert(0, q)
|
|
72
|
+
total += c
|
|
73
|
+
cur = carry
|
|
74
|
+
if cur and count(" ".join(cur + [p])) > max_tokens:
|
|
75
|
+
cur = [] # overlap would not fit: drop it
|
|
76
|
+
cur.append(p)
|
|
77
|
+
if cur:
|
|
78
|
+
chunks.append(" ".join(cur))
|
|
79
|
+
return chunks
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def _chunk_section(section: Node, count: Callable[[str], int],
|
|
83
|
+
max_tokens: int, overlap: int) -> list[Chunk]:
|
|
84
|
+
out: list[Chunk] = []
|
|
85
|
+
breadcrumb = " > ".join(section.path)
|
|
86
|
+
buffer: list[Node] = []
|
|
87
|
+
|
|
88
|
+
def emit(text: str, blocks: list[Node], kind: str) -> None:
|
|
89
|
+
idx = len(out)
|
|
90
|
+
pages = [b.page_start for b in blocks]
|
|
91
|
+
out.append(Chunk(
|
|
92
|
+
id=sha(f"{section.doc_id}|{section.id}|{idx}|{sha(text)}"),
|
|
93
|
+
doc_id=section.doc_id, section_id=section.id, text=text,
|
|
94
|
+
breadcrumb=breadcrumb, path=list(section.path),
|
|
95
|
+
page_start=min(pages), page_end=max(pages),
|
|
96
|
+
bboxes=[(b.page_start, b.bbox) for b in blocks if b.bbox is not None],
|
|
97
|
+
chunk_type=kind, order=idx, token_count=count(text),
|
|
98
|
+
))
|
|
99
|
+
|
|
100
|
+
def flush() -> None:
|
|
101
|
+
if buffer:
|
|
102
|
+
emit("\n".join(b.text for b in buffer), buffer, "text")
|
|
103
|
+
buffer.clear()
|
|
104
|
+
|
|
105
|
+
for block in (c for c in section.children if c.type == "block"):
|
|
106
|
+
text = block.text or ""
|
|
107
|
+
if not text.strip():
|
|
108
|
+
continue
|
|
109
|
+
kind = block.element_type or "text"
|
|
110
|
+
|
|
111
|
+
if kind != "text": # tables, figures, equations: one chunk, never split
|
|
112
|
+
flush()
|
|
113
|
+
emit(text, [block], kind)
|
|
114
|
+
elif count(text) > max_tokens: # too big for one chunk: split it
|
|
115
|
+
flush()
|
|
116
|
+
for piece in _split_text(text, max_tokens, overlap, count):
|
|
117
|
+
emit(piece, [block], "text")
|
|
118
|
+
else: # merge small blocks together
|
|
119
|
+
if buffer and count("\n".join([b.text for b in buffer] + [text])) > max_tokens:
|
|
120
|
+
flush()
|
|
121
|
+
buffer.append(block)
|
|
122
|
+
flush()
|
|
123
|
+
return out
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def chunk_tree(root: Node, max_tokens: int = 400, overlap_tokens: int = 40,
|
|
127
|
+
count_tokens: Callable[[str], int] = approx_tokens) -> list[Chunk]:
|
|
128
|
+
"""Turn a tree into chunks, in document order. Chunks never cross a section."""
|
|
129
|
+
if max_tokens <= 0 or not (0 <= overlap_tokens < max_tokens):
|
|
130
|
+
raise ValueError("need max_tokens > 0 and 0 <= overlap_tokens < max_tokens")
|
|
131
|
+
chunks: list[Chunk] = []
|
|
132
|
+
for node in root.walk():
|
|
133
|
+
if node.type != "block": # the document root holds any preamble blocks
|
|
134
|
+
chunks.extend(_chunk_section(node, count_tokens, max_tokens, overlap_tokens))
|
|
135
|
+
return chunks
|
|
@@ -0,0 +1,80 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import argparse
|
|
4
|
+
import sys
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
|
|
7
|
+
DEFAULT_DB = "valcyon.db"
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def _ingest(args: argparse.Namespace) -> int:
|
|
11
|
+
from valcyon.embedders.bge_m3 import BgeM3Embedder
|
|
12
|
+
from valcyon.pipeline import ingest_parsed
|
|
13
|
+
from valcyon.store import Store
|
|
14
|
+
|
|
15
|
+
src = Path(args.source)
|
|
16
|
+
if src.suffix.lower() == ".pdf":
|
|
17
|
+
from valcyon.adapters.mineru_cloud import CloudMinerURunner
|
|
18
|
+
|
|
19
|
+
out_dir = src.parent / "out"
|
|
20
|
+
CloudMinerURunner().run(str(src), str(out_dir))
|
|
21
|
+
found = sorted(out_dir.rglob("*_middle.json"))
|
|
22
|
+
if not found:
|
|
23
|
+
print(f"No *_middle.json found in {out_dir}", file=sys.stderr)
|
|
24
|
+
return 1
|
|
25
|
+
middle = found[0]
|
|
26
|
+
else:
|
|
27
|
+
middle = src
|
|
28
|
+
|
|
29
|
+
stats = ingest_parsed(str(middle), Store(args.db), BgeM3Embedder())
|
|
30
|
+
issues = stats.pop("issues")
|
|
31
|
+
print(stats)
|
|
32
|
+
for code, message in issues:
|
|
33
|
+
print(f"[{code}] {message}")
|
|
34
|
+
return 0
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def _query(args: argparse.Namespace) -> int:
|
|
38
|
+
from valcyon.embedders.bge_m3 import BgeM3Embedder
|
|
39
|
+
from valcyon.retrieve.dense import DenseRetriever
|
|
40
|
+
from valcyon.retrieve.hybrid import HybridRetriever
|
|
41
|
+
from valcyon.retrieve.keyword import KeywordRetriever
|
|
42
|
+
from valcyon.store import Store
|
|
43
|
+
|
|
44
|
+
store = Store(args.db)
|
|
45
|
+
dense = DenseRetriever(store, BgeM3Embedder())
|
|
46
|
+
retriever = HybridRetriever(dense, KeywordRetriever(store))
|
|
47
|
+
for result in retriever.search(args.question, k=args.k):
|
|
48
|
+
chunk = result.chunk
|
|
49
|
+
print(f"p.{chunk.page_start} {chunk.breadcrumb}")
|
|
50
|
+
print(" ", chunk.text[:120].replace("\n", " "))
|
|
51
|
+
return 0
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def main(argv: list[str] | None = None) -> int:
|
|
55
|
+
parser = argparse.ArgumentParser(
|
|
56
|
+
prog="valcyon",
|
|
57
|
+
description="Structure-aware retrieval for PDFs (early development).",
|
|
58
|
+
)
|
|
59
|
+
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
60
|
+
|
|
61
|
+
ingest = subparsers.add_parser(
|
|
62
|
+
"ingest",
|
|
63
|
+
help="parse, chunk, index a PDF or a MinerU *_middle.json",
|
|
64
|
+
)
|
|
65
|
+
ingest.add_argument("source")
|
|
66
|
+
ingest.add_argument("--db", default=DEFAULT_DB)
|
|
67
|
+
ingest.set_defaults(func=_ingest)
|
|
68
|
+
|
|
69
|
+
query = subparsers.add_parser("query", help="search an indexed database")
|
|
70
|
+
query.add_argument("question")
|
|
71
|
+
query.add_argument("--db", default=DEFAULT_DB)
|
|
72
|
+
query.add_argument("-k", type=int, default=3)
|
|
73
|
+
query.set_defaults(func=_query)
|
|
74
|
+
|
|
75
|
+
args = parser.parse_args(argv)
|
|
76
|
+
return args.func(args)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
if __name__ == "__main__":
|
|
80
|
+
raise SystemExit(main())
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def _parse_env(path: Path) -> dict[str, str]:
|
|
8
|
+
"""Read KEY=VALUE lines. Ignores blanks and # comments, strips quotes."""
|
|
9
|
+
values: dict[str, str] = {}
|
|
10
|
+
if not path.is_file():
|
|
11
|
+
return values
|
|
12
|
+
# utf-8-sig: Notepad on Windows may add an invisible BOM at the start
|
|
13
|
+
for line in path.read_text(encoding="utf-8-sig").splitlines():
|
|
14
|
+
line = line.strip()
|
|
15
|
+
if not line or line.startswith("#") or "=" not in line:
|
|
16
|
+
continue
|
|
17
|
+
key, _, val = line.partition("=")
|
|
18
|
+
val = val.strip()
|
|
19
|
+
if len(val) >= 2 and val[0] == val[-1] and val[0] in "\"'":
|
|
20
|
+
val = val[1:-1]
|
|
21
|
+
values[key.strip()] = val
|
|
22
|
+
return values
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def get_setting(name: str, env_file: str = ".env") -> str | None:
|
|
26
|
+
"""A real environment variable wins; otherwise look in the .env file."""
|
|
27
|
+
return os.environ.get(name) or _parse_env(Path(env_file)).get(name) or None
|
|
File without changes
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Protocol
|
|
4
|
+
|
|
5
|
+
import numpy as np
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class Embedder(Protocol):
|
|
9
|
+
name: str # stored in the index: vectors from different models must never be mixed
|
|
10
|
+
|
|
11
|
+
def embed(self, texts: list[str], kind: str = "doc") -> np.ndarray:
|
|
12
|
+
"""Return a float32 array (len(texts), dim), L2-normalized,
|
|
13
|
+
so cosine similarity is a plain dot product. kind is "doc" or "query"."""
|
|
14
|
+
...
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import numpy as np
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class BgeM3Embedder:
|
|
7
|
+
"""Local multilingual embedder. bge-m3 needs no query/document prefix,
|
|
8
|
+
so `kind` is accepted for the interface but unused."""
|
|
9
|
+
|
|
10
|
+
def __init__(self, model_name: str = "BAAI/bge-m3", device: str | None = None) -> None:
|
|
11
|
+
self.name = model_name
|
|
12
|
+
self.device = device
|
|
13
|
+
self._model = None
|
|
14
|
+
|
|
15
|
+
def _load(self):
|
|
16
|
+
if self._model is None: # lazy: the import and download happen once
|
|
17
|
+
from sentence_transformers import SentenceTransformer
|
|
18
|
+
self._model = SentenceTransformer(self.name, device=self.device)
|
|
19
|
+
return self._model
|
|
20
|
+
|
|
21
|
+
def embed(self, texts: list[str], kind: str = "doc") -> np.ndarray:
|
|
22
|
+
vecs = self._load().encode(texts, normalize_embeddings=True,
|
|
23
|
+
convert_to_numpy=True, batch_size=16)
|
|
24
|
+
return np.asarray(vecs, dtype=np.float32)
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import time
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def is_relevant(chunk_text: str, needles: list[str]) -> bool:
|
|
7
|
+
low = chunk_text.lower()
|
|
8
|
+
return any(n.lower() in low for n in needles)
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def evaluate(retriever, cases: list[dict], k: int = 5) -> dict:
|
|
12
|
+
"""cases: [{"query": str, "relevant": [str, ...]}].
|
|
13
|
+
hit_rate: share of queries with a relevant chunk in the top k.
|
|
14
|
+
hit1: share of queries whose first result is relevant.
|
|
15
|
+
mrr: mean of 1/rank of the first relevant chunk (0 if not in the top k)."""
|
|
16
|
+
hits = 0
|
|
17
|
+
first = 0
|
|
18
|
+
rr_sum = 0.0
|
|
19
|
+
failures: list[str] = []
|
|
20
|
+
ranks: list = []
|
|
21
|
+
t0 = time.perf_counter()
|
|
22
|
+
for case in cases:
|
|
23
|
+
results = retriever.search(case["query"], k=k)
|
|
24
|
+
rank = next((i for i, r in enumerate(results, 1)
|
|
25
|
+
if is_relevant(r.chunk.embed_text, case["relevant"])), None)
|
|
26
|
+
ranks.append((case["query"], rank))
|
|
27
|
+
if rank is None:
|
|
28
|
+
failures.append(case["query"])
|
|
29
|
+
else:
|
|
30
|
+
hits += 1
|
|
31
|
+
rr_sum += 1 / rank
|
|
32
|
+
first += rank == 1
|
|
33
|
+
elapsed = time.perf_counter() - t0
|
|
34
|
+
n = len(cases)
|
|
35
|
+
return {"n": n, "k": k,
|
|
36
|
+
"hit_rate": hits / n if n else 0.0,
|
|
37
|
+
"hit1": first / n if n else 0.0,
|
|
38
|
+
"mrr": rr_sum / n if n else 0.0,
|
|
39
|
+
"ms_per_query": 1000 * elapsed / n if n else 0.0,
|
|
40
|
+
"failures": failures, "ranks": ranks}
|
|
File without changes
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
import hashlib
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
def sha(text: str) -> str:
|
|
5
|
+
"""Short, stable fingerprint of a string."""
|
|
6
|
+
return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def doc_id_from_file(path: str) -> str:
|
|
10
|
+
"""Fingerprint of a file's bytes: same file -> same id."""
|
|
11
|
+
h = hashlib.sha256()
|
|
12
|
+
with open(path, "rb") as f:
|
|
13
|
+
for block in iter(lambda: f.read(1 << 20), b""):
|
|
14
|
+
h.update(block)
|
|
15
|
+
return h.hexdigest()[:16]
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def node_id(doc_id: str, position: str) -> str:
|
|
19
|
+
"""Id of a tree node, from its document and its position like '0.2.1'."""
|
|
20
|
+
return sha(f"{doc_id}|{position}")
|