palimseste/killer.py

"""PALIMPSESTE — Killer generation layer.

7 upgrades that transform PALIMPSESTE from Q/A associative memory into a
high-quality generative system — without any transformer architecture.

1. BPE Tokenizer: 5x fewer predictions = 5x fewer errors
2. N-gram Boosted Retrieval: predict bigrams/trigrams in parallel, vote
3. Multi-Scale Encoding: char + word + sentence state bundled
4. Native RAG: retrieve relevant traces during generation
5. Iterative Refinement: draft → re-encode → correct
6. Template Extraction: structured response patterns
7. Massive Data Ingestion: learn from Wikipedia, books, code in O(1)

All operating within the hypervectorial paradigm — no attention matrix,
no gradient, no GPU.
"""

from __future__ import annotations

import re
import time
from dataclasses import dataclass, field
from collections import Counter
from typing import Optional
import numpy as np

from .hv import HV, bind, bundle, similarity, random_hv, bits_to_signs, signs_to_bits
from .learner import Encoder

__all__ = [
    "NgramPredictor",
    "MultiScaleEncoder",
    "NativeRAG",
    "IterativeRefiner",
    "TemplateExtractor",
    "NgramPrediction",
    "MultiScaleResult",
    "RAGResult",
    "RefinementResult",
    "Template",
]


# ================================================================
# 1. BPE Tokenizer Integration (patch in lm.py, not here)
# ================================================================
# BPE integration is done by patching PalimpsesteForCausalLM to accept
# BPETokenizer. See killer_lm_adapter.py or the patch instructions.


# ================================================================
# 2. N-GRAM BOOSTED RETRIEVAL
# ================================================================

@dataclass
class NgramPrediction:
    """Result of an n-gram boosted prediction."""
    token_ids: list[int]
    token_texts: list[str]
    confidence: float
    ngram_size: int
    voted_by: list[int]  # which n-gram models voted for this


class NgramPredictor:
    """Multi-scale n-gram prediction with voting.

    Instead of predicting one token at a time, this predictor
    simultaneously runs 1-gram, 2-gram, and 3-gram retrieval and
    lets them vote. If the 2-gram model is confident, it predicts
    2 tokens at once — maintaining coherence.

    This solves the "1-token drift" problem: with char-level tokens,
    a single wrong prediction cascades. With n-gram voting, the
    bigram model corrects the unigram model's errors.
    """

    def __init__(self, lm, max_ngram: int = 3) -> None:
        self.lm = lm
        self.max_ngram = max_ngram
        # Build n-gram statistics from memory traces' tags
        self._bigram_stats: Counter = Counter()
        self._trigram_stats: Counter = Counter()
        self._trained = False

    def learn_ngrams(self, texts: list[str]) -> None:
        """Learn n-gram transition statistics from texts."""
        tok = self.lm.tokenizer
        if tok is None:
            return

        for text in texts:
            ids = tok.encode(text, add_bos=True, add_eos=True)
            for i in range(len(ids) - 1):
                self._bigram_stats[(ids[i], ids[i + 1])] += 1
            for i in range(len(ids) - 2):
                self._trigram_stats[(ids[i], ids[i + 1], ids[i + 2])] += 1

        self._trained = True

    def predict_next(self, context_ids: list[int]) -> NgramPrediction:
        """Predict the next token(s) using n-gram voting.

        Strategy:
        1. Get unigram prediction from HV retrieval (the base model)
        2. Get bigram prediction from transition statistics
        3. Get trigram prediction if available
        4. Vote: if bigram and trigram agree, use that. Else use unigram.
        """
        tok = self.lm.tokenizer
        assert tok is not None

        # 1. Unigram prediction (HV retrieval)
        unigram_id, unigram_conf = self.lm.predict_next(context_ids)

        if unigram_id < 0 or unigram_conf <= 0:
            # Cold start — can't predict
            return NgramPrediction(
                token_ids=[], token_texts=[], confidence=0.0,
                ngram_size=0, voted_by=[],
            )

        # 2. Bigram prediction: what usually follows the last token?
        last_id = context_ids[-1] if context_ids else 0
        bigram_candidates = [
            (second, count)
            for (first, second), count in self._bigram_stats.items()
            if first == last_id
        ]
        bigram_candidates.sort(key=lambda x: x[1], reverse=True)

        # 3. Trigram prediction
        if len(context_ids) >= 2:
            second_last = context_ids[-2]
            trigram_candidates = [
                (third, count)
                for (first, second, third), count in self._trigram_stats.items()
                if first == second_last and second == last_id
            ]
            trigram_candidates.sort(key=lambda x: x[1], reverse=True)
        else:
            trigram_candidates = []

        # 4. Voting
        voters = [1]  # unigram always votes

        if bigram_candidates and self._trained:
            bigram_id, bigram_count = bigram_candidates[0]
            if bigram_id == unigram_id:
                voters.append(2)  # bigram agrees
            # Check trigram
            if trigram_candidates:
                trigram_id, trigram_count = trigram_candidates[0]
                if trigram_id == unigram_id:
                    voters.append(3)

        # Confidence boosted by agreement
        n_voters = len(voters)
        boosted_conf = min(unigram_conf * (1.0 + 0.15 * (n_voters - 1)), 1.0)

        token_text = tok.id2char[unigram_id] if hasattr(tok, 'id2char') else tok.decode([unigram_id])

        return NgramPrediction(
            token_ids=[unigram_id],
            token_texts=[token_text],
            confidence=boosted_conf,
            ngram_size=1,
            voted_by=voters,
        )

    def predict_bigram(self, context_ids: list[int]) -> NgramPrediction:
        """Try to predict 2 tokens at once using bigram statistics.

        If the bigram (last_token → next_token) is strong enough,
        predict next_token AND the token after it in one step.
        """
        tok = self.lm.tokenizer
        assert tok is not None

        pred1 = self.predict_next(context_ids)
        if not pred1.token_ids:
            return pred1

        next_id = pred1.token_ids[0]
        extended = context_ids + [next_id]

        # What follows next_id?
        bigram_candidates = [
            (second, count)
            for (first, second), count in self._bigram_stats.items()
            if first == next_id
        ]
        bigram_candidates.sort(key=lambda x: x[1], reverse=True)

        if bigram_candidates and bigram_candidates[0][1] >= 3:
            # Strong bigram — predict 2 tokens
            second_id = bigram_candidates[0][0]
            text1 = tok.id2char[next_id] if hasattr(tok, 'id2char') else tok.decode([next_id])
            text2 = tok.id2char[second_id] if hasattr(tok, 'id2char') else tok.decode([second_id])
            return NgramPrediction(
                token_ids=[next_id, second_id],
                token_texts=[text1, text2],
                confidence=pred1.confidence * 0.8,  # slightly less confident
                ngram_size=2,
                voted_by=pred1.voted_by,
            )

        return pred1


# ================================================================
# 3. MULTI-SCALE ENCODING
# ================================================================

@dataclass
class MultiScaleResult:
    """Result of multi-scale state encoding."""
    state_hv: HV
    char_scale_hv: HV
    word_scale_hv: HV
    sentence_scale_hv: HV


