Add asyncio.sleep(0) to sample_word_2

probcomp · Nov 18, 2024 · 6d8caff · 6d8caff
1 parent 1cb07a9
commit 6d8caff
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/hfppl/chunks.py b/hfppl/chunks.py
@@ -1,3 +1,4 @@
+import asyncio
 import string
 
 from .modeling import submodel
@@ -73,6 +74,8 @@ async def sample_word_2(
     Returns:
         Tuple[str, str]: The sampled word and punctuation
     """
+    # NOTE: Yields control back to the event loop. Necessary to allow timeouts to work correctly when this method is called in a loop.
+    await asyncio.sleep(0)
 
     # This approach sometimes breaks with max_chars = 1
     if max_chars is not None: