Upload daisychain.py with huggingface_hub

daisychain.py

@@ -102,25 +102,34 @@ class DaisyChain:
         return best, bpb
 
     @torch.no_grad()
-    def generate_stream(self, domain, length=180, temperature=0.9, top_k=20, prompt=""):
-        """Yield the growing continuation base-by-base (for live streaming UIs)."""
+    def generate_stream(self, domain, length=180, temperature=1.0, top_k=40,
+                        repetition_penalty=1.3, prompt=""):
+        """Yield the growing continuation base-by-base (for live streaming UIs).
+        repetition_penalty discourages the homopolymer/repeat loops these small
+        specialists fall into; it does not erase genuine low-complexity domains."""
         m = self.models[domain]
         ids = [self.bos] + (self.tok.encode(clean(prompt), add_special_tokens=False) if prompt else [])
         t = torch.tensor([ids], device=self.dev)
-        bases = []
+        bases, emitted = [], []
         while sum(len(b) for b in bases) < length:
-            logits = m(input_ids=t).logits[:, -1, :] / max(temperature, 1e-6)
-            logits[:, :4] = -1e9
+            logits = m(input_ids=t).logits[:, -1, :].float() / max(temperature, 1e-6)
+            logits[:, :4] = -1e9                                   # never emit specials
+            if repetition_penalty and repetition_penalty != 1.0:   # penalize recent tokens
+                for tid in set(emitted[-12:]):
+                    v = logits[0, tid]
+                    logits[0, tid] = v / repetition_penalty if v > 0 else v * repetition_penalty
             if top_k > 0:
                 v, _ = torch.topk(logits, top_k)
                 logits[logits < v[:, [-1]]] = -1e9
             nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
+            ti = int(nxt); emitted.append(ti)
             t = torch.cat([t, nxt], dim=1)
-            bases.append(self.tok._ids_to_tokens[int(nxt)])
+            bases.append(self.tok._ids_to_tokens[ti])
             yield "".join(bases)[:length]
 
-    def generate(self, domain, length=180, temperature=0.9, top_k=20, prompt=""):
+    def generate(self, domain, length=180, temperature=1.0, top_k=40,
+                 repetition_penalty=1.3, prompt=""):
         out = ""
-        for out in self.generate_stream(domain, length, temperature, top_k, prompt):
+        for out in self.generate_stream(domain, length, temperature, top_k, repetition_penalty, prompt):
             pass
         return out