Add optional sentence-boundary generation stop

tiny_english/generate.py

@@ -49,10 +49,10 @@ def load_export(directory, device="cpu"):
     return model.to(device=device, dtype=torch.bfloat16).eval(), tokenizer, metadata["checkpoint_step"]
 
 
-def complete(model, tokenizer, prompt, max_new_tokens=150, temperature=0.8, top_k=40, seed=42, deadline=None, strict_context=False):
+def complete(model, tokenizer, prompt, max_new_tokens=150, temperature=0.8, top_k=40, seed=42, deadline=None, strict_context=False, stop_at_sentence_end=False, min_new_tokens=0):
     if not isinstance(prompt, str) or not prompt.strip():
         raise ValueError("Prompt must not be empty")
-    if not math.isfinite(temperature) or temperature <= 0 or top_k < 1 or max_new_tokens < 1:
+    if not math.isfinite(temperature) or temperature <= 0 or top_k < 1 or max_new_tokens < 1 or min_new_tokens < 0:
         raise ValueError("Generation numeric options must be positive and finite")
     tokens = tokenizer.encode(prompt).ids
     prompt_tokens = len(tokens)
@@ -71,10 +71,16 @@ def complete(model, tokenizer, prompt, max_new_tokens=150, temperature=0.8, top_
             logits = model(inputs)[0, -1].float() / temperature
             values, indices = logits.topk(min(top_k, logits.numel()))
             next_token = indices[torch.multinomial(values.softmax(-1), 1, generator=generator)].item()
-            if next_token == tokenizer.token_to_id("<eos>"):
-                finish_reason = "eos"
-                break
-            tokens.append(next_token)
+            if next_token == tokenizer.token_to_id("<eos>"):
+                finish_reason = "eos"
+                break
+            tokens.append(next_token)
+            if stop_at_sentence_end and len(tokens) - prompt_tokens >= min_new_tokens:
+                ending = tokenizer.decode([next_token]).rstrip()
+                completion = tokenizer.decode(tokens[prompt_tokens:])
+                if ending.endswith((".", "!", "?")) and completion.count('"') % 2 == 0:
+                    finish_reason = "sentence"
+                    break
     return {
         "text": tokenizer.decode(tokens),
         "completion": tokenizer.decode(tokens[prompt_tokens:]),