spectrafuse 0.2.0__py3-none-any.whl
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.
- spectrafuse/__init__.py +27 -0
- spectrafuse/_version.py +7 -0
- spectrafuse/cli.py +115 -0
- spectrafuse/config.py +96 -0
- spectrafuse/errors.py +52 -0
- spectrafuse/ingestion.py +118 -0
- spectrafuse/log.py +30 -0
- spectrafuse/retrieval.py +158 -0
- spectrafuse/utils/__init__.py +1 -0
- spectrafuse/utils/bm25.py +53 -0
- spectrafuse/utils/chunker.py +112 -0
- spectrafuse/utils/clusters.py +63 -0
- spectrafuse/utils/diffusion.py +105 -0
- spectrafuse/utils/edge_list.py +18 -0
- spectrafuse/utils/edges.py +38 -0
- spectrafuse/utils/edges_content.py +110 -0
- spectrafuse/utils/edges_structure.py +86 -0
- spectrafuse/utils/embedder.py +89 -0
- spectrafuse/utils/embedder_id.py +78 -0
- spectrafuse/utils/entities.py +72 -0
- spectrafuse/utils/evaluation.py +112 -0
- spectrafuse/utils/expansion.py +68 -0
- spectrafuse/utils/graph.py +90 -0
- spectrafuse/utils/graph_build.py +27 -0
- spectrafuse/utils/graph_template.py +105 -0
- spectrafuse/utils/graph_view.py +126 -0
- spectrafuse/utils/index_folder.py +78 -0
- spectrafuse/utils/llm.py +62 -0
- spectrafuse/utils/loader.py +61 -0
- spectrafuse/utils/metrics.py +61 -0
- spectrafuse/utils/prompt.py +62 -0
- spectrafuse/utils/questions.py +40 -0
- spectrafuse/utils/references.py +43 -0
- spectrafuse/utils/sections.py +98 -0
- spectrafuse/utils/seeding.py +106 -0
- spectrafuse/utils/selection.py +93 -0
- spectrafuse/utils/selectors.py +104 -0
- spectrafuse/utils/semantic_split.py +54 -0
- spectrafuse/utils/spectral.py +72 -0
- spectrafuse/utils/store.py +94 -0
- spectrafuse/utils/tokens.py +64 -0
- spectrafuse/utils/tuning.py +78 -0
- spectrafuse/utils/units.py +106 -0
- spectrafuse/utils/vector_index.py +124 -0
- spectrafuse/validation.py +84 -0
- spectrafuse-0.2.0.dist-info/METADATA +606 -0
- spectrafuse-0.2.0.dist-info/RECORD +52 -0
- spectrafuse-0.2.0.dist-info/WHEEL +5 -0
- spectrafuse-0.2.0.dist-info/entry_points.txt +2 -0
- spectrafuse-0.2.0.dist-info/licenses/LICENSE +21 -0
- spectrafuse-0.2.0.dist-info/licenses/NOTICE +23 -0
- spectrafuse-0.2.0.dist-info/top_level.txt +1 -0
spectrafuse/__init__.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""SpectraFuse: an independent implementation inspired by SpectRAG (Sahinoglu, Mohammed,
|
|
2
|
+
Kazangirler and Kazangirler, arXiv:2610.36991), spectral graph-augmented retrieval for RAG.
|
|
3
|
+
Not the official SpectRAG and not affiliated with its authors; the official implementation is
|
|
4
|
+
https://github.com/SightAndCite/SpectRAG. The two classes keep the paper's method name.
|
|
5
|
+
|
|
6
|
+
SpectRAGIngestion builds the index (offline, once per document set).
|
|
7
|
+
SpectRAGRetrieval finds the context for a question and can answer it (online, per question).
|
|
8
|
+
|
|
9
|
+
from spectrafuse import SpectRAGIngestion, SpectRAGRetrieval
|
|
10
|
+
SpectRAGIngestion("my_index").ingest("docs/")
|
|
11
|
+
print(SpectRAGRetrieval("my_index").search("my question")["context"])
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from spectrafuse._version import INDEX_FORMAT, __version__
|
|
15
|
+
from spectrafuse.config import DEFAULTS, make_settings
|
|
16
|
+
from spectrafuse.errors import (DocumentError, DocumentNotFoundError, EmbedderError, IndexError_,
|
|
17
|
+
LLMError, MissingPackageError, NoDocumentsError, SettingsError,
|
|
18
|
+
SpectRAGError, UnknownSettingError)
|
|
19
|
+
from spectrafuse.ingestion import SpectRAGIngestion
|
|
20
|
+
from spectrafuse.log import enable_logging
|
|
21
|
+
from spectrafuse.retrieval import SpectRAGRetrieval
|
|
22
|
+
|
|
23
|
+
__all__ = ["SpectRAGIngestion", "SpectRAGRetrieval", "make_settings", "DEFAULTS",
|
|
24
|
+
"enable_logging", "SpectRAGError", "SettingsError", "UnknownSettingError",
|
|
25
|
+
"IndexError_", "EmbedderError", "DocumentError", "NoDocumentsError",
|
|
26
|
+
"DocumentNotFoundError", "MissingPackageError", "LLMError",
|
|
27
|
+
"__version__", "INDEX_FORMAT"]
|
spectrafuse/_version.py
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
"""The package version (read by pyproject.toml) and the index format version."""
|
|
2
|
+
|
|
3
|
+
__version__ = "0.2.0"
|
|
4
|
+
|
|
5
|
+
# Bump INDEX_FORMAT when the files of an index change in a way older code cannot read.
|
|
6
|
+
# 1 = JSON edge lists (before the Phase 1 audit), 2 = binary .npy/.npz edges and signals.
|
|
7
|
+
INDEX_FORMAT = 2
|
spectrafuse/cli.py
ADDED
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
"""The `spectrafuse` command.
|
|
2
|
+
|
|
3
|
+
spectrafuse ingest docs/ --index my_index [--set chunk_tokens=256 --set embedder=ollama]
|
|
4
|
+
spectrafuse search "my question" --index my_index [--budget 4000] [--json]
|
|
5
|
+
spectrafuse answer "my question" --index my_index --model gemma3:4b
|
|
6
|
+
spectrafuse info --index my_index
|
|
7
|
+
spectrafuse graph --index my_index --out graph.html [--question "my question"]
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
import argparse
|
|
11
|
+
import json
|
|
12
|
+
import sys
|
|
13
|
+
|
|
14
|
+
from spectrafuse._version import __version__
|
|
15
|
+
from spectrafuse.errors import SpectRAGError
|
|
16
|
+
from spectrafuse.log import enable_logging
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def parse_sets(pairs):
|
|
20
|
+
"""["alpha=0.5", "selector=mmr"] -> {"alpha": 0.5, "selector": "mmr"} (values read as JSON)."""
|
|
21
|
+
settings = {}
|
|
22
|
+
for pair in pairs or []:
|
|
23
|
+
if "=" not in pair:
|
|
24
|
+
raise SystemExit("--set needs key=value, got %r" % pair)
|
|
25
|
+
key, value = pair.split("=", 1)
|
|
26
|
+
try:
|
|
27
|
+
settings[key] = json.loads(value)
|
|
28
|
+
except ValueError:
|
|
29
|
+
settings[key] = value # plain text, e.g. selector=mmr
|
|
30
|
+
return settings
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def command_ingest(args):
|
|
34
|
+
from spectrafuse import SpectRAGIngestion
|
|
35
|
+
info = SpectRAGIngestion(args.index, parse_sets(args.set)).ingest(args.path)
|
|
36
|
+
print(json.dumps(info, indent=1))
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def command_search(args):
|
|
40
|
+
from spectrafuse import SpectRAGRetrieval
|
|
41
|
+
result = SpectRAGRetrieval(args.index, parse_sets(args.set)).search(args.question, args.budget)
|
|
42
|
+
if args.json:
|
|
43
|
+
print(json.dumps({key: result[key] for key in ("question", "tokens", "chunks", "context")}, indent=1))
|
|
44
|
+
return
|
|
45
|
+
for chunk in result["chunks"]:
|
|
46
|
+
print("%2d. %-30s score %.3f %d tokens (%s)" % (chunk["rank"], chunk["doc_id"], chunk["score"],
|
|
47
|
+
chunk["tokens"], chunk["reason"]))
|
|
48
|
+
print("\n" + result["context"])
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def command_answer(args):
|
|
52
|
+
from spectrafuse import SpectRAGRetrieval
|
|
53
|
+
from spectrafuse.utils.llm import ollama_llm
|
|
54
|
+
search = SpectRAGRetrieval(args.index, parse_sets(args.set))
|
|
55
|
+
result = search.answer(args.question, ollama_llm(args.model), args.budget)
|
|
56
|
+
print(result["answer"])
|
|
57
|
+
print("\nSources: " + ", ".join(sorted({c["doc_id"] for c in result["chunks"]})))
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def command_info(args):
|
|
61
|
+
from spectrafuse.utils.store import check_index_folder, load_index
|
|
62
|
+
from pathlib import Path
|
|
63
|
+
check_index_folder(Path(args.index))
|
|
64
|
+
print(json.dumps(load_index(args.index, only=("info",))["info"], indent=1))
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def command_graph(args):
|
|
68
|
+
from spectrafuse.utils.graph_view import save_graph_view
|
|
69
|
+
path = save_graph_view(args.index, args.out, args.question, parse_sets(args.set), args.max_nodes)
|
|
70
|
+
print("Graph view written to %s (open it in a browser)" % path)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def build_parser():
|
|
74
|
+
parser = argparse.ArgumentParser(prog="spectrafuse", description="SpectraFuse: spectral graph RAG, an independent implementation inspired by SpectRAG (arXiv:2610.36991). Official SpectRAG: github.com/SightAndCite/SpectRAG")
|
|
75
|
+
parser.add_argument("--version", action="version", version="spectrafuse " + __version__)
|
|
76
|
+
parser.add_argument("-v", "--verbose", action="store_true", help="print progress messages")
|
|
77
|
+
commands = parser.add_subparsers(dest="command", required=True)
|
|
78
|
+
ingest = commands.add_parser("ingest", help="build an index from a file or folder")
|
|
79
|
+
ingest.add_argument("path")
|
|
80
|
+
for name, helptext in [("search", "find the context for a question"), ("answer", "answer with an LLM")]:
|
|
81
|
+
sub = commands.add_parser(name, help=helptext)
|
|
82
|
+
sub.add_argument("question")
|
|
83
|
+
sub.add_argument("--budget", type=int, default=None, help="token budget (default 12000)")
|
|
84
|
+
if name == "search":
|
|
85
|
+
sub.add_argument("--json", action="store_true", help="print JSON")
|
|
86
|
+
else:
|
|
87
|
+
sub.add_argument("--model", default="gemma3:4b", help="Ollama model")
|
|
88
|
+
info = commands.add_parser("info", help="show what an index contains")
|
|
89
|
+
graph = commands.add_parser("graph", help="write an interactive graph view (HTML)")
|
|
90
|
+
graph.add_argument("--out", default="graph.html", help="HTML file to write")
|
|
91
|
+
graph.add_argument("--question", default=None, help="show this question's retrieval path")
|
|
92
|
+
graph.add_argument("--max-nodes", type=int, default=1500, dest="max_nodes")
|
|
93
|
+
for sub in (ingest, commands.choices["search"], commands.choices["answer"], info, graph):
|
|
94
|
+
sub.add_argument("--index", required=True, help="index folder")
|
|
95
|
+
for sub in (ingest, commands.choices["search"], commands.choices["answer"], graph):
|
|
96
|
+
sub.add_argument("--set", action="append", metavar="KEY=VALUE", help="change a setting")
|
|
97
|
+
return parser
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def main(argv=None):
|
|
101
|
+
args = build_parser().parse_args(argv)
|
|
102
|
+
if args.verbose:
|
|
103
|
+
enable_logging()
|
|
104
|
+
handlers = {"ingest": command_ingest, "search": command_search, "answer": command_answer,
|
|
105
|
+
"info": command_info, "graph": command_graph}
|
|
106
|
+
try:
|
|
107
|
+
handlers[args.command](args)
|
|
108
|
+
except SpectRAGError as error:
|
|
109
|
+
print("spectrafuse: error: %s" % error, file=sys.stderr)
|
|
110
|
+
return 1
|
|
111
|
+
return 0
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
if __name__ == "__main__":
|
|
115
|
+
sys.exit(main())
|
spectrafuse/config.py
ADDED
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
"""Default settings for SpectraFuse.
|
|
2
|
+
|
|
3
|
+
Values marked "paper" come from the SpectRAG paper (Tables 1-3).
|
|
4
|
+
Values marked "assumed" are not given in the paper, so we picked a sensible value.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import copy
|
|
8
|
+
import difflib
|
|
9
|
+
|
|
10
|
+
from spectrafuse.errors import UnknownSettingError
|
|
11
|
+
from spectrafuse.validation import validate_settings
|
|
12
|
+
|
|
13
|
+
# Edge-signal weights from the paper's "signal configuration 1" (Table 9).
|
|
14
|
+
EDGE_WEIGHTS_CONFIG_1 = {
|
|
15
|
+
"semantic": 0.300,
|
|
16
|
+
"section": 0.250,
|
|
17
|
+
"entity": 0.200,
|
|
18
|
+
"adjacency": 0.084,
|
|
19
|
+
"utility": 0.083,
|
|
20
|
+
"citation": 0.083,
|
|
21
|
+
}
|
|
22
|
+
|
|
23
|
+
# Edge-signal weights from the paper's "signal configuration 2" (Table 10).
|
|
24
|
+
EDGE_WEIGHTS_CONFIG_2 = {
|
|
25
|
+
"semantic": 0.150,
|
|
26
|
+
"section": 0.200,
|
|
27
|
+
"entity": 0.250,
|
|
28
|
+
"adjacency": 0.100,
|
|
29
|
+
"utility": 0.280,
|
|
30
|
+
"citation": 0.020,
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
DEFAULTS = {
|
|
34
|
+
# --- Chunking (paper) ---
|
|
35
|
+
"chunk_tokens": 1200,
|
|
36
|
+
"overlap_tokens": 100,
|
|
37
|
+
"min_chunk_tokens": 32,
|
|
38
|
+
"semantic_percentile": 90,
|
|
39
|
+
"semantic_embedder": None, # embedder for semantic-chunking breaks; None = "embedder" (paper)
|
|
40
|
+
# --- Embedding ---
|
|
41
|
+
"embedder": "hash", # "hash" (offline), "ollama" (nomic-embed-text), or your own function
|
|
42
|
+
"embed_dims": 768, # paper: nomic-embed-text is 768-d (used by the hash embedder)
|
|
43
|
+
"ollama_model": "nomic-embed-text",
|
|
44
|
+
"ollama_url": "http://localhost:11434",
|
|
45
|
+
# --- Vector search (Section 6.2) ---
|
|
46
|
+
"ann_backend": "auto", # "exact", "faiss", "hnsw", or "auto" (exact up to 20k chunks)
|
|
47
|
+
"hnsw_m": 32, # HNSW links per node
|
|
48
|
+
"hnsw_ef_construction": 200,
|
|
49
|
+
"hnsw_ef_search": 128,
|
|
50
|
+
"dense_top_n": 1000, # with an approximate backend: dense scores for the top n only
|
|
51
|
+
# --- Graph (weights: paper, the rest: assumed) ---
|
|
52
|
+
"edge_weights": EDGE_WEIGHTS_CONFIG_1,
|
|
53
|
+
"semantic_k": 10, # assumed
|
|
54
|
+
"edge_threshold": 0.05, # assumed
|
|
55
|
+
"graph_k": 10, # assumed: each chunk keeps at most its 10 strongest edges
|
|
56
|
+
# --- Spectral (paper) ---
|
|
57
|
+
"spectral_dims": 32,
|
|
58
|
+
"n_clusters": None, # assumed: None means about sqrt(n / 2) clusters
|
|
59
|
+
# --- Retrieval (paper) ---
|
|
60
|
+
"seed_count": 20,
|
|
61
|
+
"hops": 2,
|
|
62
|
+
"max_candidates": 200,
|
|
63
|
+
"token_budget": 12000,
|
|
64
|
+
# --- Retrieval (assumed) ---
|
|
65
|
+
"rrf_k": 60, # the usual RRF constant
|
|
66
|
+
"bm25_k1": 1.5, # standard BM25 values (Robertson and Zaragoza)
|
|
67
|
+
"bm25_b": 0.75,
|
|
68
|
+
"bm25_weight": 0.5, # hybrid relevance = 0.5 * dense + 0.5 * BM25
|
|
69
|
+
"alpha": 0.3,
|
|
70
|
+
"ppr_beta": 0.15,
|
|
71
|
+
"ppr_steps": 30,
|
|
72
|
+
"diffusion_t": 1.0, # heat-kernel scale t for spectral closeness (Eq. 17)
|
|
73
|
+
"proximity_mode": "hybrid", # "hybrid" (paper), "ppr", "spectral", "heat", "none"
|
|
74
|
+
"selector": "spectrag", # "spectrag", "mmr", "protected_mmr", "anchor_expand_ce", "topk"
|
|
75
|
+
"mmr_lambda": 0.7, # MMR: weight of relevance vs. novelty
|
|
76
|
+
"core_size": 3,
|
|
77
|
+
"precision_floor": 0.4,
|
|
78
|
+
"select_weights": {"a": 1.0, "b": 0.3, "c": 0.2, "d": 0.1, "lambda": 0.5},
|
|
79
|
+
"length_power": 0.5,
|
|
80
|
+
}
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def make_settings(overrides=None):
|
|
84
|
+
"""Return a fresh copy of DEFAULTS with the user's overrides applied and checked.
|
|
85
|
+
|
|
86
|
+
Raises UnknownSettingError (a KeyError) for a misspelled name, and SettingsError
|
|
87
|
+
(a ValueError) for a value of the wrong type or out of range.
|
|
88
|
+
"""
|
|
89
|
+
settings = copy.deepcopy(DEFAULTS)
|
|
90
|
+
for key, value in (overrides or {}).items():
|
|
91
|
+
if key not in settings:
|
|
92
|
+
close = difflib.get_close_matches(key, DEFAULTS, n=1)
|
|
93
|
+
hint = " (did you mean %r?)" % close[0] if close else ""
|
|
94
|
+
raise UnknownSettingError("Unknown setting: %r%s" % (key, hint))
|
|
95
|
+
settings[key] = copy.deepcopy(value)
|
|
96
|
+
return validate_settings(settings)
|
spectrafuse/errors.py
ADDED
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
"""SpectraFuse's exceptions. Each one also inherits the matching built-in exception, so code that
|
|
2
|
+
catches ValueError, KeyError, FileNotFoundError or ConnectionError keeps working.
|
|
3
|
+
|
|
4
|
+
try:
|
|
5
|
+
search = SpectRAGRetrieval("my_index")
|
|
6
|
+
except spectrafuse.SpectRAGError as error:
|
|
7
|
+
print(error)
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class SpectRAGError(Exception):
|
|
12
|
+
"""Base class for every error raised on purpose by SpectraFuse."""
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class SettingsError(SpectRAGError, ValueError):
|
|
16
|
+
"""A setting has the wrong type, is out of range, or is not one of the allowed choices."""
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class UnknownSettingError(SettingsError, KeyError):
|
|
20
|
+
"""A setting name that SpectraFuse does not know (often a typo)."""
|
|
21
|
+
|
|
22
|
+
def __str__(self): # KeyError would otherwise print the message with quotes around it
|
|
23
|
+
return str(self.args[0]) if self.args else ""
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class IndexError_(SpectRAGError, FileNotFoundError):
|
|
27
|
+
"""The index folder is missing, incomplete, or was made by an incompatible version.
|
|
28
|
+
(The trailing underscore avoids hiding Python's built-in IndexError.)"""
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class EmbedderError(SpectRAGError, ConnectionError):
|
|
32
|
+
"""The embedder cannot be reached, or its vectors do not match the index."""
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class DocumentError(SpectRAGError, ValueError):
|
|
36
|
+
"""Documents given to ingest_documents() are not in the expected form."""
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class NoDocumentsError(DocumentError):
|
|
40
|
+
"""There was no document with text to ingest."""
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class DocumentNotFoundError(DocumentError, FileNotFoundError):
|
|
44
|
+
"""The file or folder given to ingest() does not exist."""
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class MissingPackageError(SpectRAGError, ImportError):
|
|
48
|
+
"""An optional package is needed (e.g. faiss-cpu or pypdf); the message says what to install."""
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
class LLMError(SpectRAGError, ConnectionError):
|
|
52
|
+
"""The LLM (e.g. a local Ollama model) cannot be reached or did not reply in time."""
|
spectrafuse/ingestion.py
ADDED
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
"""SpectRAGIngestion: the offline pipeline (paper Table 1). Run it once per document set.
|
|
2
|
+
|
|
3
|
+
ingest = SpectRAGIngestion("my_index")
|
|
4
|
+
info = ingest.ingest("examples/docs")
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import time
|
|
8
|
+
|
|
9
|
+
import numpy as np
|
|
10
|
+
|
|
11
|
+
from spectrafuse._version import INDEX_FORMAT, __version__
|
|
12
|
+
from spectrafuse.config import make_settings
|
|
13
|
+
from spectrafuse.errors import DocumentError, NoDocumentsError
|
|
14
|
+
from spectrafuse.log import log_step
|
|
15
|
+
from spectrafuse.utils.chunker import chunk_documents
|
|
16
|
+
from spectrafuse.utils.edges import build_edge_signals, count_edges
|
|
17
|
+
from spectrafuse.utils.embedder import make_embedder
|
|
18
|
+
from spectrafuse.utils.embedder_id import embed_probe, saved_settings
|
|
19
|
+
from spectrafuse.utils.entities import add_entities
|
|
20
|
+
from spectrafuse.utils.graph_build import build_graph, build_spectral
|
|
21
|
+
from spectrafuse.utils.index_folder import prepare_folder, replace_folder
|
|
22
|
+
from spectrafuse.utils.loader import load_docs
|
|
23
|
+
from spectrafuse.utils.questions import add_utility_questions
|
|
24
|
+
from spectrafuse.utils.references import add_references
|
|
25
|
+
from spectrafuse.utils.store import edges_to_array, save_index
|
|
26
|
+
from spectrafuse.utils.vector_index import build_vector_index, save_vector_index
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class SpectRAGIngestion:
|
|
30
|
+
"""Load -> chunk -> embed -> edge signals -> weighted graph -> spectral -> save."""
|
|
31
|
+
|
|
32
|
+
def __init__(self, index_dir, settings=None, entity_fn=None, question_fn=None):
|
|
33
|
+
self.index_dir = index_dir
|
|
34
|
+
self.settings = make_settings(settings)
|
|
35
|
+
self.embed = make_embedder(self.settings)
|
|
36
|
+
# Semantic chunking can use a faster embedder for sentence distances (None = the same).
|
|
37
|
+
split_choice = self.settings["semantic_embedder"]
|
|
38
|
+
self.split_embed = self.embed if split_choice is None else \
|
|
39
|
+
make_embedder(dict(self.settings, embedder=split_choice))
|
|
40
|
+
self.entity_fn = entity_fn # optional: text -> list of entities
|
|
41
|
+
self.question_fn = question_fn # optional: text -> list of utility questions
|
|
42
|
+
|
|
43
|
+
def make_chunks(self, docs):
|
|
44
|
+
"""Split documents into chunks with entities, references and utility questions."""
|
|
45
|
+
chunks = chunk_documents(docs, self.settings, embed_fn=self.split_embed)
|
|
46
|
+
add_entities(chunks, self.entity_fn)
|
|
47
|
+
add_references(chunks)
|
|
48
|
+
add_utility_questions(chunks, self.question_fn)
|
|
49
|
+
return chunks
|
|
50
|
+
|
|
51
|
+
def make_graph(self, signals):
|
|
52
|
+
"""Fuse the signals (Eq. 10), drop weak edges, keep each chunk's strongest graph_k."""
|
|
53
|
+
return build_graph(signals, self.settings)
|
|
54
|
+
|
|
55
|
+
def make_spectral(self, edges, node_count):
|
|
56
|
+
"""Laplacian (Eq. 11) -> spectral coordinates (Eq. 12-13) -> clusters."""
|
|
57
|
+
return build_spectral(edges, node_count, self.settings)
|
|
58
|
+
|
|
59
|
+
def ingest(self, path):
|
|
60
|
+
"""Ingest a file or a folder (.txt, .md, .pdf), save the index, return a summary."""
|
|
61
|
+
docs = load_docs(path)
|
|
62
|
+
if not docs:
|
|
63
|
+
raise NoDocumentsError("No .txt, .md or .pdf file with text found in " + str(path))
|
|
64
|
+
return self.ingest_documents(docs)
|
|
65
|
+
|
|
66
|
+
def ingest_documents(self, docs):
|
|
67
|
+
"""Ingest documents already in memory: [{"doc_id": ..., "text": ...}, ...]."""
|
|
68
|
+
check_documents(docs)
|
|
69
|
+
start = begin = time.time()
|
|
70
|
+
chunks = self.make_chunks(docs)
|
|
71
|
+
start = log_step("chunking", start, documents=len(docs), chunks=len(chunks))
|
|
72
|
+
embeddings = self.embed([chunk["text"] for chunk in chunks], kind="document")
|
|
73
|
+
vector_index = build_vector_index(embeddings, self.settings) # exact, FAISS or HNSW
|
|
74
|
+
start = log_step("embedding", start, backend=vector_index["backend"])
|
|
75
|
+
signals, question_vectors, owners = build_edge_signals(
|
|
76
|
+
chunks, embeddings, self.embed, self.settings, vector_index)
|
|
77
|
+
start = log_step("edge signals", start, **count_edges(signals))
|
|
78
|
+
edges, weights = self.make_graph(signals)
|
|
79
|
+
coords, eigenvalues, clusters = self.make_spectral(edges, len(chunks))
|
|
80
|
+
start = log_step("graph and spectral", start, edges=len(edges), dims=coords.shape[1])
|
|
81
|
+
for chunk, cluster in zip(chunks, clusters):
|
|
82
|
+
chunk["cluster"] = int(cluster)
|
|
83
|
+
info = {"documents": len({c["doc_id"] for c in chunks}), "chunks": len(chunks),
|
|
84
|
+
"signal_edges": count_edges(signals), "graph_edges": len(edges),
|
|
85
|
+
"weights_used": {name: round(w, 4) for name, w in weights.items()},
|
|
86
|
+
"spectral_dims": coords.shape[1], "ann_backend": vector_index["backend"],
|
|
87
|
+
"clusters": len(set(clusters.tolist())),
|
|
88
|
+
"format_version": INDEX_FORMAT, "spectrafuse_version": __version__,
|
|
89
|
+
"seconds": round(time.time() - begin, 2)}
|
|
90
|
+
partial = prepare_folder(self.index_dir)
|
|
91
|
+
save_index(partial, {
|
|
92
|
+
"chunks": chunks, "embeddings": embeddings, "edges": edges_to_array(edges),
|
|
93
|
+
"signals": {name: edges_to_array(found) for name, found in signals.items()},
|
|
94
|
+
"spectral": coords, "eigenvalues": eigenvalues,
|
|
95
|
+
"question_vectors": question_vectors,
|
|
96
|
+
"question_owners": np.asarray(owners, dtype=np.int64),
|
|
97
|
+
"settings": saved_settings(self.settings), "info": info, # functions saved by name
|
|
98
|
+
"embedder_probe": embed_probe(self.embed), # fingerprint, checked when opened
|
|
99
|
+
})
|
|
100
|
+
save_vector_index(vector_index, partial)
|
|
101
|
+
replace_folder(partial, self.index_dir)
|
|
102
|
+
log_step("saved index", start, folder=self.index_dir)
|
|
103
|
+
return info
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def check_documents(docs):
|
|
107
|
+
"""Every document needs a text and a unique doc_id."""
|
|
108
|
+
if not docs:
|
|
109
|
+
raise NoDocumentsError("No documents with text to ingest")
|
|
110
|
+
seen = set()
|
|
111
|
+
for number, doc in enumerate(docs):
|
|
112
|
+
if not isinstance(doc, dict) or not isinstance(doc.get("text"), str) or "doc_id" not in doc:
|
|
113
|
+
raise DocumentError("Document %d must be a dict with 'doc_id' and 'text' (a string)" % number)
|
|
114
|
+
if doc["doc_id"] in seen:
|
|
115
|
+
raise DocumentError("Two documents have the same doc_id: %r" % (doc["doc_id"],))
|
|
116
|
+
seen.add(doc["doc_id"])
|
|
117
|
+
if not any(doc["text"].strip() for doc in docs):
|
|
118
|
+
raise NoDocumentsError("Every document is empty")
|
spectrafuse/log.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
"""Logging. SpectraFuse is silent by default (like any library); switch messages on with:
|
|
2
|
+
|
|
3
|
+
import spectrafuse
|
|
4
|
+
spectrafuse.enable_logging() # INFO: one line per pipeline step, with its time
|
|
5
|
+
spectrafuse.enable_logging("DEBUG") # DEBUG: also one line per query
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import logging
|
|
9
|
+
import time
|
|
10
|
+
|
|
11
|
+
logger = logging.getLogger("spectrafuse")
|
|
12
|
+
logger.addHandler(logging.NullHandler())
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def enable_logging(level="INFO"):
|
|
16
|
+
"""Print SpectraFuse's log messages to the terminal (stderr)."""
|
|
17
|
+
if not any(getattr(h, "spectrafuse", False) for h in logger.handlers):
|
|
18
|
+
handler = logging.StreamHandler()
|
|
19
|
+
handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s"))
|
|
20
|
+
handler.spectrafuse = True
|
|
21
|
+
logger.addHandler(handler)
|
|
22
|
+
logger.setLevel(level)
|
|
23
|
+
return logger
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def log_step(name, start, **details):
|
|
27
|
+
"""Log one finished step: 'chunking: 1.23 s (chunks=8)'. Returns the time now."""
|
|
28
|
+
extra = ", ".join("%s=%s" % item for item in details.items())
|
|
29
|
+
logger.info("%s: %.2f s%s", name, time.time() - start, " (%s)" % extra if extra else "")
|
|
30
|
+
return time.time()
|
spectrafuse/retrieval.py
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
"""SpectRAGRetrieval: the online pipeline (paper Table 2, Algorithm 1). Run it per question.
|
|
2
|
+
|
|
3
|
+
search = SpectRAGRetrieval("my_index")
|
|
4
|
+
result = search.search("Which company did the co-founder of Aurora Labs join?")
|
|
5
|
+
print(result["context"])
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import time
|
|
9
|
+
|
|
10
|
+
from spectrafuse._version import INDEX_FORMAT
|
|
11
|
+
from spectrafuse.config import make_settings
|
|
12
|
+
from spectrafuse.errors import EmbedderError, IndexError_, SettingsError
|
|
13
|
+
from spectrafuse.log import logger
|
|
14
|
+
from spectrafuse.utils.bm25 import build_bm25
|
|
15
|
+
from spectrafuse.utils.diffusion import score_candidates
|
|
16
|
+
from spectrafuse.utils.embedder import make_embedder
|
|
17
|
+
from spectrafuse.utils.embedder_id import (CUSTOM, check_embedder, is_saved_function,
|
|
18
|
+
read_saved_settings)
|
|
19
|
+
from spectrafuse.utils.graph import degrees, neighbors
|
|
20
|
+
from spectrafuse.utils.graph_build import build_graph, build_spectral
|
|
21
|
+
from spectrafuse.utils.index_folder import readable_folder
|
|
22
|
+
from spectrafuse.utils.prompt import build_prompt, format_context
|
|
23
|
+
from spectrafuse.utils.seeding import find_seeds
|
|
24
|
+
from spectrafuse.utils.selectors import select
|
|
25
|
+
from spectrafuse.utils.store import array_to_edges, check_index_folder, load_index
|
|
26
|
+
from spectrafuse.utils.vector_index import load_vector_index
|
|
27
|
+
|
|
28
|
+
# These change the graph itself, so they need rebuild_graph() instead of a per-call override.
|
|
29
|
+
GRAPH_SETTINGS = ("edge_weights", "edge_threshold", "graph_k", "spectral_dims", "n_clusters")
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class SpectRAGRetrieval:
|
|
33
|
+
"""Seeds -> graph expansion -> proximity re-scoring -> token-budget selection -> context."""
|
|
34
|
+
|
|
35
|
+
def __init__(self, index_dir, settings=None, embed_fn=None, rerank_fn=None):
|
|
36
|
+
"""Load an index made by SpectRAGIngestion.
|
|
37
|
+
|
|
38
|
+
settings: settings to change, e.g. {"token_budget": 4000, "selector": "mmr"}.
|
|
39
|
+
Graph settings (edge_threshold, graph_k, ...) rebuild the graph from the
|
|
40
|
+
saved signals, without embedding anything again.
|
|
41
|
+
embed_fn: pass the same function again if the index was built with your own embedder.
|
|
42
|
+
rerank_fn: optional (question, texts) -> scores, e.g. a cross-encoder, used by the
|
|
43
|
+
"anchor_expand_ce" selector.
|
|
44
|
+
"""
|
|
45
|
+
self.index_dir = readable_folder(index_dir) # '<index>.old' only after an interrupted ingest
|
|
46
|
+
check_index_folder(self.index_dir)
|
|
47
|
+
self.index = load_index(self.index_dir, skip=("signals",)) # raw signals: loaded on demand
|
|
48
|
+
made_with = self.index["info"].get("format_version", 1)
|
|
49
|
+
if made_with > INDEX_FORMAT:
|
|
50
|
+
raise IndexError_("This index has format %d; this SpectraFuse version reads up to %d. "
|
|
51
|
+
"Upgrade spectrafuse." % (made_with, INDEX_FORMAT))
|
|
52
|
+
self.index["edges"] = array_to_edges(self.index["edges"])
|
|
53
|
+
self.index["question_owners"] = [int(owner) for owner in self.index["question_owners"]]
|
|
54
|
+
built_with = make_settings(read_saved_settings(self.index["settings"]))
|
|
55
|
+
self.settings = make_settings(dict(built_with, **(settings or {})))
|
|
56
|
+
if embed_fn is not None:
|
|
57
|
+
self.settings["embedder"] = embed_fn
|
|
58
|
+
elif is_saved_function(self.settings["embedder"]):
|
|
59
|
+
raise SettingsError("This index was built with your own embed function (%s). "
|
|
60
|
+
"Pass it again: SpectRAGRetrieval(index_dir, embed_fn=...)"
|
|
61
|
+
% self.settings["embedder"][len(CUSTOM):])
|
|
62
|
+
self.embed = make_embedder(self.settings)
|
|
63
|
+
check_embedder(self.embed, self.index.get("embedder_probe"), built_with["embedder"])
|
|
64
|
+
self.rerank_fn = rerank_fn
|
|
65
|
+
self.bm25 = build_bm25([chunk["text"] for chunk in self.index["chunks"]])
|
|
66
|
+
self.vector_index = load_vector_index(self.index_dir, self.index["embeddings"], self.settings)
|
|
67
|
+
if any(self.settings[key] != built_with[key] for key in GRAPH_SETTINGS):
|
|
68
|
+
self.rebuild_graph()
|
|
69
|
+
self.refresh_graph()
|
|
70
|
+
|
|
71
|
+
def refresh_graph(self):
|
|
72
|
+
"""Neighbour table (CSR) and degrees for the current edges."""
|
|
73
|
+
count = len(self.index["chunks"])
|
|
74
|
+
self.table = neighbors(self.index["edges"], count)
|
|
75
|
+
self.index["degree"] = degrees(self.index["edges"], count)
|
|
76
|
+
|
|
77
|
+
def rebuild_graph(self, overrides=None):
|
|
78
|
+
"""Change graph settings without re-embedding: fuse the saved signals again (Eq. 10)
|
|
79
|
+
and recompute spectral coordinates and clusters (Eq. 11-13). In memory only."""
|
|
80
|
+
self.settings = make_settings(dict(self.settings, **(overrides or {})))
|
|
81
|
+
saved = load_index(self.index_dir, only=("signals",))["signals"]
|
|
82
|
+
signals = {name: array_to_edges(edges) for name, edges in saved.items()}
|
|
83
|
+
edges, _ = build_graph(signals, self.settings)
|
|
84
|
+
coords, eigenvalues, clusters = build_spectral(edges, len(self.index["chunks"]), self.settings)
|
|
85
|
+
self.index.update({"edges": edges, "spectral": coords, "eigenvalues": eigenvalues})
|
|
86
|
+
for chunk, cluster in zip(self.index["chunks"], clusters):
|
|
87
|
+
chunk["cluster"] = int(cluster)
|
|
88
|
+
self.refresh_graph()
|
|
89
|
+
|
|
90
|
+
def settings_for(self, overrides):
|
|
91
|
+
"""The settings for one call: the current settings plus per-call overrides."""
|
|
92
|
+
if not overrides:
|
|
93
|
+
return self.settings
|
|
94
|
+
changed = [key for key in GRAPH_SETTINGS
|
|
95
|
+
if key in overrides and overrides[key] != self.settings[key]]
|
|
96
|
+
if changed:
|
|
97
|
+
raise SettingsError("Use rebuild_graph() to change graph settings: " + ", ".join(changed))
|
|
98
|
+
return make_settings(dict(self.settings, **overrides))
|
|
99
|
+
|
|
100
|
+
def retrieve(self, question, settings=None, query_vector=None):
|
|
101
|
+
"""Stages 1-4. Returns (seed result, scored candidates, picks)."""
|
|
102
|
+
settings = self.settings_for(settings)
|
|
103
|
+
start = time.time()
|
|
104
|
+
if query_vector is None:
|
|
105
|
+
query_vector = self.embed([question], kind="query")[0]
|
|
106
|
+
if len(query_vector) != self.index["embeddings"].shape[1]:
|
|
107
|
+
raise EmbedderError("The query vector has %d dimensions but the index has %d: search "
|
|
108
|
+
"with the same embedder the index was built with."
|
|
109
|
+
% (len(query_vector), self.index["embeddings"].shape[1]))
|
|
110
|
+
seeds = find_seeds(question, query_vector, self.index, self.bm25, settings,
|
|
111
|
+
self.vector_index)
|
|
112
|
+
scored = score_candidates(seeds["seeds"], seeds["relevance"], self.index,
|
|
113
|
+
self.table, settings)
|
|
114
|
+
rerank = None
|
|
115
|
+
if settings["selector"] == "anchor_expand_ce" and self.rerank_fn is not None:
|
|
116
|
+
texts = [self.index["chunks"][c]["text"] for c in scored["candidates"]]
|
|
117
|
+
rerank = self.rerank_fn(question, texts)
|
|
118
|
+
picks = select(scored, self.index, settings, self.table, rerank)
|
|
119
|
+
logger.debug("query %r: %d seeds, %d candidates, %d picked in %.1f ms", question[:60],
|
|
120
|
+
len(seeds["seeds"]), len(scored["candidates"]), len(picks),
|
|
121
|
+
1000 * (time.time() - start))
|
|
122
|
+
return seeds, scored, picks
|
|
123
|
+
|
|
124
|
+
def describe(self, k, reason, rank, scored):
|
|
125
|
+
"""One selected chunk with all its scores (the 'ranked source metadata')."""
|
|
126
|
+
chunk = self.index["chunks"][scored["candidates"][k]]
|
|
127
|
+
info = {key: chunk[key] for key in ["id", "doc_id", "position", "section", "tokens", "text"]}
|
|
128
|
+
info.update({"rank": rank, "reason": reason, "hop": scored["hop"][k],
|
|
129
|
+
"score": float(scored["scores"][k]),
|
|
130
|
+
"relevance": float(scored["relevance"][k]),
|
|
131
|
+
"proximity": float(scored["proximity"][k])})
|
|
132
|
+
return info
|
|
133
|
+
|
|
134
|
+
def search(self, question, token_budget=None, settings=None, query_vector=None):
|
|
135
|
+
"""Find the context for a question. Returns a dict with 'context', 'chunks', ...
|
|
136
|
+
|
|
137
|
+
settings: per-call overrides, e.g. {"selector": "mmr", "alpha": 0.5}.
|
|
138
|
+
"""
|
|
139
|
+
overrides = dict(settings or {})
|
|
140
|
+
if token_budget is not None:
|
|
141
|
+
overrides["token_budget"] = token_budget
|
|
142
|
+
seeds, scored, picks = self.retrieve(question, overrides, query_vector)
|
|
143
|
+
chunks = [self.describe(k, reason, rank, scored)
|
|
144
|
+
for rank, (k, reason) in enumerate(picks, start=1)]
|
|
145
|
+
context = format_context(chunks)
|
|
146
|
+
return {"question": question, "context": context, "chunks": chunks,
|
|
147
|
+
"prompt": build_prompt(question, context),
|
|
148
|
+
"tokens": sum(chunk["tokens"] for chunk in chunks),
|
|
149
|
+
"seeds": seeds["seeds"], "candidates": scored["candidates"]}
|
|
150
|
+
|
|
151
|
+
def answer(self, question, llm_fn, token_budget=None, settings=None):
|
|
152
|
+
"""search() + generation (Algorithm 1 line 13). llm_fn is any function: prompt -> text.
|
|
153
|
+
|
|
154
|
+
Returns the search result with an extra "answer" key; result["chunks"] are the sources.
|
|
155
|
+
"""
|
|
156
|
+
result = self.search(question, token_budget, settings)
|
|
157
|
+
result["answer"] = llm_fn(result["prompt"]).strip()
|
|
158
|
+
return result
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Small helper functions used by SpectRAGIngestion and SpectRAGRetrieval."""
|