From e83cd73ed654ed7706905629d7acbc664110f54b Mon Sep 17 00:00:00 2001
From: "Brandon T. Willard" <brandonwillard@users.noreply.github.com>
Date: Mon, 11 Sep 2023 15:51:32 -0500
Subject: [PATCH] Reset FSM states in Regex after completion

---
 outlines/text/generate/regex.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/outlines/text/generate/regex.py b/outlines/text/generate/regex.py
index 76a0a4c38..15e28152d 100644
--- a/outlines/text/generate/regex.py
+++ b/outlines/text/generate/regex.py
@@ -146,6 +146,10 @@ def create_proposal(
 
         return logits + mask
 
+    def postprocess_completions(self, completions: List[str]) -> List[str]:
+        self.pstates.clear()
+        return super().postprocess_completions(completions)
+
 
 def regex(model, regex_string: str, max_tokens: Optional[int] = None):
     """Generate text sequences that match the input regex.