Switch sampled decoder to nucleus/top-p (Carbon's decoder) instead of top-k — escapes low-complexity loops

daisychain.py

@@ -103,7 +103,7 @@ class DaisyChain:
 
     @torch.no_grad()
     def generate_stream(self, domain, length=180, temperature=1.0, top_k=40,
-                        repetition_penalty=1.3, prompt="", greedy=False):
+                        repetition_penalty=1.3, prompt="", greedy=False, top_p=0.9):
         """Yield the growing continuation base-by-base (for live streaming UIs).
         greedy=True takes the argmax 6-mer each step (deterministic — the model's single
         best guess, same decoding as the sequence-recovery metric; the coding domains
@@ -127,20 +127,26 @@ class DaisyChain:
             # self-reinforcing loop (argmax keeps re-picking the same 6-mer); the penalty plus
             # a hard block on the last few emitted tokens forces it onward to its next-best,
             # non-repeating guess — which is what actually de-degenerates the output.
-            if repetition_penalty and repetition_penalty != 1.0:
-                for tid in set(emitted[-12:]):
-                    v = logits[0, tid]
-                    logits[0, tid] = v / repetition_penalty if v > 0 else v * repetition_penalty
             if greedy:
-                for tid in set(emitted[-8:]):                      # block longer low-complexity cycles
+                # argmax with repeat penalty + a hard block on recent tokens (greedy alone loops)
+                if repetition_penalty and repetition_penalty != 1.0:
+                    for tid in set(emitted[-12:]):
+                        v = logits[0, tid]
+                        logits[0, tid] = v / repetition_penalty if v > 0 else v * repetition_penalty
+                for tid in set(emitted[-8:]):
                     logits[0, tid] = -1e9
                 ti = int(logits.argmax())
             else:
+                # nucleus (top-p) sampling — the same decoder Carbon uses. Adapts the candidate
+                # set to the distribution (keeps only tokens carrying mass), which escapes the
+                # low-complexity GC/AT loops that a fixed top-k falls into. No repetition penalty.
                 logits = logits / max(temperature, 1e-6)
-                if top_k > 0:
-                    v, _ = torch.topk(logits, top_k)
-                    logits[logits < v[:, [-1]]] = -1e9
-                ti = int(torch.multinomial(F.softmax(logits, dim=-1), 1))
+                sl, si = torch.sort(logits[0], descending=True)
+                cum = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
+                rm = cum > top_p
+                rm[1:] = rm[:-1].clone(); rm[0] = False
+                logits[0, si[rm]] = -1e9
+                ti = int(torch.multinomial(F.softmax(logits[0], dim=-1), 1))
             emitted.append(ti)
             t = torch.cat([t, torch.tensor([[ti]], device=self.dev)], dim=1)
             bases.append(self.tok._ids_to_tokens[ti])