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.
Files changed (44) hide show
  1. valcyon-0.1.0/.gitignore +10 -0
  2. valcyon-0.1.0/LICENSE +0 -0
  3. valcyon-0.1.0/PKG-INFO +24 -0
  4. valcyon-0.1.0/README.md +7 -0
  5. valcyon-0.1.0/pyproject.toml +32 -0
  6. valcyon-0.1.0/src/valcyon/__init__.py +1 -0
  7. valcyon-0.1.0/src/valcyon/adapters/__init__.py +0 -0
  8. valcyon-0.1.0/src/valcyon/adapters/base.py +7 -0
  9. valcyon-0.1.0/src/valcyon/adapters/errors.py +2 -0
  10. valcyon-0.1.0/src/valcyon/adapters/mineru.py +76 -0
  11. valcyon-0.1.0/src/valcyon/adapters/mineru_cloud.py +138 -0
  12. valcyon-0.1.0/src/valcyon/chunking.py +135 -0
  13. valcyon-0.1.0/src/valcyon/cli.py +80 -0
  14. valcyon-0.1.0/src/valcyon/config.py +27 -0
  15. valcyon-0.1.0/src/valcyon/embedders/__init__.py +0 -0
  16. valcyon-0.1.0/src/valcyon/embedders/base.py +14 -0
  17. valcyon-0.1.0/src/valcyon/embedders/bge_m3.py +24 -0
  18. valcyon-0.1.0/src/valcyon/embedders/local.py +0 -0
  19. valcyon-0.1.0/src/valcyon/eval/__init__.py +0 -0
  20. valcyon-0.1.0/src/valcyon/eval/metrics.py +40 -0
  21. valcyon-0.1.0/src/valcyon/eval/runner.py +0 -0
  22. valcyon-0.1.0/src/valcyon/ids.py +20 -0
  23. valcyon-0.1.0/src/valcyon/models.py +49 -0
  24. valcyon-0.1.0/src/valcyon/pipeline.py +46 -0
  25. valcyon-0.1.0/src/valcyon/rerankers/__init__.py +18 -0
  26. valcyon-0.1.0/src/valcyon/rerankers/bge.py +23 -0
  27. valcyon-0.1.0/src/valcyon/retrieve/__init__.py +0 -0
  28. valcyon-0.1.0/src/valcyon/retrieve/dense.py +53 -0
  29. valcyon-0.1.0/src/valcyon/retrieve/fusion.py +0 -0
  30. valcyon-0.1.0/src/valcyon/retrieve/hybrid.py +29 -0
  31. valcyon-0.1.0/src/valcyon/retrieve/keyword.py +41 -0
  32. valcyon-0.1.0/src/valcyon/retrieve/rerank.py +23 -0
  33. valcyon-0.1.0/src/valcyon/retrieve/search.py +0 -0
  34. valcyon-0.1.0/src/valcyon/retrieve/tree_ops.py +0 -0
  35. valcyon-0.1.0/src/valcyon/store/__init__.py +3 -0
  36. valcyon-0.1.0/src/valcyon/store/keyword.py +41 -0
  37. valcyon-0.1.0/src/valcyon/store/sqlite.py +181 -0
  38. valcyon-0.1.0/src/valcyon/store/vectors.py +0 -0
  39. valcyon-0.1.0/src/valcyon/tree/__init__.py +0 -0
  40. valcyon-0.1.0/src/valcyon/tree/builder.py +52 -0
  41. valcyon-0.1.0/src/valcyon/tree/levels.py +60 -0
  42. valcyon-0.1.0/src/valcyon/tree/promote.py +61 -0
  43. valcyon-0.1.0/src/valcyon/tree/restructure.py +89 -0
  44. valcyon-0.1.0/src/valcyon/tree/validate.py +97 -0
@@ -0,0 +1,10 @@
1
+ .venv/
2
+ __pycache__/
3
+ *.egg-info/
4
+ dist/
5
+ *.db
6
+ .pytest_cache/
7
+ .env
8
+ out/
9
+ out_vlm/
10
+ .venv312/
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.
@@ -0,0 +1,7 @@
1
+ # Valcyon
2
+
3
+ <p align="center">
4
+ <img src="assets/kingfisher-logo.png" alt="pdftree logo" width="220">
5
+ </p>
6
+
7
+ Deterministic hierarchical PDF extraction and retrieval toolkit. This project is in early development.
@@ -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,7 @@
1
+ from typing import Protocol
2
+
3
+ from valcyon.models import Element
4
+
5
+
6
+ class Parser(Protocol):
7
+ def parse(self, source: str) -> list[Element]: ...
@@ -0,0 +1,2 @@
1
+ class MinerUError(RuntimeError):
2
+ """No token. Put MINERU_TOKEN in a .env file or set it as an environment variable."""
@@ -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}")