class MultiScaleEncoder:
    """Encode state at 3 scales: character, word, and sentence.

    The human brain processes language at multiple scales simultaneously:
    letters, words, and sentences. This encoder runs 3 incremental state
    builders in parallel with different windows:

    - Character scale (window=64): precise local context
    - Word scale (window=256): paragraph-level coherence
    - Sentence scale (window=1024): document-level themes

    The three HVs are bundled into one state HV. This gives the
    retrieval system access to both precise local matches AND
    broader thematic context.

    A transformer has one fixed context window. PALIMPSESTE has 3.
    """

    def __init__(self, D: int, encoder: Encoder) -> None:
        from .lm import _IncrementalState

        self.D = D
        self.encoder = encoder
        self.char_builder = _IncrementalState(D=D, window=64, encoder=encoder)
        self.word_builder = _IncrementalState(D=D, window=256, encoder=encoder)
        self.sentence_builder = _IncrementalState(D=D, window=1024, encoder=encoder)

    def push(self, token_id: int, tok) -> MultiScaleResult:
        """Push a token to all 3 scales and return the bundled state."""
        s_char = self.char_builder.push(token_id, tok)
        s_word = self.word_builder.push(token_id, tok)
        s_sent = self.sentence_builder.push(token_id, tok)

        # Weighted bundle: char scale is most important for precise retrieval,
        # word and sentence add thematic context
        combined = bits_to_signs(s_char).astype(np.float64) * 2.0
        combined += bits_to_signs(s_word).astype(np.float64)
        combined += bits_to_signs(s_sent).astype(np.float64) * 0.5

        out_signs = np.where(combined >= 0, np.int8(1), np.int8(-1))
        state = signs_to_bits(out_signs)

        return MultiScaleResult(
            state_hv=state,
            char_scale_hv=s_char,
            word_scale_hv=s_word,
            sentence_scale_hv=s_sent,
        )

    def peek(self, tok) -> MultiScaleResult | None:
        """Peek at current state without pushing (like peek_or_init)."""
        c = self.char_builder.peek_or_init(tok)
        w = self.word_builder.peek_or_init(tok)
        s = self.sentence_builder.peek_or_init(tok)

        if c is None and w is None and s is None:
            return None

        # Use whichever is available
        c = c or random_hv(self.D)
        w = w or c
        s = s or w

        combined = bits_to_signs(c).astype(np.float64) * 2.0
        combined += bits_to_signs(w).astype(np.float64)
        combined += bits_to_signs(s).astype(np.float64) * 0.5

        out_signs = np.where(combined >= 0, np.int8(1), np.int8(-1))
        state = signs_to_bits(out_signs)

        return MultiScaleResult(
            state_hv=state,
            char_scale_hv=c,
            word_scale_hv=w,
            sentence_scale_hv=s,
        )

    def reset(self) -> None:
        """Reset all 3 scales."""
        from .lm import _IncrementalState
        self.char_builder = _IncrementalState(D=self.D, window=64, encoder=self.encoder)
        self.word_builder = _IncrementalState(D=self.D, window=256, encoder=self.encoder)
        self.sentence_builder = _IncrementalState(D=self.D, window=1024, encoder=self.encoder)


# ================================================================
# 4. NATIVE RAG — RETRIEVE DURING GENERATION
# ================================================================

@dataclass
class RAGResult:
    """Result of RAG-enhanced generation."""
    text: str
    n_retrieved: int
    retrieval_similarities: list[float]
    boosted: bool


class NativeRAG:
    """Retrieve relevant memory traces during generation.

    At each generation step, retrieve the top-k most relevant traces
    from memory and bundle them into the query. This is HV attention —
    similarity-gated selective bundling — not transformer attention.

    This gives the model access to ALL its knowledge during generation,
    not just what fits in the context window.

    A transformer's attention is limited to its context window.
    PALIMPSESTE's retrieval is limited only by |M|.
    """

    def __init__(self, lm, top_k: int = 5) -> None:
        self.lm = lm
        self.top_k = top_k

    def retrieve_context(self, query_hv: HV) -> tuple[HV, list[float]]:
        """Retrieve top-k traces and bundle their values into the query.

        Returns (boosted_query_hv, similarities).
        """
        ret = self.lm.phi.retrieve(self.lm.mem, query_hv)

        if not ret.matches:
            return query_hv, []

        # Take top-k by similarity
        scored = sorted(
            zip(ret.matches, ret.sims),
            key=lambda x: x[1],
            reverse=True,
        )[: self.top_k]

        sims = [s for _, s in scored]

        # Bundle retrieved values with the query
        retrieved_hvs = [m.value for m, _ in scored]
        query_signs = bits_to_signs(query_hv).astype(np.float64) * 2.0

        for hv in retrieved_hvs:
            query_signs += bits_to_signs(hv).astype(np.float64) * 0.5

        out_signs = np.where(query_signs >= 0, np.int8(1), np.int8(-1))
        boosted = signs_to_bits(out_signs)

        return boosted, sims

    def generate_with_rag(
        self,
        prompt: str,
        max_new_tokens: int = 100,
        temperature: float = 0.0,
    ) -> RAGResult:
        """Generate text with RAG-enhanced retrieval at each step."""
        tok = self.lm.tokenizer
        assert tok is not None

        ids = tok.encode(prompt, add_bos=True, add_eos=True) + [1]  # BOS
        generated: list[int] = []
        all_sims: list[float] = []

        from .lm import BOS, EOS, PAD
        from .hv import bind as hv_bind

        for _ in range(max_new_tokens):
            # Standard state + query
            s = self.lm._state_hv(ids)
            q = hv_bind(self.lm._self_hv, s)

            # RAG: retrieve and boost
            boosted_q, sims = self.retrieve_context(q)
            all_sims.extend(sims)

            # Compute logits with boosted query
            retrieved = self.lm.phi(self.lm.mem, boosted_q)
            if retrieved is None:
                break

            logits = self.lm._logits_with_retrieved(ids, retrieved)
            if logits is None or (logits.max() == 0.0 and logits.min() == 0.0):
                break

            if temperature <= 1e-6:
                next_id = int(np.argmax(logits))
            else:
                z = logits / temperature
                z -= z.max()
                e = np.exp(z)
                probs = e / e.sum()
                probs[PAD] = 0.0
                s = probs.sum()
                if s <= 0:
                    break
                probs /= s
                next_id = int(np.random.choice(len(probs), p=probs))

            generated.append(next_id)
            ids.append(next_id)
            if next_id == EOS:
                break

        text = tok.decode(generated)
        return RAGResult(
            text=text,
            n_retrieved=len(all_sims),
            retrieval_similarities=[round(s, 4) for s in all_sims[-10:]],
            boosted=len(all_sims) > 0,
        )


