Spaces:
Runtime error
Runtime error
File size: 4,606 Bytes
da4758c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import numpy as np
from src import config
from src.chunking import split_text
from src.embeddings import build_cosine_index, embed_texts
class EvidenceError(Exception):
"""Raised when there is no readable text to build evidence from."""
class EvidenceStore:
"""All source text, chunked and indexed in an in-memory FAISS cosine index.
Used three ways: retrieve passages per sub-question for the Writer (`gather`),
and pull the passages that back a specific claim for the Checker (`support_for`).
"""
def __init__(self, sources: list, embedder=None):
self.sources = sources
self.embedder = embedder
self.records = []
for source_idx, source in enumerate(sources):
for piece in split_text(source["content"], config.CHUNK_SIZE, config.CHUNK_OVERLAP):
self.records.append({"id": len(self.records), "text": piece, "source_idx": source_idx})
if not self.records:
raise EvidenceError("The sources contained no readable text.")
self.vectors = embed_texts([r["text"] for r in self.records], embedder)
self.index = build_cosine_index(self.vectors)
self.numbering = {} # source_idx -> citation number, assigned by gather()
self.selected = []
def _search(self, query: str, k: int):
query_vector = embed_texts([query], self.embedder)
scores, ids = self.index.search(query_vector, min(len(self.records), k))
return [(self.records[i], float(s)) for s, i in zip(scores[0], ids[0]) if i >= 0]
def gather(self, questions: list, chunks_per_question: int) -> tuple:
"""Selects the best passages for each sub-question; returns (selected_chunks, numbered_sources)."""
chosen, chosen_ids, per_source_total = [], set(), {}
for q_idx, question in enumerate(questions):
kept, per_source_q = 0, {}
for record, score in self._search(question, chunks_per_question * 6):
if kept >= chunks_per_question or len(chosen) >= config.MAX_EVIDENCE_CHUNKS:
break
sid = record["source_idx"]
if per_source_q.get(sid, 0) >= config.MAX_PER_SOURCE_PER_QUESTION:
continue
if record["id"] in chosen_ids:
# Already picked for an earlier question: it still covers this one, so count it
# instead of padding this question with weaker passages.
if score >= config.MIN_SCORE:
per_source_q[sid] = per_source_q.get(sid, 0) + 1
kept += 1
continue
if per_source_total.get(sid, 0) >= config.MAX_CHUNKS_PER_SOURCE:
continue
if score < config.MIN_SCORE and kept >= config.MIN_KEEP_PER_QUESTION:
continue
chosen_ids.add(record["id"])
chosen.append({**record, "score": score, "question_idx": q_idx})
per_source_q[sid] = per_source_q.get(sid, 0) + 1
per_source_total[sid] = per_source_total.get(sid, 0) + 1
kept += 1
if not chosen:
raise EvidenceError("No relevant passages were found in the sources.")
# Number sources by first appearance so citations read [1], [2], [3]... in order
self.numbering = {}
for chunk in chosen:
self.numbering.setdefault(chunk["source_idx"], len(self.numbering) + 1)
for chunk in chosen:
chunk["cite"] = self.numbering[chunk["source_idx"]]
self.selected = chosen
numbered = [
{
"n": n,
"title": self.sources[sid]["title"],
"url": self.sources[sid].get("url", ""),
"kind": self.sources[sid].get("kind", "web"),
}
for sid, n in sorted(self.numbering.items(), key=lambda kv: kv[1])
]
return chosen, numbered
def support_for(self, claim: str, cites: list, k: int = 2) -> list:
"""The k passages from the cited sources that are closest to a claim (what the Checker reads)."""
wanted = {sid for sid, n in self.numbering.items() if n in set(cites)}
candidates = [r["id"] for r in self.records if r["source_idx"] in wanted]
if not candidates:
return []
claim_vector = embed_texts([claim], self.embedder)[0]
sims = self.vectors[candidates] @ claim_vector
best = np.argsort(-sims)[:k]
return [self.records[candidates[i]]["text"] for i in best]
|