convmemory 0.4.0__py3-none-any.whl

convmemory/reranker.py

@@ -0,0 +1,253 @@
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+import numpy as np
+import torch
+
+from .scoring import (
+    build_memory_to_windows,
+    cosine_scores,
+    rerank_candidates,
+    score_ce_lite,
+    window_scores,
+)
+
+
+@dataclass
+class RerankConfig:
+    window_size: int = 5
+    stride: int = 1
+    candidate_top_n: int = 500
+    raw_weight: float = 0.0
+    dca_router_block_size: int = 32
+    lexical_features: bool = True
+    window_mode: str = "full"
+
+
+@dataclass
+class RerankResult:
+    memory_id: str
+    score: float
+    raw_score: float
+    rank: int
+    text: Optional[str] = None
+
+
+def sliding_windows(num_items: int, window_size: int, stride: int):
+    if num_items <= 0:
+        return []
+    if num_items <= window_size:
+        return [list(range(num_items))]
+
+    windows = []
+    for start in range(0, num_items - window_size + 1, stride):
+        windows.append(list(range(start, start + window_size)))
+    last = list(range(num_items - window_size, num_items))
+    if windows[-1] != last:
+        windows.append(last)
+    return windows
+
+
+def normalize_rows(matrix):
+    matrix = np.asarray(matrix, dtype=np.float32)
+    norms = np.linalg.norm(matrix, axis=1, keepdims=True) + 1e-8
+    return matrix / norms
+
+
+def window_tensor(memory_embeddings, windows):
+    if not windows:
+        return torch.zeros((0, 0, memory_embeddings.shape[1]), dtype=torch.float32)
+    if all(len(window) == len(windows[0]) for window in windows):
+        indices = np.asarray(windows, dtype=np.int64)
+        return torch.as_tensor(memory_embeddings[indices], dtype=torch.float32)
+    max_len = max(len(window) for window in windows)
+    batch = np.zeros((len(windows), max_len, memory_embeddings.shape[1]), dtype=np.float32)
+    for i, window in enumerate(windows):
+        batch[i, : len(window)] = memory_embeddings[window]
+    return torch.tensor(batch, dtype=torch.float32)
+
+
+def candidate_local_windows(num_items, candidate_indices, window_size):
+    if num_items <= 0:
+        return []
+    if num_items <= window_size:
+        return [list(range(num_items))]
+
+    windows = []
+    seen = set()
+    half = window_size // 2
+    for idx in candidate_indices:
+        idx = int(idx)
+        start = idx - half
+        end = start + window_size
+        if start < 0:
+            start = 0
+            end = window_size
+        if end > num_items:
+            end = num_items
+            start = max(0, end - window_size)
+        window = tuple(range(start, end))
+        if window not in seen:
+            seen.add(window)
+            windows.append(list(window))
+    return windows
+
+
+class ConvMemoryReranker:
+    """Plug-in reranker over precomputed memory embeddings.
+
+    The class deliberately does not own text embedding. In production, users can
+    bring any retriever or embedding model, pass the raw top-k candidates here,
+    and get a ConvMemory-enhanced ordering back.
+    """
+
+    def __init__(self, conv_model, scorer, config=None, device="cpu"):
+        self.conv_model = conv_model
+        self.scorer = scorer
+        self.config = config or RerankConfig()
+        self.device = device
+
+    def make_item(
+        self,
+        query_embedding,
+        memory_embeddings,
+        memory_ids,
+        memory_texts: Optional[Iterable[str]] = None,
+        query: str = "",
+    ):
+        memory_ids = [str(x) for x in memory_ids]
+        memory_embeddings = normalize_rows(memory_embeddings)
+        query_embedding = np.asarray(query_embedding, dtype=np.float32)
+        query_embedding = query_embedding / (np.linalg.norm(query_embedding) + 1e-8)
+        if memory_texts is None:
+            memory_texts = ["" for _ in memory_ids]
+        memories = [
+            {"id": memory_id, "text": text}
+            for memory_id, text in zip(memory_ids, memory_texts)
+        ]
+        windows = sliding_windows(
+            len(memory_ids),
+            self.config.window_size,
+            self.config.stride,
+        )
+
+        return {
+            "question_id": "query",
+            "question_type": "unknown",
+            "query": query,
+            "query_embedding": query_embedding.astype(np.float32),
+            "memory_embeddings": memory_embeddings.astype(np.float32),
+            "memory_ids": memory_ids,
+            "memories": memories,
+            "windows": windows,
+            "window_tensor": None,
+            "gold_memory_ids": [],
+        }
+
+    def rerank_embeddings(
+        self,
+        query_embedding,
+        memory_embeddings,
+        memory_ids,
+        memory_texts: Optional[Iterable[str]] = None,
+        query: str = "",
+        candidate_indices=None,
+        candidate_top_n: Optional[int] = None,
+        raw_weight: Optional[float] = None,
+        window_mode: Optional[str] = None,
+    ):
+        item = self.make_item(
+            query_embedding=query_embedding,
+            memory_embeddings=memory_embeddings,
+            memory_ids=memory_ids,
+            memory_texts=memory_texts,
+            query=query,
+        )
+        return self.rerank_item(
+            item,
+            candidate_indices=candidate_indices,
+            candidate_top_n=candidate_top_n,
+            raw_weight=raw_weight,
+            window_mode=window_mode,
+        )
+
+    def rerank_item(
+        self,
+        item,
+        candidate_indices=None,
+        candidate_top_n: Optional[int] = None,
+        raw_weight: Optional[float] = None,
+        window_mode: Optional[str] = None,
+    ):
+        raw_scores = cosine_scores(item["query_embedding"], item["memory_embeddings"])
+        if candidate_indices is None:
+            top_n = candidate_top_n or self.config.candidate_top_n
+            candidate_indices = np.argsort(-raw_scores)[: min(top_n, len(raw_scores))]
+        else:
+            candidate_indices = np.asarray(candidate_indices, dtype=np.int64)
+
+        scoring_item = item
+        selected_window_mode = window_mode or self.config.window_mode
+        if selected_window_mode == "candidate_local":
+            local_windows = candidate_local_windows(
+                len(item["memory_ids"]),
+                candidate_indices,
+                self.config.window_size,
+            )
+            scoring_item = {
+                **item,
+                "windows": local_windows,
+                "window_tensor": window_tensor(item["memory_embeddings"], local_windows),
+            }
+        elif selected_window_mode != "full":
+            raise ValueError(f"Unknown window_mode: {selected_window_mode}")
+        elif scoring_item.get("window_tensor") is None:
+            scoring_item = {
+                **item,
+                "window_tensor": window_tensor(item["memory_embeddings"], item["windows"]),
+            }
+
+        with torch.no_grad():
+            conv_tensor = window_scores(self.conv_model, scoring_item, self.device)
+        memory_to_windows = build_memory_to_windows(scoring_item["windows"])
+        _, _, ce_lite_scores = score_ce_lite(
+            self.conv_model,
+            self.scorer,
+            scoring_item,
+            candidate_indices,
+            self.device,
+            raw_scores_all=raw_scores,
+            window_logits=conv_tensor,
+            memory_to_windows=memory_to_windows,
+            dca_router_block_size=self.config.dca_router_block_size,
+            lexical_features=self.config.lexical_features,
+        )
+        ranked_ids = rerank_candidates(
+            raw_scores,
+            candidate_indices,
+            ce_lite_scores,
+            item["memory_ids"],
+            raw_weight=self.config.raw_weight if raw_weight is None else raw_weight,
+        )
+        score_by_id = {
+            item["memory_ids"][int(idx)]: float(score)
+            for idx, score in zip(candidate_indices, ce_lite_scores)
+        }
+        raw_by_id = {
+            item["memory_ids"][idx]: float(score)
+            for idx, score in enumerate(raw_scores)
+        }
+        text_by_id = {
+            str(memory.get("id", idx)): memory.get("text")
+            for idx, memory in enumerate(item.get("memories", []))
+        }
+        return [
+            RerankResult(
+                memory_id=memory_id,
+                score=score_by_id.get(memory_id, raw_by_id[memory_id]),
+                raw_score=raw_by_id[memory_id],
+                rank=rank,
+                text=text_by_id.get(memory_id),
+            )
+            for rank, memory_id in enumerate(ranked_ids, start=1)
+        ]

