diff --git a/turbo_alignment/dataset/chat/chat.py b/turbo_alignment/dataset/chat/chat.py index b874c3f..c96af58 100755 --- a/turbo_alignment/dataset/chat/chat.py +++ b/turbo_alignment/dataset/chat/chat.py @@ -94,10 +94,8 @@ def __keep_start( inference: bool, max_tokens: int | None = None, ) -> tuple[int, int]: - print(max_tokens, conversation.messages) continuation = False for i, (message, end_index) in enumerate(zip(conversation.messages[::-1], replicas_cum_len[::-1])): - print(i, message, end_index) if self.settings.only_answer_loss: if inference and message.role == ChatMessageRole.BOT and not continuation: continuation = True @@ -107,7 +105,6 @@ def __keep_start( continuation = False if max_tokens is None or end_index < max_tokens: - print('RETURN', len(replicas_cum_len) - i) return 0, len(replicas_cum_len) - i raise ValueError('Can\'t trim dialogue to fit all requirements') @@ -203,7 +200,6 @@ def _truncate_and_merge( labels = np.array([]) truncated_conversation_messages = conversation.messages[left_bound:right_bound] - print("TRUNCATED", truncated_conversation_messages) truncated_tokenized_replicas = tokenized_replicas[left_bound:right_bound] if self.source.system_prompt is not None and self.settings.keep_end and left_bound != 0: