Match Carbon FNS BpLogitsProcessor decoding

daisychain.py

@@ -173,19 +173,28 @@ class DaisyChain:
         return out
 
     # ---- FNS base-pair-level decode + score (Carbon's factorized-nucleotide approach) ----
-    def _bp_marginals(self, logits):
-        """Marginalize a 6-mer logit vector into six 4-way per-position base distributions [6,4]."""
-        p = F.softmax(logits[self._kmer_ids], dim=-1)          # [4096] over the full 6-mers only
+    def _bp_marginals(self, logits, top_p=1.0):
+        """Marginalize a 6-mer logit vector into six 4-way per-position base distributions [6,4].
+        Matching Carbon's _BPLogitsProcessor: any top-p filtering happens at the 6-MER level
+        (before marginalizing), not per-base."""
+        kl = logits[self._kmer_ids]                            # [4096] 6-mer logits
+        if top_p < 1.0:                                        # nucleus on the 6-mers first
+            sk, si = torch.sort(kl, descending=True)
+            rm = torch.cumsum(F.softmax(sk, -1), -1) > top_p
+            rm[1:] = rm[:-1].clone(); rm[0] = False
+            kl = kl.clone(); kl[si[rm]] = float("-inf")
+        p = F.softmax(kl, dim=-1)
         bp = torch.zeros(6, 4, device=logits.device)
         for pos in range(6):
             bp[pos].scatter_add_(0, self._base_at[pos], p)
         return bp
 
     @torch.no_grad()
-    def generate_baselevel_stream(self, domain, length=180, temperature=1.0, top_p=0.9, prompt=""):
-        """Base-pair-level generation: at each step marginalize the 6-mer softmax to six 4-way
-        nucleotide distributions and sample each base with temperature/top-p (Carbon FNS-style).
-        Deciding each base over 4 options — not 6 bases at once over 4096 — avoids whole-6-mer loops."""
+    def generate_baselevel_stream(self, domain, length=180, temperature=1.0, top_p=0.9,
+                                  prompt="", greedy=False):
+        """Base-pair generation exactly as Carbon's FNS BpLogitsProcessor: apply temperature +
+        top-p at the 6-mer level, marginalize the filtered 6-mer softmax to six 4-way base
+        distributions, then pick each base by multinomial (sampling) or argmax (greedy)."""
         m = self.models[domain]
         p = clean(prompt) if prompt else ""
         p = p[-1020:]; p = p[len(p) % 6:]
@@ -194,16 +203,11 @@ class DaisyChain:
         out = ""
         while len(out) < length:
             logits = m(input_ids=t).logits[0, -1].float() / max(temperature, 1e-6)
-            bp = self._bp_marginals(logits)                    # [6,4]
-            six = []
-            for pos in range(6):
-                probs = bp[pos]
-                s, si = torch.sort(probs, descending=True)
-                rm = torch.cumsum(s, -1) > top_p
-                rm[1:] = rm[:-1].clone(); rm[0] = False
-                q = probs.clone(); q[si[rm]] = 0; q = q / q.sum()
-                six.append(self._idx2base[int(torch.multinomial(q, 1))])
-            six = "".join(six)
+            bp = self._bp_marginals(logits, top_p=(1.0 if greedy else top_p))   # [6,4]
+            if greedy:
+                six = "".join(self._idx2base[int(bp[pos].argmax())] for pos in range(6))
+            else:
+                six = "".join(self._idx2base[int(torch.multinomial(bp[pos], 1))] for pos in range(6))
             t = torch.cat([t, torch.tensor([[self._kmer_id_of[six]]], device=self.dev)], dim=1)
             out += six
             yield out[:length]
@@ -211,10 +215,15 @@ class DaisyChain:
     @torch.no_grad()
     def score(self, domain, seq):
         """Carbon `score_sequence` equivalent: mean per-base log-prob of the observed sequence under
-        the base-level (marginalized) distribution. Higher = more likely. bits/base = -score/ln2."""
+        the base-level (marginalized) distribution. Higher = more likely. bits/base = -score/ln2.
+        Right-pads to a multiple of 6 with 'A' (Carbon's convention) and scores only real bases."""
         m = self.models[domain]
-        s = clean(seq); s = s[len(s) % 6:]
-        if len(s) < 12:
+        s = clean(seq)
+        orig = len(s)
+        r = len(s) % 6
+        if r:
+            s = s + "A" * (6 - r)                              # right-pad like score_sequence
+        if orig < 12:
             return float("nan")
         ids = [self.bos] + self.tok.encode(s, add_special_tokens=False)
         t = torch.tensor([ids], device=self.dev)
@@ -222,8 +231,14 @@ class DaisyChain:
         b2i = {"A": 0, "T": 1, "C": 2, "G": 3}
         tot, n = 0.0, 0
         for i in range(len(ids) - 1):
+            if i * 6 >= orig:                                  # don't score the padding
+                break
             bp = self._bp_marginals(logits[i].float())         # [6,4] predicted bases of next 6-mer
             nxt = s[i * 6:(i + 1) * 6]
             for pos, ch in enumerate(nxt):
+                if i * 6 + pos >= orig:                          # padding base — don't score
+                    break
+                if ch not in b2i:                               # skip N / ambiguous bases
+                    continue
                 tot += math.log(max(float(bp[pos, b2i[ch]]), 1e-12)); n += 1
-        return tot / n                                          # mean per-base logp
+        return tot / max(n, 1)                                  # mean per-base logp