convmemory/routing.py

@@ -0,0 +1,208 @@
+from dataclasses import dataclass
+from typing import Iterable, List, Mapping, Optional, Sequence
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class CompressionRouteConfig:
+    """Configuration for note-to-memory candidate routing.
+
+    The defaults are intentionally conservative and reflect the stable
+    LoCoMo v0.31 setting: use compressed notes to select a smaller raw-memory
+    pool, then let ConvMemory rerank that pool.
+    """
+
+    note_depth: int = 240
+    max_sources_per_note: int = 5
+    max_candidates: int = 450
+    raw_anchor: int = 80
+
+
+@dataclass(frozen=True)
+class CompressionRouteResult:
+    candidate_indices: List[int]
+    candidate_ids: List[str]
+    note_indices: List[int]
+    raw_anchor_count: int
+
+
+@dataclass(frozen=True)
+class CompressedNoteConfig:
+    """Configuration for lightweight raw-memory note construction."""
+
+    mode: str = "session"
+    block_size: int = 32
+    representatives: int = 3
+    strategy: str = "central"
+    session_key: str = "session_id"
+
+
+class CompressionRouter:
+    """Route raw memories through compressed note blocks.
+
+    Compressed memories are dictionaries with at least:
+    - `text`: note text used for embedding
+    - `source_ids`: raw memory ids covered by this note
+
+    The router does not call ConvMemory directly. It only returns candidate ids
+    so it can be plugged into any retrieval or agent-memory pipeline.
+    """
+
+    def __init__(self, config: Optional[CompressionRouteConfig] = None):
+        self.config = config or CompressionRouteConfig()
+
+    def route(
+        self,
+        query_embedding,
+        memory_embeddings,
+        memory_ids: Sequence[str],
+        compressed_embeddings,
+        compressed_memories: Iterable[Mapping],
+    ) -> CompressionRouteResult:
+        memory_ids = [str(memory_id) for memory_id in memory_ids]
+        compressed_memories = list(compressed_memories)
+        if len(compressed_memories) != len(compressed_embeddings):
+            raise ValueError("compressed_memories and compressed_embeddings must have the same length")
+
+        query = _normalize_vector(query_embedding)
+        memories = _normalize_matrix(memory_embeddings)
+        notes = _normalize_matrix(compressed_embeddings)
+
+        raw_scores = memories @ query
+        raw_order = np.argsort(-raw_scores)
+        note_scores = notes @ query if len(notes) else np.asarray([], dtype=np.float32)
+        note_order = np.argsort(-note_scores)[: max(0, int(self.config.note_depth))]
+
+        id_to_index = {memory_id: idx for idx, memory_id in enumerate(memory_ids)}
+        selected: List[int] = []
+        seen = set()
+
+        for idx in raw_order[: max(0, int(self.config.raw_anchor))]:
+            self._add_candidate(int(idx), selected, seen)
+            if len(selected) >= self.config.max_candidates:
+                return self._result(selected, memory_ids, note_order)
+
+        for note_idx in note_order:
+            source_indices = []
+            for source_id in compressed_memories[int(note_idx)].get("source_ids", []):
+                source_key = str(source_id)
+                if source_key in id_to_index:
+                    source_indices.append(id_to_index[source_key])
+            source_indices.sort(key=lambda idx: -float(raw_scores[idx]))
+            limit = int(self.config.max_sources_per_note)
+            if limit > 0:
+                source_indices = source_indices[:limit]
+
+            for idx in source_indices:
+                self._add_candidate(int(idx), selected, seen)
+                if len(selected) >= self.config.max_candidates:
+                    return self._result(selected, memory_ids, note_order)
+
+        return self._result(selected, memory_ids, note_order)
+
+    @staticmethod
+    def _add_candidate(idx: int, selected: List[int], seen) -> None:
+        if idx in seen:
+            return
+        selected.append(idx)
+        seen.add(idx)
+
+    def _result(self, selected: List[int], memory_ids: Sequence[str], note_order) -> CompressionRouteResult:
+        return CompressionRouteResult(
+            candidate_indices=list(selected),
+            candidate_ids=[memory_ids[idx] for idx in selected],
+            note_indices=[int(idx) for idx in note_order],
+            raw_anchor_count=min(len(selected), max(0, int(self.config.raw_anchor))),
+        )
+
+
+def _normalize_vector(x):
+    arr = np.asarray(x, dtype=np.float32)
+    return arr / (np.linalg.norm(arr) + 1e-8)
+
+
+def _normalize_matrix(x):
+    arr = np.asarray(x, dtype=np.float32)
+    if arr.ndim == 1:
+        arr = arr.reshape(1, -1)
+    return arr / (np.linalg.norm(arr, axis=1, keepdims=True) + 1e-8)
+
+
+def build_compressed_notes(
+    memories: Iterable[Mapping],
+    memory_embeddings,
+    config: Optional[CompressedNoteConfig] = None,
+):
+    """Build compressed notes from raw memories.
+
+    This helper is intentionally simple: it groups an ordered memory stream by
+    session or fixed-size blocks, chooses representative turns, and keeps
+    `source_ids` so the note can be expanded back to raw memories.
+    """
+
+    cfg = config or CompressedNoteConfig()
+    memories = list(memories)
+    embeddings = _normalize_matrix(memory_embeddings)
+    if len(memories) != len(embeddings):
+        raise ValueError("memories and memory_embeddings must have the same length")
+    if cfg.mode not in {"session", "block"}:
+        raise ValueError("CompressedNoteConfig.mode must be 'session' or 'block'")
+    if cfg.strategy not in {"central", "first"}:
+        raise ValueError("CompressedNoteConfig.strategy must be 'central' or 'first'")
+
+    groups = _session_groups(memories, cfg.session_key)
+    if cfg.mode == "block":
+        groups = [
+            group[start : start + cfg.block_size]
+            for group in groups
+            for start in range(0, len(group), cfg.block_size)
+        ]
+
+    notes = []
+    for note_idx, group in enumerate(groups):
+        if not group:
+            continue
+        reps = _representative_indices(group, embeddings, cfg.strategy, cfg.representatives)
+        text = " ".join(str(memories[idx].get("text", "")) for idx in reps)
+        source_ids = [str(memories[idx].get("id", idx)) for idx in group]
+        session_id = str(memories[group[0]].get(cfg.session_key, ""))
+        notes.append(
+            {
+                "id": f"{cfg.mode}:{note_idx}",
+                "text": text,
+                "source_ids": source_ids,
+                cfg.session_key: session_id,
+                "granularity": cfg.mode,
+            }
+        )
+    return notes
+
+
+def _session_groups(memories, session_key):
+    groups = []
+    current_value = None
+    current = []
+    for idx, item in enumerate(memories):
+        value = item.get(session_key, "")
+        if current and value != current_value:
+            groups.append(current)
+            current = []
+        current_value = value
+        current.append(idx)
+    if current:
+        groups.append(current)
+    return groups
+
+
+def _representative_indices(group, embeddings, strategy, representatives):
+    count = max(1, min(int(representatives), len(group)))
+    if strategy == "first":
+        return group[:count]
+
+    local_embeddings = embeddings[group]
+    centroid = local_embeddings.mean(axis=0)
+    centroid = centroid / (np.linalg.norm(centroid) + 1e-8)
+    scores = local_embeddings @ centroid
+    picked = np.argsort(-scores)[:count]
+    return [group[int(i)] for i in sorted(picked)]