doclighter 0.1.0__py3-none-any.whl

doclighter/__init__.py

@@ -0,0 +1,21 @@
+"""Doclighter — a semantic Ctrl+F that paints your document with relevance.
+
+See https://github.com/pratyush272/doclighter for docs.
+"""
+from .core import Doclighter, SearchResult
+from .chunking import Chunk, make_chunks
+from .scoring import word_heatmap, aggregate_multi_query
+from .render import render_html, score_to_hex
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Doclighter",
+    "SearchResult",
+    "Chunk",
+    "make_chunks",
+    "word_heatmap",
+    "aggregate_multi_query",
+    "render_html",
+    "score_to_hex",
+]

doclighter/chunking.py

@@ -0,0 +1,58 @@
+"""Rolling window chunking.
+
+Small word-window chunks (default 12 words, 50% overlap) are the unit of
+semantic match. This is deliberately finer than typical RAG chunking
+(256-1024 tokens) because Doclighter is a visualization tool, not a
+context-window filler — fine chunks give fine spatial resolution.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass(frozen=True)
+class Chunk:
+    """A single rolling window over the document word list."""
+
+    text: str
+    start: int   # inclusive word index
+    end: int     # exclusive word index
+
+
+def make_chunks(
+    words: List[str],
+    chunk_size: int = 12,
+    overlap: float = 0.5,
+) -> List[Chunk]:
+    """Split a word list into rolling windows.
+
+    Parameters
+    ----------
+    words : list of str
+        Tokenized document (typically ``text.split()``).
+    chunk_size : int
+        Words per window. Default 12 — small enough that semantic units rarely
+        get cut, large enough that MiniLM produces a useful embedding.
+    overlap : float
+        Fraction of overlap between consecutive windows, in [0, 1).
+        Default 0.5 means 50%% overlap (step = chunk_size / 2).
+
+    Returns
+    -------
+    list of Chunk
+    """
+    if not 0 <= overlap < 1:
+        raise ValueError(f"overlap must be in [0, 1), got {overlap}")
+    if chunk_size < 1:
+        raise ValueError(f"chunk_size must be >= 1, got {chunk_size}")
+
+    step = max(1, int(chunk_size * (1 - overlap)))
+    chunks: List[Chunk] = []
+    for i in range(0, len(words), step):
+        window = words[i : i + chunk_size]
+        if window:
+            chunks.append(Chunk(text=" ".join(window), start=i, end=i + len(window)))
+        if i + chunk_size >= len(words):
+            break  # last window already covers tail
+    return chunks

doclighter/core.py

@@ -0,0 +1,297 @@
+"""The main Doclighter API.
+
+Typical usage::
+
+    from doclighter import Doclighter
+
+    doc = Doclighter.from_pdf("contract.pdf")
+    result = doc.search("termination clauses")
+
+    # In Jupyter:
+    from IPython.display import HTML, display
+    display(HTML(result.to_html()))
+
+    # Anywhere else:
+    print(result.top_chunks(k=5))
+    scores = result.word_scores  # numpy array, shape (n_words,)
+"""
+from __future__ import annotations
+
+import pickle
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+
+from . import extract
+from .chunking import Chunk, make_chunks
+from .embedding import Embedder, default_embedder, wrap_callable
+from .index import FlatIndex, QuantizedIndex
+from .render import render_html
+from .scoring import aggregate_multi_query, word_heatmap
+
+PathLike = Union[str, Path]
+
+
+@dataclass
+class SearchResult:
+    """The output of a single Doclighter search.
+
+    Attributes
+    ----------
+    query : str or list of str
+        The query that produced this result. List form for multi-query searches.
+    word_scores : np.ndarray, shape (n_words,)
+        Per-word relevance in [0, 1]. The thing you'd visualize.
+    chunk_scores : np.ndarray, shape (n_chunks,)
+        Raw cosine similarity per chunk (before proximity smoothing).
+    elapsed_ms : float
+        Search latency, including all postprocessing.
+    decay_sigma : float
+        The sigma used for this search (recorded for reproducibility).
+    """
+
+    query: Union[str, List[str]]
+    word_scores: np.ndarray
+    chunk_scores: np.ndarray
+    elapsed_ms: float
+    decay_sigma: float
+    _words: List[str] = field(repr=False)
+    _chunks: List[Chunk] = field(repr=False)
+
+    def top_chunks(self, k: int = 10) -> List[Tuple[str, float, Tuple[int, int]]]:
+        """Top-k chunks by raw semantic score.
+
+        Returns list of ``(chunk_text, score, (start_word, end_word))``.
+        """
+        top_idx = np.argsort(-self.chunk_scores)[:k]
+        return [
+            (self._chunks[i].text, float(self.chunk_scores[i]),
+             (self._chunks[i].start, self._chunks[i].end))
+            for i in top_idx
+        ]
+
+    def to_html(self, **render_kwargs) -> str:
+        """Render the heatmap as an HTML fragment (string).
+
+        Extra kwargs are forwarded to ``doclighter.render.render_html``
+        (e.g. ``max_height_px``, ``quantize_levels``).
+        """
+        query_str = self.query if isinstance(self.query, str) else " | ".join(self.query)
+        return render_html(
+            words=self._words,
+            word_scores=self.word_scores,
+            query=query_str,
+            decay_sigma=self.decay_sigma,
+            elapsed_ms=self.elapsed_ms,
+            **render_kwargs,
+        )
+
+    def _repr_html_(self) -> str:
+        """Jupyter automatically calls this for rich display."""
+        return self.to_html()
+
+
+class Doclighter:
+    """A searchable, visualizable view of a single document.
+
+    Build once (embedding the document is the slow step), then search many
+    times — each search is sub-100ms on typical documents.
+
+    Parameters are mostly set at build time; ``decay_sigma`` can be overridden
+    per-search, which lets you toggle between fine word-level highlights
+    (small sigma) and broad thematic regions (large sigma) without re-indexing.
+    """
+
+    def __init__(
+        self,
+        text: str,
+        *,
+        chunk_size: int = 12,
+        chunk_overlap: float = 0.5,
+        embedder: Optional[Union[Embedder, Callable]] = None,
+        embedding_model: str = "all-MiniLM-L6-v2",
+        decay_sigma: float = 20.0,
+        quantize: bool = False,
+        quantize_rerank_k: int = 200,
+    ):
+        if not text or not text.strip():
+            raise ValueError("text is empty — nothing to index")
+
+        self.text = text
+        self.words = text.split()
+        self.chunks = make_chunks(self.words, chunk_size=chunk_size, overlap=chunk_overlap)
+        if not self.chunks:
+            raise ValueError("No chunks produced — document too short?")
+
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.decay_sigma = decay_sigma  # default; per-search override allowed
+        self.embedding_model_name = embedding_model
+
+        # Embedder: user-supplied callable or default MiniLM
+        if embedder is None:
+            self._embedder = default_embedder(embedding_model)
+        elif callable(embedder):
+            self._embedder = wrap_callable(embedder)
+        else:
+            raise TypeError(f"embedder must be callable, got {type(embedder)}")
+
+        # Embed all chunks (the slow step)
+        chunk_texts = [c.text for c in self.chunks]
+        self.chunk_embeddings = self._embedder(chunk_texts)
+
+        # Build index
+        if quantize:
+            self.index = QuantizedIndex(self.chunk_embeddings, rerank_k=quantize_rerank_k)
+        else:
+            self.index = FlatIndex(self.chunk_embeddings)
+
+    # ---------- alternate constructors ----------
+
+    @classmethod
+    def from_text(cls, text: str, **kwargs) -> "Doclighter":
+        """Build from a raw string. Same as the default constructor."""
+        return cls(text, **kwargs)
+
+    @classmethod
+    def from_pdf(cls, path: PathLike, **kwargs) -> "Doclighter":
+        """Build from a PDF file on disk."""
+        return cls(extract.from_pdf_path(path), **kwargs)
+
+    @classmethod
+    def from_url(cls, url: str, timeout: int = 30, **kwargs) -> "Doclighter":
+        """Build from a PDF served over HTTP(S)."""
+        return cls(extract.from_url(url, timeout=timeout), **kwargs)
+
+    # ---------- search ----------
+
+    def search(
+        self,
+        query: Union[str, Sequence[str]],
+        *,
+        decay_sigma: Optional[float] = None,
+        multi_query_aggregate: str = "max",
+    ) -> SearchResult:
+        """Search the document and return a SearchResult.
+
+        Parameters
+        ----------
+        query : str or sequence of str
+            A single query, or multiple sub-queries (e.g.
+            ``["termination", "indemnification"]``). Multi-query mode runs
+            each sub-query and combines the heatmaps.
+        decay_sigma : float, optional
+            Override the proximity decay scale for this search. Smaller =
+            sharper word-level highlights; larger = broader regions.
+            Defaults to the value set at construction.
+        multi_query_aggregate : {"max", "mean", "sum"}
+            How to combine sub-query heatmaps. See ``scoring.aggregate_multi_query``.
+
+        Returns
+        -------
+        SearchResult
+        """
+        t0 = time.perf_counter()
+        sigma = decay_sigma if decay_sigma is not None else self.decay_sigma
+
+        is_multi = not isinstance(query, str)
+        queries = list(query) if is_multi else [query]
+        q_embs = self._embedder(queries)  # shape (n_queries, dim)
+
+        heatmaps = []
+        all_chunk_scores = []
+        for q_emb in q_embs:
+            chunk_scores = self.index.all_scores(q_emb)
+            all_chunk_scores.append(chunk_scores)
+            hm = word_heatmap(
+                chunk_scores=chunk_scores,
+                chunks=self.chunks,
+                n_words=len(self.words),
+                decay_sigma=sigma,
+            )
+            heatmaps.append(hm)
+
+        if is_multi:
+            final_heatmap = aggregate_multi_query(heatmaps, mode=multi_query_aggregate)
+            # For chunk_scores, use the max across queries (most informative single number)
+            final_chunk_scores = np.max(np.stack(all_chunk_scores), axis=0)
+        else:
+            final_heatmap = heatmaps[0]
+            final_chunk_scores = all_chunk_scores[0]
+
+        elapsed_ms = (time.perf_counter() - t0) * 1000
+        return SearchResult(
+            query=query if is_multi else queries[0],
+            word_scores=final_heatmap,
+            chunk_scores=final_chunk_scores,
+            elapsed_ms=elapsed_ms,
+            decay_sigma=sigma,
+            _words=self.words,
+            _chunks=self.chunks,
+        )
+
+    # ---------- persistence ----------
+
+    def save(self, path: PathLike) -> None:
+        """Save the indexed document to disk (skip re-embedding next time).
+
+        Note: this saves embeddings and chunk metadata, not the embedder
+        itself — when you load, you'll need to be able to embed query
+        strings, so the same model must be available.
+        """
+        state = {
+            "text": self.text,
+            "words": self.words,
+            "chunks": self.chunks,
+            "chunk_embeddings": self.chunk_embeddings,
+            "chunk_size": self.chunk_size,
+            "chunk_overlap": self.chunk_overlap,
+            "decay_sigma": self.decay_sigma,
+            "embedding_model_name": self.embedding_model_name,
+        }
+        with open(path, "wb") as f:
+            pickle.dump(state, f)
+
+    @classmethod
+    def load(
+        cls,
+        path: PathLike,
+        *,
+        embedder: Optional[Callable] = None,
+        quantize: bool = False,
+    ) -> "Doclighter":
+        """Load a saved index. Re-instantiates the embedder for query encoding."""
+        with open(path, "rb") as f:
+            state = pickle.load(f)
+
+        obj = cls.__new__(cls)
+        obj.text = state["text"]
+        obj.words = state["words"]
+        obj.chunks = state["chunks"]
+        obj.chunk_embeddings = state["chunk_embeddings"]
+        obj.chunk_size = state["chunk_size"]
+        obj.chunk_overlap = state["chunk_overlap"]
+        obj.decay_sigma = state["decay_sigma"]
+        obj.embedding_model_name = state["embedding_model_name"]
+
+        if embedder is None:
+            obj._embedder = default_embedder(obj.embedding_model_name)
+        else:
+            obj._embedder = wrap_callable(embedder)
+
+        if quantize:
+            obj.index = QuantizedIndex(obj.chunk_embeddings)
+        else:
+            obj.index = FlatIndex(obj.chunk_embeddings)
+        return obj
+
+    # ---------- diagnostics ----------
+
+    def __repr__(self) -> str:
+        return (
+            f"Doclighter(words={len(self.words)}, chunks={len(self.chunks)}, "
+            f"chunk_size={self.chunk_size}, model={self.embedding_model_name!r})"
+        )

doclighter/embedding.py

@@ -0,0 +1,57 @@
+"""Embedding backend.
+
+Default: sentence-transformers ``all-MiniLM-L6-v2`` (384-dim, fast, ~80MB).
+Users can pass any callable matching ``Embedder`` (e.g. their own model,
+an API call, or a different sentence-transformers checkpoint).
+"""
+from __future__ import annotations
+
+from typing import Callable, List, Protocol
+
+import numpy as np
+
+
+class Embedder(Protocol):
+    """A callable that turns texts into a (N, dim) float32 numpy array.
+
+    Vectors should be L2-normalized — Doclighter uses inner-product cosine
+    similarity and assumes unit norm.
+    """
+
+    def __call__(self, texts: List[str]) -> np.ndarray: ...
+
+
+def default_embedder(model_name: str = "all-MiniLM-L6-v2") -> Embedder:
+    """Build a sentence-transformers-backed embedder.
+
+    The model is loaded once and reused. Batches of 64 by default; override
+    by wrapping with your own callable if you need finer control.
+    """
+    from sentence_transformers import SentenceTransformer  # local import
+
+    model = SentenceTransformer(model_name)
+
+    def encode(texts: List[str]) -> np.ndarray:
+        return model.encode(
+            texts,
+            normalize_embeddings=True,
+            batch_size=64,
+            convert_to_numpy=True,
+            show_progress_bar=len(texts) > 100,
+        ).astype("float32")
+
+    return encode
+
+
+def wrap_callable(fn: Callable[[List[str]], np.ndarray]) -> Embedder:
+    """Wrap a user-supplied callable, ensuring float32 + L2 normalization."""
+
+    def encode(texts: List[str]) -> np.ndarray:
+        out = np.asarray(fn(texts), dtype="float32")
+        if out.ndim != 2:
+            raise ValueError(f"Embedder must return 2D array, got shape {out.shape}")
+        norms = np.linalg.norm(out, axis=1, keepdims=True)
+        norms[norms == 0] = 1.0
+        return out / norms
+
+    return encode

doclighter/extract.py

@@ -0,0 +1,38 @@
+"""Text loaders for Doclighter.
+
+Supports PDFs (local path or URL) and raw text. Kept deliberately small —
+users with exotic formats (docx, html) should convert upstream and pass text.
+"""
+from __future__ import annotations
+
+from io import BytesIO
+from pathlib import Path
+from typing import Union
+
+
+def from_pdf_bytes(data: bytes) -> str:
+    """Extract text from PDF bytes using pypdf."""
+    from pypdf import PdfReader  # local import keeps base install lean
+
+    reader = PdfReader(BytesIO(data))
+    pages = [p.extract_text() for p in reader.pages if p.extract_text()]
+    if not pages:
+        raise RuntimeError(
+            "No extractable text in PDF. It may be image-only — "
+            "run OCR (e.g. pytesseract) and pass the result to Doclighter.from_text()."
+        )
+    return "\n".join(pages)
+
+
+def from_pdf_path(path: Union[str, Path]) -> str:
+    """Read a PDF from disk and extract text."""
+    return from_pdf_bytes(Path(path).read_bytes())
+
+
+def from_url(url: str, timeout: int = 30) -> str:
+    """Download a PDF from a URL and extract text."""
+    import requests
+
+    r = requests.get(url, timeout=timeout)
+    r.raise_for_status()
+    return from_pdf_bytes(r.content)

doclighter/index.py

@@ -0,0 +1,72 @@
+"""Vector index.
+
+Default: flat exact inner-product (fast enough for documents up to ~50K
+chunks, which covers virtually all single-document use cases — a 500-page
+book at 12-word chunks is ~30K chunks).
+
+Opt-in: SQ8 scalar-quantized index for very large docs or memory-constrained
+environments. Trades a tiny bit of ranking precision for 4x memory savings.
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+import numpy as np
+
+
+class FlatIndex:
+    """Exact inner-product search. Trivial, fast, no dependencies beyond numpy."""
+
+    def __init__(self, embeddings: np.ndarray):
+        self.embeddings = embeddings.astype("float32")
+        self.ntotal = embeddings.shape[0]
+        self.dim = embeddings.shape[1]
+
+    def search(self, query: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        """Return (scores, indices), both shape (k,), sorted descending."""
+        scores = self.embeddings @ query.reshape(-1)
+        k = min(k, self.ntotal)
+        # argpartition is O(n); full sort only on the top-k slice
+        top_idx = np.argpartition(-scores, k - 1)[:k]
+        top_idx = top_idx[np.argsort(-scores[top_idx])]
+        return scores[top_idx], top_idx
+
+    def all_scores(self, query: np.ndarray) -> np.ndarray:
+        """Exact scores for every vector — used for full word-heatmap rendering."""
+        return self.embeddings @ query.reshape(-1)
+
+
+class QuantizedIndex:
+    """SQ8 + exact rerank. ~4x smaller in memory, near-identical ranking.
+
+    Only useful for very large documents; for the default visualization use
+    case, FlatIndex is simpler and the speed difference is invisible.
+    """
+
+    def __init__(self, embeddings: np.ndarray, rerank_k: int = 200):
+        import faiss  # local import — only required for this path
+
+        self.embeddings = embeddings.astype("float32")  # kept for exact rerank
+        self.ntotal = embeddings.shape[0]
+        self.dim = embeddings.shape[1]
+        self.rerank_k = rerank_k
+
+        self._sq = faiss.IndexScalarQuantizer(
+            self.dim, faiss.ScalarQuantizer.QT_8bit, faiss.METRIC_INNER_PRODUCT
+        )
+        self._sq.train(self.embeddings)
+        self._sq.add(self.embeddings)
+
+    def search(self, query: np.ndarray, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        q = query.reshape(1, -1).astype("float32")
+        cand_k = min(max(k, self.rerank_k), self.ntotal)
+        _, cand_idx = self._sq.search(q, cand_k)
+        cand_idx = cand_idx[0]
+        # Exact rerank
+        exact_scores = self.embeddings[cand_idx] @ q[0]
+        order = np.argsort(-exact_scores)[:k]
+        return exact_scores[order], cand_idx[order]
+
+    def all_scores(self, query: np.ndarray) -> np.ndarray:
+        """Exact scores everywhere — we keep float32 embeddings for this."""
+        return self.embeddings @ query.reshape(-1)

doclighter/render.py

@@ -0,0 +1,124 @@
+"""HTML rendering of word-level heatmaps.
+
+The renderer is deliberately decoupled from Jupyter — ``to_html()`` returns
+a string. Jupyter users wrap with ``display(HTML(...))``; Streamlit users
+pass to ``st.markdown(..., unsafe_allow_html=True)``; FastAPI users return
+it from a route.
+"""
+from __future__ import annotations
+
+import html
+from typing import List, Tuple
+
+import numpy as np
+
+# Smooth gradient: grey -> indigo -> cyan -> yellow -> orange -> red
+# Tuples are (score_threshold, (r, g, b)) — Tailwind-esque palette that
+# reads well on both light and dark backgrounds.
+_COLOR_STOPS: List[Tuple[float, Tuple[int, int, int]]] = [
+    (1.00, (220, 38, 38)),    # red-600 — core match
+    (0.80, (249, 115, 22)),   # orange-500
+    (0.60, (234, 179, 8)),    # yellow-500
+    (0.40, (6, 182, 212)),    # cyan-500
+    (0.20, (79, 70, 229)),    # indigo-600
+    (0.00, (156, 163, 175)),  # gray-400 — unrelated
+]
+
+
+def score_to_hex(score: float) -> str:
+    """Interpolate a score in [0, 1] to a hex color along the gradient."""
+    s = max(0.0, min(1.0, score))
+    for i in range(len(_COLOR_STOPS) - 1):
+        hi_s, hi_c = _COLOR_STOPS[i]
+        lo_s, lo_c = _COLOR_STOPS[i + 1]
+        if s >= lo_s:
+            t = (s - lo_s) / (hi_s - lo_s) if hi_s > lo_s else 0.0
+            r = int(lo_c[0] + t * (hi_c[0] - lo_c[0]))
+            g = int(lo_c[1] + t * (hi_c[1] - lo_c[1]))
+            b = int(lo_c[2] + t * (hi_c[2] - lo_c[2]))
+            return f"#{r:02x}{g:02x}{b:02x}"
+    return "#9ca3af"
+
+
+def _legend_html(n_words: int, decay_sigma: float, elapsed_ms: float) -> str:
+    swatches = " → ".join(
+        f'<span style="color:{score_to_hex(s)}">■</span>'
+        for s in (1.0, 0.8, 0.6, 0.4, 0.2, 0.0)
+    )
+    return (
+        '<div style="font-family:monospace;margin-bottom:8px;font-size:0.9em">'
+        f'<b>Relevance:</b>&nbsp;{swatches}&nbsp;'
+        '<span style="color:#666">(hot → cold)</span>'
+        f'&nbsp;&nbsp;<span style="color:#888;font-size:0.85em">'
+        f'({elapsed_ms:.1f} ms · {n_words} words · σ={decay_sigma})</span>'
+        '</div>'
+    )
+
+
+def render_html(
+    words: List[str],
+    word_scores: np.ndarray,
+    query: str,
+    decay_sigma: float,
+    elapsed_ms: float,
+    quantize_levels: int = 64,
+    max_height_px: int = 540,
+) -> str:
+    """Render a heatmap-colored document as an HTML string.
+
+    Parameters
+    ----------
+    words : list of str
+        The full document, tokenized.
+    word_scores : np.ndarray, shape (n_words,), values in [0, 1]
+    query : str
+        Shown in the result header.
+    decay_sigma : float
+        Shown in the legend (informational).
+    elapsed_ms : float
+        Search latency in milliseconds, shown in the legend.
+    quantize_levels : int
+        Group consecutive words of similar score into a single <span> to
+        keep HTML compact. Higher = finer color gradient, larger HTML.
+    max_height_px : int
+        Scroll the document container if it exceeds this height.
+
+    Returns
+    -------
+    str
+        A self-contained HTML fragment. Safe to embed; words are escaped.
+    """
+    # Quantize scores so consecutive words of similar warmth share a span
+    def quant(s: float) -> int:
+        return int(s * quantize_levels)
+
+    spans: List[str] = []
+    cur_q: int | None = None
+    cur_words: List[str] = []
+    for i, w in enumerate(words):
+        q = quant(float(word_scores[i]))
+        if q == cur_q:
+            cur_words.append(w)
+        else:
+            if cur_words:
+                col = score_to_hex(cur_q / quantize_levels if cur_q is not None else 0.0)
+                spans.append(
+                    f'<span style="color:{col}">{html.escape(" ".join(cur_words))}</span>'
+                )
+            cur_q, cur_words = q, [w]
+    if cur_words:
+        col = score_to_hex(cur_q / quantize_levels if cur_q is not None else 0.0)
+        spans.append(
+            f'<span style="color:{col}">{html.escape(" ".join(cur_words))}</span>'
+        )
+
+    header = f'<h4 style="margin-bottom:4px">Results for: <em>{html.escape(query)}</em></h4>'
+    legend = _legend_html(len(words), decay_sigma, elapsed_ms)
+    doc = (
+        f'<div style="font-family:Georgia,serif;font-size:0.95em;line-height:1.8;'
+        f'white-space:pre-wrap;border:1px solid #ddd;border-radius:6px;'
+        f'padding:16px;max-height:{max_height_px}px;overflow-y:auto">'
+        + " ".join(spans)
+        + '</div>'
+    )
+    return header + legend + doc

doclighter/scoring.py

@@ -0,0 +1,118 @@
+"""Proximity-decayed word-level scoring.
+
+This is the algorithmic heart of Doclighter. Given per-chunk semantic scores,
+we compute a per-word relevance score by letting each chunk's score "radiate"
+outward with exponential decay:
+
+    word_score[w] = max over chunks c of  raw[c] * exp(-distance(w, c) / sigma)
+
+Distance is measured in words to the nearest edge of the chunk (0 if inside).
+Max-aggregation (rather than sum) means a word's color reflects its single
+strongest semantic neighbor, which matches the visual intuition: a word is
+"hot" if any nearby region is hot, not because many lukewarm regions sum up.
+"""
+from __future__ import annotations
+
+from typing import List
+
+import numpy as np
+
+from .chunking import Chunk
+
+
+def word_heatmap(
+    chunk_scores: np.ndarray,
+    chunks: List[Chunk],
+    n_words: int,
+    decay_sigma: float = 20.0,
+    threshold_frac: float = 0.05,
+) -> np.ndarray:
+    """Spread chunk scores into per-word scores via exponential decay.
+
+    Parameters
+    ----------
+    chunk_scores : np.ndarray, shape (n_chunks,)
+        Semantic similarity score per chunk. Negatives are clipped to 0
+        (irrelevant regions shouldn't radiate warmth).
+    chunks : list of Chunk
+        Word-span metadata for each chunk.
+    n_words : int
+        Total number of words in the document.
+    decay_sigma : float
+        Decay length scale in words. exp(-d/sigma): at d=sigma, weight = 0.37;
+        at d=2*sigma, weight = 0.14. Small sigma (~5) = sharp word-level
+        highlights. Large sigma (~100) = broad thematic regions.
+    threshold_frac : float
+        Skip chunks with score < threshold_frac * max_score. Pure optimization:
+        chunks with negligible scores can't affect the final heatmap.
+
+    Returns
+    -------
+    np.ndarray, shape (n_words,), values in [0, 1]
+        Per-word relevance score, normalized so max = 1.
+    """
+    raw = np.maximum(chunk_scores, 0.0).astype(np.float32)
+    if raw.max() == 0:
+        return np.zeros(n_words, dtype=np.float32)
+
+    threshold = float(raw.max()) * threshold_frac
+    word_idx = np.arange(n_words, dtype=np.float32)
+    word_scores = np.zeros(n_words, dtype=np.float32)
+
+    for ci, chunk in enumerate(chunks):
+        score = raw[ci]
+        if score < threshold:
+            continue
+        # Distance from each word to nearest edge of this chunk
+        # (0 inside the chunk, positive outside)
+        dist = np.where(
+            word_idx < chunk.start,
+            chunk.start - word_idx,
+            np.where(word_idx >= chunk.end, word_idx - (chunk.end - 1), 0.0),
+        )
+        influence = score * np.exp(-dist / decay_sigma)
+        np.maximum(word_scores, influence, out=word_scores)
+
+    mx = word_scores.max()
+    if mx > 0:
+        word_scores /= mx
+    return word_scores
+
+
+def aggregate_multi_query(
+    per_query_heatmaps: List[np.ndarray],
+    mode: str = "max",
+) -> np.ndarray:
+    """Combine word-level heatmaps from multiple sub-queries.
+
+    Parameters
+    ----------
+    per_query_heatmaps : list of np.ndarray
+        Each shape (n_words,), values in [0, 1].
+    mode : {"max", "mean", "sum"}
+        "max" — strongest single match wins per word (default; closest to
+        ColBERT MaxSim spirit).
+        "mean" — average influence across queries.
+        "sum" — additive, then renormalized. Useful when you want words that
+        match *multiple* sub-queries to outshine words matching just one.
+
+    Returns
+    -------
+    np.ndarray, shape (n_words,), values in [0, 1]
+    """
+    if not per_query_heatmaps:
+        raise ValueError("per_query_heatmaps must be non-empty")
+    stack = np.stack(per_query_heatmaps, axis=0)
+    if mode == "max":
+        combined = stack.max(axis=0)
+    elif mode == "mean":
+        combined = stack.mean(axis=0)
+    elif mode == "sum":
+        combined = stack.sum(axis=0)
+    else:
+        raise ValueError(f"mode must be 'max', 'mean', or 'sum'; got {mode!r}")
+
+    mx = combined.max()
+    if mx > 0:
+        combined = combined / mx
+    return combined.astype(np.float32)

doclighter-0.1.0.dist-info/METADATA

@@ -0,0 +1,245 @@
+Metadata-Version: 2.4
+Name: doclighter
+Version: 0.1.0
+Summary: A semantic Ctrl+F that paints your document with a relevance gradient.
+Author-email: Pratyush <pratyush272@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/pratyush272/doclighter
+Project-URL: Repository, https://github.com/pratyush272/doclighter
+Project-URL: Issues, https://github.com/pratyush272/doclighter/issues
+Keywords: semantic-search,embeddings,visualization,rag,nlp,pdf
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Classifier: Topic :: Text Processing :: Indexing
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.21
+Requires-Dist: sentence-transformers>=2.2
+Requires-Dist: pypdf>=4.0
+Requires-Dist: requests>=2.25
+Provides-Extra: quantize
+Requires-Dist: faiss-cpu>=1.7; extra == "quantize"
+Provides-Extra: streamlit
+Requires-Dist: streamlit>=1.28; extra == "streamlit"
+Provides-Extra: dev
+Requires-Dist: pytest>=7; extra == "dev"
+Requires-Dist: pytest-cov>=4; extra == "dev"
+Requires-Dist: ruff>=0.1; extra == "dev"
+Provides-Extra: all
+Requires-Dist: doclighter[dev,quantize,streamlit]; extra == "all"
+Dynamic: license-file
+
+# doclighter
+
+[![tests](https://github.com/pratyush272/doclighter/actions/workflows/test.yml/badge.svg)](https://github.com/pratyush272/doclighter/actions/workflows/test.yml)
+[![PyPI](https://img.shields.io/pypi/v/doclighter.svg)](https://pypi.org/project/doclighter/)
+[![Python](https://img.shields.io/pypi/pyversions/doclighter.svg)](https://pypi.org/project/doclighter/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+**A semantic Ctrl+F that paints your document with a relevance gradient.**
+
+`doclighter` is what you reach for when you need to *see* where a topic lives in a document — not be told an answer. It embeds your document at fine granularity, then projects query relevance back onto every word as a heatmap. No LLM, no hallucination, no top-K cliff.
+
+```python
+from doclighter import Doclighter
+
+doc = Doclighter.from_pdf("contract.pdf")
+result = doc.search("termination clauses")
+result  # In Jupyter: renders the whole document, color-coded by relevance
+```
+
+## Why this exists
+
+Traditional RAG hands an LLM the top 3–10 chunks and asks it to generate an answer. That's great when you trust the answer and don't need to read the source. But sometimes you *need to read the source* — legal review, paper skimming, contract diffing, due diligence — and you want a tool that helps you *navigate* a long document, not summarize it away.
+
+`doclighter` is for that. It treats the whole document as the output, and re-colors it by semantic relevance to your query. Hot regions deserve your attention; cold regions you can skim past.
+
+It's deterministic, fast (sub-100ms per query after indexing), and shows you the long tail — including the case where your query *doesn't* match anything (everything stays cold blue, which is itself useful information that RAG hides).
+
+## Install
+
+```bash
+pip install doclighter
+```
+
+Optional extras:
+```bash
+pip install "doclighter[quantize]"   # FAISS SQ8 index for very large docs
+pip install "doclighter[streamlit]"  # for the interactive demo app
+pip install "doclighter[dev]"        # for contributors
+```
+
+## Quickstart
+
+### Load a document
+
+```python
+from doclighter import Doclighter
+
+# From a PDF on disk
+doc = Doclighter.from_pdf("contract.pdf")
+
+# From a PDF URL
+doc = Doclighter.from_url("https://example.com/contract.pdf")
+
+# From raw text
+doc = Doclighter.from_text(open("paper.txt").read())
+```
+
+The first call downloads the default embedding model (~80 MB MiniLM) and embeds your document. For a ~10K word doc this takes ~25 seconds. Subsequent searches reuse the index.
+
+### Search
+
+```python
+result = doc.search("termination clauses")
+
+result.word_scores      # numpy array, shape (n_words,), values in [0, 1]
+result.top_chunks(k=10) # list of (chunk_text, score, (start, end))
+result.elapsed_ms       # ~10-50ms for typical docs
+result.to_html()        # HTML string for display anywhere
+```
+
+In Jupyter, just put `result` on the last line of a cell — it renders the heatmap inline.
+
+### Zoom: the `decay_sigma` knob
+
+The differentiating feature. `decay_sigma` controls how far semantic warmth spreads from a matched region:
+
+```python
+narrow = doc.search("termination", decay_sigma=5.0)   # sharp word-level highlights
+broad  = doc.search("termination", decay_sigma=80.0)  # broad thematic regions
+```
+
+Same index, no re-embedding. Drag the σ slider in the Streamlit demo to feel what this does.
+
+### Multi-query
+
+```python
+result = doc.search(
+    ["termination", "indemnification", "labour wages"],
+    multi_query_aggregate="max",  # or "sum" to favor regions matching multiple
+)
+```
+
+### Save / load the index
+
+Embedding is the slow step. Save once, reuse:
+
+```python
+doc.save("contract.idx")
+doc = Doclighter.load("contract.idx")
+```
+
+### Bring your own embedder
+
+Any callable mapping `list[str] -> np.ndarray` of shape `(N, dim)` works:
+
+```python
+from sentence_transformers import SentenceTransformer
+
+bge = SentenceTransformer("BAAI/bge-small-en-v1.5")
+doc = Doclighter.from_text(text, embedder=bge.encode)
+```
+
+## Streamlit demo
+
+```bash
+pip install "doclighter[streamlit]"
+streamlit run examples/streamlit_app.py
+```
+
+A working UI with PDF upload, query box, σ slider, and live re-rendering.
+
+## How it works
+
+1. **Chunk** the document into small rolling windows (default: 12 words, 50% overlap).
+2. **Embed** each chunk with sentence-transformers (default: `all-MiniLM-L6-v2`).
+3. **Score** each chunk against your query via cosine similarity.
+4. **Project** chunk scores back onto every word via exponential proximity decay:
+
+   ```
+   word_score[w] = max over chunks c of  raw[c] × exp(-distance(w, c) / sigma)
+   ```
+
+5. **Render** the document as colored HTML — words inherit warmth from their nearest semantically matched chunk.
+
+Step 4 is the interesting one. Max-aggregation (rather than sum) means a word's color reflects its single strongest semantic neighbor — visually intuitive and resistant to "many lukewarm chunks add up to red" noise.
+
+## How this compares to RAG
+
+`doclighter` and RAG solve different problems:
+
+| | RAG | doclighter |
+|---|---|---|
+| Output | LLM-generated answer | Document, recolored |
+| Best for | "What's the answer?" | "Where in the doc?" |
+| When query has no answer | LLM hedges / hallucinates | Document stays cold (honest) |
+| Hides chunk boundaries | No — top-K cliff | Yes — gradient smooths over |
+| Cost per query | LLM tokens | Free (one matmul) |
+| Determinism | Sampling-dependent | Fully deterministic |
+
+They're complementary. Use RAG when you trust the LLM with the question; use `doclighter` when you need to read the source yourself.
+
+The algorithmic kernel is conceptually related to [ColBERT](https://github.com/stanford-futuredata/ColBERT)'s MaxSim late-interaction, but applied to *visualization* rather than ranking, and at the word-level rather than the token-level.
+
+## API reference
+
+### `Doclighter(text, **kwargs)`
+
+| Parameter | Default | Description |
+|---|---|---|
+| `text` | required | The document as a string |
+| `chunk_size` | `12` | Words per rolling window |
+| `chunk_overlap` | `0.5` | Fraction overlap between windows |
+| `embedder` | `None` | Custom embedder callable (default: MiniLM) |
+| `embedding_model` | `"all-MiniLM-L6-v2"` | sentence-transformers model name |
+| `decay_sigma` | `20.0` | Default proximity decay scale (words) |
+| `quantize` | `False` | Use SQ8 FAISS index instead of flat exact |
+| `quantize_rerank_k` | `200` | If quantize=True, rerank top-K with exact |
+
+Alternate constructors: `Doclighter.from_text(...)`, `Doclighter.from_pdf(path, ...)`, `Doclighter.from_url(url, ...)`.
+
+### `doc.search(query, **kwargs) -> SearchResult`
+
+| Parameter | Default | Description |
+|---|---|---|
+| `query` | required | A string, or list of strings for multi-query |
+| `decay_sigma` | `None` | Override the doc's default sigma |
+| `multi_query_aggregate` | `"max"` | `"max"`, `"mean"`, or `"sum"` for multi-query |
+
+### `SearchResult`
+
+| Attribute / method | Returns |
+|---|---|
+| `.word_scores` | `np.ndarray` of shape `(n_words,)`, in `[0, 1]` |
+| `.chunk_scores` | `np.ndarray` of raw per-chunk cosine scores |
+| `.top_chunks(k=10)` | `list[(text, score, (start, end))]` |
+| `.to_html(**kwargs)` | HTML string of the heatmap-colored document |
+| `.elapsed_ms` | Search latency |
+
+## Development
+
+```bash
+git clone https://github.com/pratyush272/doclighter
+cd doclighter
+pip install -e ".[dev]"
+pytest
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE).
+
+## Citation / acknowledgement
+
+If you use `doclighter` in research, a link back is appreciated. The proximity-decay scoring idea borrows from passage-retrieval literature; max-aggregation over fine-grained matches is in spirit closest to ColBERT.

doclighter-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+doclighter/__init__.py,sha256=kzREXwu7PcSRSSEOKxefkl7wqq3Tf92iRZ8GyrJ8aws,522
+doclighter/chunking.py,sha256=Zaicqk9eNpkPlulCJPB_TkIOGCSag7p6bLoEG9jN42o,1823
+doclighter/core.py,sha256=rz2ZOIB9FY1ibtZnn5qvWZCasy-KpD2wOpfOGA0chso,10271
+doclighter/embedding.py,sha256=hd2e5K75HvBC0BM83z6gwjMqXwuoS6tgDzRffknNpE0,1816
+doclighter/extract.py,sha256=NZoHDD8VaeqUeuw6Mszb91Qp375k-5z9gNsf54lPKDo,1201
+doclighter/index.py,sha256=AHCS4w3FR-zq8JEfzXMF81CToswkuZJkk7ffqSElnnk,2830
+doclighter/render.py,sha256=fzkqYC7tfyW6SgtUeG9VPnc4Q-5yamWJUpTeOyh_M80,4473
+doclighter/scoring.py,sha256=MZGGIFuWK3nW3hGpX38aBUjiDa9uq1nBe1WXMLzZXwI,4048
+doclighter-0.1.0.dist-info/licenses/LICENSE,sha256=k412KnI3imf_ScR1LlhRFjnHFq62gZSbHO88cwa_Ljk,1065
+doclighter-0.1.0.dist-info/METADATA,sha256=V3Pjm6w2MtLzEQsHi3RuqF4Eb5yu1jDxeZNNPr2heYc,9401
+doclighter-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+doclighter-0.1.0.dist-info/top_level.txt,sha256=Sd9T52y2LYw0Vz5aIRk_C3rD6JAktygliw5KakMp4TA,11
+doclighter-0.1.0.dist-info/RECORD,,

doclighter-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

doclighter-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Pratyush
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

doclighter-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+doclighter