diff --git a/aborist/distill/__init__.py b/aborist/distill/__init__.py index 4b034dc..2756c55 100644 --- a/aborist/distill/__init__.py +++ b/aborist/distill/__init__.py @@ -2,11 +2,19 @@ from aborist.distill.base import DistillationResult, Distiller from aborist.distill.first_sentence import FirstSentenceDistiller +from aborist.distill.tfidf import TfidfKeywordDistiller -__all__ = ["DistillationResult", "Distiller", "FirstSentenceDistiller"] +__all__ = [ + "DistillationResult", + "Distiller", + "FirstSentenceDistiller", + "TfidfKeywordDistiller", +] def get_distiller(name: str) -> Distiller: if name == FirstSentenceDistiller.name: return FirstSentenceDistiller() + if name == TfidfKeywordDistiller.name: + return TfidfKeywordDistiller() raise ValueError(f"unknown distiller: {name}") diff --git a/aborist/distill/tfidf.py b/aborist/distill/tfidf.py new file mode 100644 index 0000000..5d2d115 --- /dev/null +++ b/aborist/distill/tfidf.py @@ -0,0 +1,96 @@ +"""TF-IDF top-K keyword distiller. + +Pure-Python, stdlib-only. Uses chunk-level "documents" as the corpus +baseline so a term that appears in every chunk of the same source gets +penalized relative to a term concentrated in fewer chunks. + +Output core content is a deterministic comma-separated keyword list — +extreme compression toward the "tweet/haiku" end of the planet metaphor. +""" + +from __future__ import annotations + +import math +import re +from collections import Counter + +from aborist.distill.base import DistillationResult, Distiller +from aborist.document import Document + + +_TOKEN_RE = re.compile(r"\b[a-zA-Z][a-zA-Z\-']{2,}\b") +_STOPWORDS = frozenset( + """ + the and or in of to is are was were for with on as by an be this that + these those from at but not have has had been they their his her its + our we you he she it all any if no more than such also can may will + would should could do does did between which where when what who how + some many most into through during above below after before since + while well much even only still about other only own same so very + just over under + """.split() +) + + +def _tokenize(text: str) -> list[str]: + return [t.lower() for t in _TOKEN_RE.findall(text) if t.lower() not in _STOPWORDS] + + +class TfidfKeywordDistiller(Distiller): + name = "tfidf-keywords-v1" + + def __init__(self, top_k: int = 16): + self.top_k = top_k + + def distill( + self, source: Document, source_chunks: list[str] + ) -> DistillationResult: + if not source_chunks: + return self._empty_result(source) + + chunk_tokens = [_tokenize(c) for c in source_chunks] + n = len(chunk_tokens) + + df: Counter[str] = Counter() + for toks in chunk_tokens: + df.update(set(toks)) + + tf: Counter[str] = Counter() + for toks in chunk_tokens: + tf.update(toks) + + scores: dict[str, float] = {} + for t, f in tf.items(): + idf = math.log((n + 1) / (df[t] + 1)) + 1.0 + scores[t] = f * idf + + ranked = sorted(scores.items(), key=lambda kv: (-kv[1], kv[0])) + keywords = [t for t, _ in ranked[: self.top_k]] + kw_set = set(keywords) + + contributing: set[int] = set() + for i, toks in enumerate(chunk_tokens): + if any(t in kw_set for t in toks): + contributing.add(i) + + core_text = ", ".join(keywords) + return DistillationResult( + core=Document( + uri=f"{source.uri}#core/{self.name}", + content=core_text, + source_type=f"core:{self.name}", + title=(source.title or "") + " [KEYWORDS]", + ), + contributing_chunk_indices=sorted(contributing), + ) + + def _empty_result(self, source: Document) -> DistillationResult: + return DistillationResult( + core=Document( + uri=f"{source.uri}#core/{self.name}", + content="", + source_type=f"core:{self.name}", + title=(source.title or "") + " [KEYWORDS]", + ), + contributing_chunk_indices=[], + ) diff --git a/tests/test_tfidf.py b/tests/test_tfidf.py new file mode 100644 index 0000000..fdefbf9 --- /dev/null +++ b/tests/test_tfidf.py @@ -0,0 +1,97 @@ +"""TF-IDF keyword distiller tests.""" + +from __future__ import annotations + +import json +from typing import Iterator + +from aborist.distill import TfidfKeywordDistiller +from aborist.distill.runner import distill_existing +from aborist.document import Document +from aborist.ingest import ingest_source +from aborist.merkle import proof_from_dict, verify_proof +from aborist.source import Source +from aborist.store import connect + + +class FakeSource(Source): + source_type = "test" + + def __init__(self, docs: list[Document]): + self.docs = docs + + def iter_documents(self) -> Iterator[Document]: + yield from self.docs + + +def test_tfidf_picks_distinctive_terms(): + """A term concentrated in one chunk should outrank one spread across all.""" + distiller = TfidfKeywordDistiller(top_k=4) + chunks = [ + "merkle merkle merkle providence cryptography", + "the quick brown fox common common common common", + "wikipedia article history history history", + ] + src = Document(uri="x", content="", source_type="t", title="t") + result = distiller.distill(src, chunks) + keywords = result.core.content.split(", ") + # 'merkle' (3x in chunk 0, 0 elsewhere) should rank high. + assert "merkle" in keywords + # 'history' should also rank high (3x in chunk 2). + assert "history" in keywords + + +def test_tfidf_deterministic(): + """Same input -> same output.""" + d = TfidfKeywordDistiller(top_k=8) + src = Document(uri="x", content="", source_type="t", title="t") + chunks = [ + "merkle providence proves provenance cryptographically", + "providence cache stores answer records", + ] + a = d.distill(src, chunks) + b = d.distill(src, chunks) + assert a.core.content == b.core.content + assert a.contributing_chunk_indices == b.contributing_chunk_indices + + +def test_tfidf_distillation_full_round_trip(tmp_path): + db = tmp_path / "tfidf.db" + conn = connect(db) + try: + long_text = ( + "merkle providence reverse rag verifies provenance cryptographically. " * 30 + + "\n\n" + + "wikipedia article from 2003 about anarchism. " * 30 + + "\n\n" + + "the eight forms of capital include living, social, intellectual. " * 30 + ) + ingest_source(conn, FakeSource([Document(uri="t://x", content=long_text, source_type="test", title="X")])) + result = distill_existing(conn, TfidfKeywordDistiller(top_k=8)) + assert result["distilled"] == 1 + + core = conn.execute( + "SELECT * FROM documents WHERE kind='core'" + ).fetchone() + assert core["source_type"] == "core:tfidf-keywords-v1" + assert "[KEYWORDS]" in (core["title"] or "") + + # Core content is comma-separated tokens. + chunk_row = conn.execute( + "SELECT content FROM chunks WHERE document_root=? ORDER BY idx", + (core["document_root"],), + ).fetchone() + keywords = chunk_row["content"].split(", ") + assert len(keywords) <= 8 + # Distinctive terms should appear. + assert any(k in {"merkle", "providence", "anarchism", "wikipedia", "capital"} for k in keywords) + + # Per-chunk Merkle proofs in derivation must reconstruct source root. + deriv = conn.execute("SELECT proof_blob, src_root FROM derivations").fetchone() + blob = json.loads(deriv["proof_blob"]) + for entry in blob["contributing"]: + p = proof_from_dict(entry["proof"]) + assert verify_proof(p) + assert p.root.hex() == deriv["src_root"] + finally: + conn.close()