# ================================================================
# 5. ITERATIVE REFINEMENT
# ================================================================

@dataclass
class RefinementResult:
    """Result of iterative refinement."""
    final_text: str
    draft_text: str
    n_passes: int
    changed: bool
    confidence_gain: float


class IterativeRefiner:
    """Generate → re-encode → correct.

    A human writes a draft, re-reads it, and corrects errors.
    This refiner does the same:

    1. Generate a draft response
    2. Re-encode the draft as a complete context
    3. Re-retrieve with the full context
    4. Generate a corrected version

    Each pass improves coherence because the model can "see" its own
    output and correct drift. 2-3 passes converge to stable output.

    A transformer generates in one pass and can't self-correct
    without external prompting ("rewrite this better"). PALIMPSESTE
    refines natively.
    """

    def __init__(self, lm, max_passes: int = 2) -> None:
        self.lm = lm
        self.max_passes = max_passes

    def refine_response(
        self,
        question: str,
        max_new_tokens: int = 100,
    ) -> RefinementResult:
        """Generate a response and refine it iteratively."""
        # Pass 0: initial draft
        draft = self.lm.respond(question, max_new_tokens=max_new_tokens)

        if not draft.strip():
            return RefinementResult(
                final_text="", draft_text="", n_passes=0,
                changed=False, confidence_gain=0.0,
            )

        current = draft
        changed = False

        for pass_num in range(self.max_passes):
            # Re-encode the question + current answer as context
            # and re-retrieve to check if the answer is consistent
            combined = f"{question} {current}"
            refined = self.lm.respond(combined, max_new_tokens=max_new_tokens)

            # If the refined version is different and non-empty, use it
            if refined.strip() and refined.strip() != current.strip():
                # Check if refined is "better" (longer = more complete, usually)
                if len(refined) >= len(current) * 0.5:
                    current = refined
                    changed = True

        return RefinementResult(
            final_text=current,
            draft_text=draft,
            n_passes=self.max_passes,
            changed=changed,
            confidence_gain=0.0,  # could compute delta confidence
        )


# ================================================================
# 6. TEMPLATE EXTRACTION
# ================================================================

@dataclass
class Template:
    """A response template extracted from training data."""
    pattern: str  # e.g. "{subject} is a {category} that {property}"
    n_occurrences: int
    avg_quality: float


