Upload folder using huggingface_hub

app/config.py

@@ -47,7 +47,12 @@ class Settings(BaseSettings):
     europepmc_url: str = "https://www.ebi.ac.uk/europepmc/webservices/rest"
     fulltext_timeout: float = 20.0
     fulltext_max_lines: int = Field(default=160, ge=20, le=600)
-    # Gap map: every embedded study, scanned in pages and clustered once.
+    # Gap map: the studies nearest the question, fetched by kNN and clustered per
+    # question. A few hundred nodes is what a reader can take in and what a force
+    # layout can settle in a second; the full-index scan below only runs when
+    # there is no question to center on.
+    gapmap_neighborhood: int = Field(default=800, ge=50, le=5000)
+    # Whole-index build (no question): every embedded study, scanned in pages.
     # The region count is a resolution, and what counts as an empty band depends on
     # it: measured on 6000 documents of the live index, 24 regions leave the corpus
     # continuous (no band under the occupancy threshold) while 80 keep a median of
@@ -68,6 +73,12 @@ class Settings(BaseSettings):
     # regions; a browser cannot paint two million marks, so the drawn subset is
     # thinned by a hash of the document id.
     gapmap_points: int = Field(default=6000, ge=100, le=50000)
+    # Edges the canvas draws between drawn studies: each study joins its most
+    # similar few, and only when the cosine reaches the floor. The cap is what
+    # keeps a dense literature from becoming one hub; the floor is what leaves an
+    # isolated paper isolated instead of tethered to its nearest stranger.
+    gapmap_edge_neighbors: int = Field(default=3, ge=1, le=12)
+    gapmap_edge_floor: float = Field(default=0.8, ge=-1.0, le=1.0)
     # Rows the 2-D projection is fitted on. The basis is a drawing choice, and a
     # full-corpus SVD would stall the build for minutes to move points by pixels.
     # Nothing can be drawn before it exists, so it is fitted on the first pages

app/gapmap.py

@@ -233,6 +233,44 @@ def project_with(mean: np.ndarray, basis: np.ndarray, vectors: np.ndarray) -> np
     return (np.atleast_2d(np.asarray(vectors, dtype=np.float32)) - mean) @ basis.T
 
 
+def neighbor_edges(
+    vectors: np.ndarray,
+    start: int = 0,
+    neighbors: int = 3,
+    floor: float = 0.0,
+    chunk: int = 1024,
+) -> list[tuple[int, int, float]]:
+    """Cosine-nearest-neighbor edges among unit ``vectors``.
+
+    Each row from ``start`` on is joined to its ``neighbors`` most similar rows
+    anywhere in the matrix, provided the cosine reaches ``floor``; rows before
+    ``start`` already had their turn on an earlier call, so a stream can call
+    this once per frame with only the new rows and never repeats an edge. Pairs
+    are returned once each as ``(i, j, cosine)`` with ``i < j``, indices being
+    positions in ``vectors``. The degree cap is what keeps a dense literature
+    from becoming one hub; the floor is what keeps an isolated paper isolated.
+    """
+    total = len(vectors)
+    if total < 2 or start >= total or neighbors < 1:
+        return []
+    matrix = np.asarray(vectors, dtype=np.float32)
+    edges: dict[tuple[int, int], float] = {}
+    for begin in range(start, total, chunk):
+        end = min(begin + chunk, total)
+        scores = matrix[begin:end] @ matrix.T
+        scores[np.arange(end - begin), np.arange(begin, end)] = -np.inf
+        top = min(neighbors, total - 1)
+        candidates = np.argpartition(-scores, top - 1, axis=1)[:, :top]
+        for offset, row in enumerate(range(begin, end)):
+            for column in candidates[offset].tolist():
+                cosine = float(scores[offset, column])
+                if cosine < floor:
+                    continue
+                key = (row, column) if row < column else (column, row)
+                edges[key] = cosine
+    return [(i, j, cosine) for (i, j), cosine in edges.items()]
+
+
 def project(vectors: np.ndarray) -> np.ndarray:
     """2-D coordinates for ``vectors`` under a basis fitted on themselves."""
     if not len(vectors):

app/gapmap_service.py

