fix CFG Generation

outlines-dev · May 5, 2024 · 6198be1 · 6198be1
1 parent e295631
commit 6198be1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/outlines/generate/generator.py b/outlines/generate/generator.py
@@ -86,8 +86,9 @@ def sequence_generator(
         token_ids = update_token_ids(token_ids, next_token_ids, ancestors)
         attention_masks = update_attention_masks(attention_masks, ancestors)
         kv_cache = reorder_kv_cache(kv_cache, ancestors)
-        fsms = reorder_fsms(fsms, ancestors)
-        fsm_states = reorder_fsm_states(fsm_states, ancestors)
+        if len(ancestors) > 1:
+            fsms = reorder_fsms(fsms, ancestors)
+            fsm_states = reorder_fsm_states(fsm_states, ancestors)
 
         fsm_states = get_next_fsm_states(fsms, fsm_states, next_token_ids)
         is_finished = is_generation_finished(fsms, fsm_states)