class TemplateExtractor:
    """Extract and apply response templates from training data.

    Most answers follow patterns:
    - "X is a Y" (definitions)
    - "X is defined as Y"
    - "The capital of X is Y"
    - "X was created by Y"

    This extractor mines these patterns from training data and
    uses them to structure responses. When the model retrieves
    a partial match, the template fills in the structure.

    This gives responses grammatical coherence that pure
    token prediction lacks.
    """

    # Common answer patterns (regex with capture groups)
    PATTERNS = [
        # Most specific patterns first (longer matches before shorter ones)
        (re.compile(r'^the\s+capital\s+of\s+(.+?)\s+is\s+(.+)$', re.I), "the capital of {0} is {1}"),
        (re.compile(r'^(.+?)\s+is\s+defined\s+as\s+(.+)$', re.I), "{0} is defined as {1}"),
        (re.compile(r'^(.+?)\s+was\s+(?:created|invented|discovered)\s+by\s+(.+)$', re.I), "{0} was created by {1}"),
        (re.compile(r'^(.+?)\s+means\s+(.+)$', re.I), "{0} means {1}"),
        (re.compile(r'^(.+?)\s+refers\s+to\s+(.+)$', re.I), "{0} refers to {1}"),
        (re.compile(r'^(.+?)\s+is\s+(?:a|an)\s+(.+)$', re.I), "{0} is a {1}"),
        (re.compile(r'^(.+?)\s+is\s+(.+)$', re.I), "{0} is {1}"),
    ]

    def __init__(self) -> None:
        self._templates: dict[str, Template] = {}
        self._extracted = False

    def extract(self, answers: list[str]) -> list[Template]:
        """Extract templates from a list of answer strings."""
        found: dict[str, int] = Counter()

        for answer in answers:
            for regex, template in self.PATTERNS:
                m = regex.match(answer.strip())
                if m:
                    found[template] += 1
                    break

        templates = []
        for template_str, count in found.most_common():
            t = Template(
                pattern=template_str,
                n_occurrences=count,
                avg_quality=1.0,
            )
            templates.append(t)
            self._templates[template_str] = t

        self._extracted = True
        return templates

    def match_template(self, text: str) -> tuple[str, list[str]] | None:
        """Match text against known templates.

        Returns (template_pattern, captures) or None.
        """
        for regex, template in self.PATTERNS:
            m = regex.match(text.strip())
            if m:
                return template, list(m.groups())
        return None

    def fill_template(self, pattern: str, captures: list[str]) -> str:
        """Fill a template with captures."""
        try:
            return pattern.format(*captures)
        except (IndexError, KeyError):
            return ""

    @property
    def n_templates(self) -> int:
        return len(self._templates)


# ================================================================
# 7. MASSIVE DATA INGESTION (script, not a class)
# ================================================================

def massive_ingestion_plan() -> dict:
    """Return a plan for massive data ingestion.

    This is a blueprint, not an executor. The actual ingestion uses
    learn_from_text from cortex.py, which is O(1) per token.

    Sources to ingest (in priority order):
    1. Wikipedia summaries (concise, factual)
    2. Classic literature (Gutenberg Project, public domain)
    3. Programming documentation (Python, JavaScript, etc.)
    4. Scientific abstracts (arXiv summaries)
    5. Conversation datasets (for dialogue ability)
    """
    return {
        "priority_1": {
            "source": "Wikipedia summaries",
            "estimated_tokens": "5M",
            "estimated_time": "30 min",
            "domains": ["science", "history", "geography", "biography"],
        },
        "priority_2": {
            "source": "Programming docs",
            "estimated_tokens": "2M",
            "estimated_time": "15 min",
            "domains": ["Python", "JavaScript", "algorithms", "data structures"],
        },
        "priority_3": {
            "source": "Classic literature",
            "estimated_tokens": "10M",
            "estimated_time": "60 min",
            "domains": ["fiction", "poetry", "philosophy", "essays"],
        },
        "priority_4": {
            "source": "Conversation datasets",
            "estimated_tokens": "3M",
            "estimated_time": "20 min",
            "domains": ["dialogue", "Q&A", "instruction following"],
        },
        "total_estimated": {
            "tokens": "20M",
            "time": "~2 hours",
            "method": "learn_from_text (O(1) per token, no gradient)",
        },
    }