@@ -1,18 +1,21 @@
-"""Serves the gap map: cluster the whole embedded index, describe it, place ideas on it.
-
-The map is a property of the corpus, not of a query, so it is computed once and
-cached. Placing an idea then costs one query embedding and two matrix products,
-which is what makes the "where does this land" answer instant while still being
-computed over every embedded paper rather than the ten a chat model can read.
-
-Covering the whole index means the build is long enough to be worth watching, so
-it is written as a stream: documents are scanned in pages, provisional centroids
-move as pages arrive, and full-corpus k-means then runs to convergence, emitting
-the real intermediate state after every iteration. Nothing emitted is
-interpolated or replayed — a partial map is labelled partial, and only the final
-event describes the whole corpus. Measured on the live index (2,096,736 embedded
-studies): 264 s end to end — roughly 130 s to scan at sixteen slices and 25
-bounded k-means iterations over the whole corpus in the rest, peaking at 6.9 GB.
+"""Serves the gap map: cluster the studies around a question, describe them, place the question.
+
+A map with a question is built from that question's neighborhood: the
+``gapmap_neighborhood`` studies nearest its vector, fetched by kNN, clustered
+and drawn. That is a few hundred documents, so the build takes about a second
+and the picture is one a reader can take in. Its coverage says exactly that —
+the N nearest studies out of the embedded index — and never more.
+
+Without a question there is nothing to center on, so the map falls back to the
+whole embedded index. Covering the whole index means the build is long enough
+to be worth watching, so it is written as a stream: documents are scanned in
+pages, provisional centroids move as pages arrive, and full-corpus k-means then
+runs to convergence, emitting the real intermediate state after every iteration.
+Nothing emitted is interpolated or replayed — a partial map is labelled partial,
+and only the final event describes the whole corpus. Measured on the live index
+(2,096,736 embedded studies): 264 s end to end — roughly 130 s to scan at sixteen
+slices and 25 bounded k-means iterations over the whole corpus in the rest,
+peaking at 6.9 GB.
 """
 
 from __future__ import annotations
@@ -38,6 +41,7 @@ from app.gapmap import (
     find_gaps,
     fit_projection,
     kmeans_steps,
+    neighbor_edges,
     normalize,
     place,
     project_with,
@@ -48,6 +52,9 @@ logger = logging.getLogger(__name__)
 
 Progress = Callable[[dict], Awaitable[None]]
 
+# Question maps kept in memory beside the corpus map.
+MAP_CACHE = 16
+
 
 def _public(region: dict) -> dict:
     """Region without its centroid: a 384-float vector is not UI payload."""
@@ -143,10 +150,11 @@ class GapMapService:
         self.repo = repository
         self.config = config
         self.embedder = QueryEmbedder(config)
-        self._built: dict | None = None
+        # One cached map per question (keyed by its vector) plus the corpus map.
+        self._maps: dict[bytes | None, dict] = {}
         self._lock = asyncio.Lock()
         self._subscribers: set[asyncio.Queue] = set()
-        self._task: asyncio.Task | None = None
+        self._tasks: dict[bytes | None, asyncio.Task] = {}
         # The plane the current build is drawing on, so a streamed question can be
         # placed on the same one before the build finishes.
         self._basis: tuple[np.ndarray, np.ndarray] | None = None
@@ -162,26 +170,94 @@ class GapMapService:
                 continue
             queue.put_nowait(payload)
 
-    async def build(self, refresh: bool = False) -> dict:
-        """The cached map, building it (once, shared) if it is not there yet.
+    @property
+    def _built(self) -> dict | None:
+        """The cached corpus map, if it has been built."""
+        return self._maps.get(None)
+
+    async def build(self, refresh: bool = False, vector: np.ndarray | None = None) -> dict:
+        """The cached map for this question (or the corpus), building it once if needed.
 
         The build is a task of its own rather than a coroutine of whichever
         request arrived first: a viewer who closes the tab after twenty seconds
-        must not throw away a three-minute clustering run that other viewers,
-        and the cache, are waiting for.
+        must not throw away a clustering run that other viewers, and the cache,
+        are waiting for.
         """
-        if self._built is not None and not refresh:
-            return self._built
+        key = None if vector is None else np.asarray(vector, dtype=np.float32).tobytes()
+        if key in self._maps and not refresh:
+            return self._maps[key]
         async with self._lock:
