File size: 8,703 Bytes
1f71c7d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
"""PALIMPSESTE — Creative generation via mixture logits.

The core problem: _logits() collapses K retrieved candidates into ONE HV via
majority sign, then takes argmax. This is structurally incapable of producing
novel tokens — it can only echo the majority.

The solution: Mixture Logits. Instead of one bundle, we use ALL K retrieved
candidates as weighted votes over the vocabulary:

    mixture_logits[t] = sum_k( sim_k * similarity(token_hv[t], value_hv_k) )

A token that's close to MULTIPLE candidates gets boosted beyond what any single
candidate would predict. This is how novel compositions emerge — the model
finds tokens at the intersection of multiple knowledge fragments.

Additionally: anti-repetition via n-gram blocking prevents attractor loops.
"""

from __future__ import annotations

import numpy as np
from dataclasses import dataclass
from typing import Optional

from .hv import HV, bind, similarity, bits_to_signs

__all__ = ["CreativeGenerator", "CreativeResult"]


@dataclass
class CreativeResult:
    """Result of creative generation."""
    text: str
    n_candidates_per_step: list[int]
    novel_tokens: int
    temperature: float
    mixture_diversity: float


class CreativeGenerator:
    """Generate text using mixture logits from multiple retrieved candidates.

    Instead of retrieving one bundled HV, this generator retrieves K matches
    and builds a mixture distribution over the vocabulary. Tokens at the
    intersection of multiple candidates get boosted — enabling novel
    compositions that were never explicitly stored.

    This is how PALIMPSESTE creates, not just repeats.
    """

    def __init__(self, lm, top_k: int = 5, diversity: float = 0.3) -> None:
        """
        Parameters
        ----------
        lm : PalimpsesteForCausalLM
            The language model.
        top_k : int
            Number of retrieved candidates to use per step.
        diversity : float
            How much to boost intersection tokens (0=none, 1=max).
        """
        self.lm = lm
        self.top_k = top_k
        self.diversity = diversity

    def mixture_logits(self, context_ids: list[int]) -> np.ndarray:
        """Compute mixture logits from K retrieved candidates.

        For each of the top-K retrieved traces, compute similarity between
        the trace's value HV and every token HV. Weight by retrieval sim.
        Sum to get mixture logits.

        This produces a RICHER distribution than single-HV retrieval because
        tokens at the intersection of candidates get extra boost.
        """
        assert self.lm.tokenizer is not None

        s = self.lm._state_hv(context_ids)
        q = bind(self.lm._self_hv, s)

        # Retrieve K matches (NOT the bundled result)
        ret = self.lm.phi.retrieve(self.lm.mem, q)

        V = getattr(self.lm.tokenizer, 'vocab_size_actual', None) or self.lm.config.vocab_size
        logits = np.zeros(V, dtype=np.float64)

        if not ret.matches:
            return logits

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

        if not top_matches:
            return logits

        # Ensure token bits cache is built
        packed_len = len(top_matches[0][0].value.bits)
        if (self.lm._token_bits_cache is None or
                self.lm._token_bits_cache.shape[0] != V or
                self.lm._token_bits_cache.shape[1] != packed_len):
            self.lm._token_bits_cache = np.empty((V, packed_len), dtype=np.uint8)
            for tid in range(V):
                self.lm._token_bits_cache[tid] = self.lm.tokenizer.token_hv(tid).bits
        tok_bits = self.lm._token_bits_cache

        _POPCOUNT = np.array([bin(i).count('1') for i in range(256)], dtype=np.uint16)
        D = self.lm.config.D

        # For each candidate, compute token logits and add weighted
        for match, sim in top_matches:
            value_bits = match.value.bits
            xored = np.bitwise_xor(tok_bits, value_bits[np.newaxis, :])
            hamming_dists = _POPCOUNT[xored].sum(axis=1)
            candidate_logits = 1.0 - 2.0 * hamming_dists / D

            # Weight by retrieval similarity (map [-1,1] → [0,1])
            weight = (sim + 1.0) / 2.0

            # Add weighted logits
            logits += weight * candidate_logits

            # Diversity boost: tokens that are above threshold in this candidate
            # AND already high in logits get extra boost (intersection effect)
            if self.diversity > 0:
                high_mask = candidate_logits > 0.0
                intersection = logits > 0.0
                boost_mask = high_mask & intersection
                logits[boost_mask] += self.diversity * 0.1

        # Normalize
        if logits.max() > logits.min():
            logits = (logits - logits.min()) / (logits.max() - logits.min())

        return logits

    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 100,
        temperature: float = 0.0,
        ngram_block: int = 3,
        seed: int | None = None,
    ) -> CreativeResult:
        """Generate text creatively using mixture logits.

        Parameters
        ----------
        prompt : str
            Input text to continue.
        max_new_tokens : int
            Maximum tokens to generate.
        temperature : float
            0.0 = greedy on mixture, >0 = sample from mixture.
        ngram_block : int
            Block repetition of n-grams of this length (3 = no trigram repeats).
        seed : int | None
            RNG seed for reproducibility.
        """
        from .lm import BOS, EOS, PAD
        assert self.lm.tokenizer is not None

        tok = self.lm.tokenizer
        rng = np.random.default_rng(seed) if seed is not None else np.random.default_rng()

        ids = tok.encode(prompt, add_bos=True, add_eos=True) + [BOS]
        generated: list[int] = []
        candidates_per_step: list[int] = []
        novel = 0

        for step in range(max_new_tokens):
            logits = self.mixture_logits(ids)

            if logits.max() == 0.0 and logits.min() == 0.0:
                break  # cold start

            # N-gram blocking: penalize tokens that would complete a recent n-gram
            if ngram_block > 0 and len(generated) >= ngram_block - 1:
                recent_ngram = tuple(generated[-(ngram_block - 1):]) if ngram_block > 1 else ()
                # Check if this n-gram appeared before
                for i in range(len(generated) - ngram_block + 1):
                    ng = tuple(generated[i:i + ngram_block - 1]) if ngram_block > 1 else ()
                    if ng == recent_ngram and i + ngram_block - 1 < len(generated):
                        blocked_token = generated[i + ngram_block - 1]
                        logits[blocked_token] *= 0.3  # penalize but don't forbid

            # Select next token
            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(rng.choice(len(probs), p=probs))

            # Track if this token is "novel" (different from what standard
            # retrieval would predict)
            standard_logits = self.lm._logits(ids)
            standard_next = int(np.argmax(standard_logits))
            if next_id != standard_next:
                novel += 1

            generated.append(next_id)
            ids.append(next_id)
            candidates_per_step.append(self.top_k)

            if next_id == EOS:
                break

        text = tok.decode(generated)
        avg_diversity = novel / max(len(generated), 1)

        return CreativeResult(
            text=text,
            n_candidates_per_step=candidates_per_step,
            novel_tokens=novel,
            temperature=temperature,
            mixture_diversity=avg_diversity,
        )

    def respond(self, question: str, max_new_tokens: int = 100,
                temperature: float = 0.0) -> str:
        """Respond to a question using creative generation."""
        assert self.lm.tokenizer is not None
        q_ids = self.lm.tokenizer.encode(question, add_bos=True, add_eos=True)
        ctx = q_ids + [1]  # BOS
        result = self.generate(question, max_new_tokens=max_new_tokens,
                               temperature=temperature)
        return result.text