naming

turbo-llm · Nov 15, 2024 · 005b64e · 005b64e
1 parent a31c01e
commit 005b64e
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/turbo_alignment/dataset/chat/chat.py b/turbo_alignment/dataset/chat/chat.py
@@ -94,15 +94,15 @@ def __keep_start(
         inference: bool,
         max_tokens: int | None = None,
     ) -> tuple[int, int]:
-        continuation = False
+        bot_message_was_presented = False
         for i, (message, end_index) in enumerate(zip(conversation.messages[::-1], replicas_cum_len[::-1])):
             if self.settings.only_answer_loss:
-                if inference and message.role == ChatMessageRole.BOT and not continuation:
-                    continuation = True
+                if inference and message.role == ChatMessageRole.BOT and not bot_message_was_presented:
+                    bot_message_was_presented = True
                     continue
                 if not inference and message.role != ChatMessageRole.BOT:
                     continue
-                continuation = False
+                bot_message_was_presented = False
 
             if max_tokens is None or end_index < max_tokens:
                 return 0, len(replicas_cum_len) - i