-            if self._built is not None and not refresh:
-                return self._built
-            if self._task is None or self._task.done():
-                self._task = asyncio.create_task(self._build())
-            task = self._task
-        self._built = await asyncio.shield(task)
-        return self._built
-
-    async def _build(self) -> dict:
+            if key in self._maps and not refresh:
+                return self._maps[key]
+            task = self._tasks.get(key)
+            if task is None or task.done():
+                task = asyncio.create_task(self._build(vector))
+                self._tasks[key] = task
+        built = await asyncio.shield(task)
+        if len(self._maps) >= MAP_CACHE and key not in self._maps:
+            for stale in [k for k in self._maps if k is not None][: len(self._maps) - MAP_CACHE + 1]:
+                del self._maps[stale]
+        self._maps[key] = built
+        return built
+
+    async def _build(self, vector: np.ndarray | None) -> dict:
+        if vector is None:
+            return await self._build_corpus()
+        return await self._build_neighborhood(vector)
+
+    async def _build_neighborhood(self, vector: np.ndarray) -> dict:
+        """Fetch the studies nearest the question, cluster them, describe the result."""
+        counting = asyncio.create_task(self.repo.count_embedded())
+        size = self.config.gapmap_neighborhood
+        query = normalize(np.asarray(vector, dtype=np.float32)[None, :])[0]
+        documents = await self.repo.knn_studies(query.tolist(), size)
+        build = MapBuild(0, capacity=size)
+        build.add(documents)
+        stride = max(1, len(build) // max(1, self.config.gapmap_points))
+        drawn = [
+            row
+            for row in range(len(build))
+            if _drawn(str(build.documents[row].get("id", row)), stride)
+        ]
+        basis: tuple[np.ndarray, np.ndarray] | None = None
+        if len(build):
+            basis = await asyncio.to_thread(fit_projection, build.vectors)
+            self._basis = basis
+        corpus = max(0, await counting)
+        build.corpus = max(corpus, len(build))
+        scope = {"scope": "neighborhood", "neighborhood": size}
+        centroids = None
+        labels = np.zeros(len(build), dtype=np.int64)
+        iteration = 0
+        emitted = 0
+        if len(build):
+            steps = kmeans_steps(
+                build.vectors,
+                self.config.gapmap_regions,
+                seed=self.config.gapmap_seed,
+                iterations=self.config.gapmap_iterations,
+            )
+            while True:
+                step = await asyncio.to_thread(next, steps, None)
+                if step is None:
+                    break
+                centroids, labels = step
+                iteration += 1
+                if basis is not None:
+                    emitted = await self._emit_partial(
+                        build,
+                        centroids,
+                        basis,
+                        drawn,
+                        emitted,
+                        len(build),
+                        iteration=iteration,
+                        labels=labels,
+                        extra={"coverage": scope},
+                    )
+        return await self._finish(build, centroids, labels, basis, drawn, iteration, scope)
+
+    async def _build_corpus(self) -> dict:
         """Scan the embedded corpus, cluster it, and describe the result."""
         # Counted beside the scan rather than before it: the count is a query of
         # its own over two million documents, and nothing can be drawn while it
@@ -283,6 +359,22 @@ class GapMapService:
                         labels=labels,
                     )
 
+        return await self._finish(
+            build, centroids, labels, basis, drawn, iteration, {"scope": "corpus"}
+        )
+
+    async def _finish(
+        self,
+        build: MapBuild,
+        centroids: np.ndarray | None,
+        labels: np.ndarray,
+        basis: tuple[np.ndarray, np.ndarray] | None,
+        drawn: list[int],
+        iteration: int,
+        scope: dict,
+    ) -> dict:
+        """Describe a clustered build: regions, gaps, drawn points, edges, coverage."""
+        corpus = build.corpus
         clusters_used = len(centroids) if centroids is not None else 0
         regions = await asyncio.to_thread(
             describe_regions, build.documents, build.vectors, labels, clusters_used
@@ -290,7 +382,8 @@ class GapMapService:
         gaps = await asyncio.to_thread(
             find_gaps, build.documents, regions, vectors=build.vectors
         )
-        points = self._points(build, regions, labels, basis, drawn)
+        points, kept = self._points(build, regions, labels, basis, drawn)
+        edges = await asyncio.to_thread(self._edges, build, kept)
         if basis is not None and regions:
             centre, plane = basis
             positions = project_with(
@@ -306,6 +399,7 @@ class GapMapService:
             "regions": regions,
             "gaps": gaps,
             "points": points,
+            "edges": edges,
             "basis": basis,
             "coverage": {
                 "clustered": len(build),
@@ -313,11 +407,13 @@ class GapMapService:
                 "regions": len(regions),
                 "drawn": len(points),
                 "complete": True,
+                **scope,
             },
             "version": MAP_VERSION,
         }
         logger.info(
-            "Gap map built over %s of %s embedded studies in %s regions after %s iterations",
+            "Gap map (%s) built over %s of %s embedded studies in %s regions after %s iterations",
+            scope.get("scope"),
             len(build),
             corpus,
             len(regions),
@@ -336,15 +432,18 @@ class GapMapService:
         target: int,
         iteration: int | None = None,
         labels: np.ndarray | None = None,
+        publish: Progress | None = None,
+        extra: dict | None = None,
     ) -> int:
         """Publish the map as it stands: new points, where the regions are now.
 
         Only the points the canvas has not seen are sent in full; the ones it
         already holds are recoloured through ``pointRegions``, because a region
         a point belongs to changes on every k-means iteration while its position
-        does not.
+        does not. Frames go to ``publish`` when given (one viewer's own build)
+        and to every subscriber otherwise (the shared corpus build).
         """
-        if not self._subscribers:
+        if publish is None and not self._subscribers:
             return emitted
         centre, plane = basis
         rows = np.array(drawn, dtype=np.int64)
@@ -354,10 +453,12 @@ class GapMapService:
             member = assign(build.vectors[rows], centroids)
         else:
             member = labels[rows]
-        fresh = []
+        fresh: list[dict] = []
+        edges: list[list[float]] = []
         if len(rows) > emitted:
             new_rows = rows[emitted:]
             coords = project_with(centre, plane, build.vectors[new_rows])
+            edges = await asyncio.to_thread(self._edges, build, rows.tolist(), emitted)
             for offset, row in enumerate(new_rows.tolist()):
                 document = build.documents[row]
                 title = str(document.get("title") or "")
@@ -374,8 +475,9 @@ class GapMapService:
                 )
         positions = project_with(centre, plane, centroids)
         counts = np.bincount(member, minlength=len(centroids))
-        await self._publish(
+        await (publish or self._publish)(
             {
+                **(extra or {}),
                 "event": "progress",
                 "stage": "clustering" if iteration else "scanning",
                 "iteration": iteration,
@@ -386,8 +488,10 @@ class GapMapService:
                     "drawn": len(rows),
                     "target": target,
                     "complete": False,
+                    **(extra or {}).get("coverage", {}),
                 },
                 "points": fresh,
+                "edges": edges,
                 "pointRegions": [int(value) for value in member.tolist()],
                 "regions": [
                     {
@@ -410,19 +514,22 @@ class GapMapService:
         labels: np.ndarray,
         basis: tuple[np.ndarray, np.ndarray] | None,
         drawn: list[int],
-    ) -> list[dict]:
-        """2-D positions for the drawn subset, each carrying its region."""
+    ) -> tuple[list[dict], list[int]]:
+        """2-D positions for the drawn subset, each carrying its region.
+
+        Also returns the rows drawn, in point order, so edges can be indexed
+        the same way the canvas indexes points.
+        """
         if basis is None or not regions or not drawn:
-            return []
+            return [], []
         centre, plane = basis
-        rows = np.array(drawn, dtype=np.int64)
-        coords = project_with(centre, plane, build.vectors[rows])
         known = {region["id"] for region in regions}
+        kept = [row for row in drawn if int(labels[row]) in known]
+        rows = np.array(kept, dtype=np.int64)
+        coords = project_with(centre, plane, build.vectors[rows])
         points = []
         for offset, row in enumerate(rows.tolist()):
             region = int(labels[row])
-            if region not in known:
-                continue
             document = build.documents[row]
             title = str(document.get("title") or "")
             points.append(
@@ -436,7 +543,23 @@ class GapMapService:
                     "year": document.get("year") if isinstance(document.get("year"), int) else None,
                 }
             )
-        return points
+        return points, kept
+
+    def _edges(self, build: MapBuild, drawn: list[int], start: int = 0) -> list[list[float]]:
+        """Similarity edges among drawn points, indexed by position in ``drawn``.
+
+        ``start`` is how many of them already have their edges, so a streamed
+        frame only pays for the rows it adds. Each edge is ``[i, j, cosine]``.
+        """
+        if len(drawn) < 2:
+            return []
+        found = neighbor_edges(
+            build.vectors[np.array(drawn, dtype=np.int64)],
+            start=start,
+            neighbors=self.config.gapmap_edge_neighbors,
+            floor=self.config.gapmap_edge_floor,
+        )
+        return [[i, j, round(cosine, 3)] for i, j, cosine in found]
 
     # ---------------------------------------------------------------- serving
 
@@ -453,15 +576,10 @@ class GapMapService:
             placement["point"] = {"x": float(coords[0, 0]), "y": float(coords[0, 1])}
         return placement
 
-    async def _place_expression(
-        self, expression: MapArithmetic, built: dict, warnings: list[str]
-    ) -> dict | None:
-        """Place ``start − remove + add``: embed each phrase, combine, locate.
-
-        The combined vector is treated like any other point on the map — what
-        comes back are the real papers nearest it, so the honesty of the answer
-        does not depend on how cleanly the analogy worked.
-        """
+    async def _expression_vector(
+        self, expression: MapArithmetic, warnings: list[str]
+    ) -> np.ndarray | None:
+        """The unit vector of ``start − remove + add``: embed each phrase, combine."""
         terms = [expression.start, *expression.remove, *expression.add]
         vectors = await asyncio.gather(*(self.embed(term) for term in terms))
         if any(vector is None for vector in vectors):
@@ -475,7 +593,16 @@ class GapMapService:
             warnings.append(
                 "The expression cancels itself out; the combined vector has no direction."
             )
-            return None
+        return combined
+
+    async def _place_expression(
+        self, expression: MapArithmetic, combined: np.ndarray | list[float], built: dict
+    ) -> dict:
+        """Place the combined vector like any other point on the map.
+
+        What comes back are the real papers nearest it, so the honesty of the
+        answer does not depend on how cleanly the analogy worked.
+        """
         placement = await asyncio.to_thread(
             place,
             combined,
@@ -498,10 +625,23 @@ class GapMapService:
         cutoff_year: int | None = None,
         refresh: bool = False,
     ) -> dict:
-        built = await self.build(refresh=refresh)
         warnings: list[str] = []
+        vector: np.ndarray | list[float] | None = None
+        if arithmetic is not None:
+            vector = await self._expression_vector(arithmetic, warnings)
+        elif idea:
+            vector = await self.embed(idea)
+            if vector is None:
+                warnings.append("Embeddings are disabled, so the idea could not be placed.")
+        built = await self.build(refresh=refresh, vector=vector)
         coverage = built["coverage"]
-        if coverage["corpus"] > coverage["clustered"]:
+        if coverage.get("scope") == "neighborhood":
+            warnings.append(
+                f"The map shows the {coverage['clustered']} embedded studies nearest this "
+                f"question, out of {coverage['corpus']} in the index: the question's "
+                "neighborhood, not the whole index."
+            )
+        elif coverage["corpus"] > coverage["clustered"]:
             warnings.append(
                 f"The map clusters {coverage['clustered']} of {coverage['corpus']} "
                 "embedded studies, not the whole index."
@@ -512,22 +652,18 @@ class GapMapService:
                 f"draws {coverage['drawn']} of them."
             )
         placement = None
-        if arithmetic is not None:
-            placement = await self._place_expression(arithmetic, built, warnings)
-        elif idea:
-            vector = await self.embed(idea)
-            if vector is None:
-                warnings.append("Embeddings are disabled, so the idea could not be placed.")
-            else:
-                placement = await asyncio.to_thread(
-                    place,
-                    vector,
-                    built["documents"],
-                    built["regions"],
-                    built["gaps"],
-                    vectors=built["vectors"],
-                )
-                placement = self._decorate(placement, vector, built)
+        if vector is not None and arithmetic is not None:
+            placement = await self._place_expression(arithmetic, vector, built)
+        elif vector is not None:
+            placement = await asyncio.to_thread(
+                place,
+                vector,
+                built["documents"],
+                built["regions"],
+                built["gaps"],
+                vectors=built["vectors"],
+            )
+            placement = self._decorate(placement, vector, built)
         calibration = None
         if cutoff_year is not None:
             calibration = await asyncio.to_thread(
@@ -544,6 +680,7 @@ class GapMapService:
             "regions": [_public(region) for region in built["regions"]],
             "gaps": built["gaps"],
             "points": built["points"],
+            "edges": built.get("edges", []),
             "placement": placement,
             "calibration": calibration,
             "warnings": warnings,
@@ -584,9 +721,12 @@ class GapMapService:
         self, idea: str | None, arithmetic: MapArithmetic | None
     ) -> np.ndarray | None:
         """The unit vector of this viewer's question, if there is one to embed."""
-        if arithmetic is not None or not idea:
+        if arithmetic is not None:
+            vector = await self._expression_vector(arithmetic, [])
+        elif idea:
+            vector = await self.embed(idea)
+        else:
             return None
-        vector = await self.embed(idea)
         if vector is None:
             return None
         return normalize(np.asarray(vector, dtype=np.float32)[None, :])[0]