llama fixes

Author: jacklanchantin

src/fairseq2/recipes/lm/_online_finetune/_generative_judge.py

@@ -377,7 +377,7 @@ def remove_think_tags(self, rollout_text):
             text_without_thought = text_without_thought.replace("<|endoftext|>", "")
             return text_without_thought
         else:
-            return ""  # set rollout to empty string if it doesn't contain thought or has multiple
+            return "DUMMY"  # set rollout to dummy string if it doesn't contain thought or has multiple
 
     @override
     def format_prompt(

src/fairseq2/recipes/lm/_online_finetune/_grpo.py

@@ -74,7 +74,80 @@ class GRPOBatch:
     rewards: torch.Tensor
 
 
-def clip_outputs_after_think_token(rollouts, tokenizer, think_tokens, num_tokens):
+# def clip_outputs_after_think_token(rollouts, tokenizer, num_tokens):
+#     """
+#     Clip token_ids and logprobs to keep only num_tokens after the </think> token sequence ends.
+#     If </think> is not found, clip to just the first num_tokens.
+#     Recompute the text from clipped tokens.
+
+#     Args:
+#         rollouts: List of rollout objects
+#         tokenizer: Tokenizer instance
+#         think_tokens: List of token IDs for </think>
+#         num_tokens: Number of tokens to keep after </think> token sequence ends (or from start if no </think>)
+
+#     Returns:
+#         List of modified rollout objects
+#     """
+#     think_tokens = tokenizer.encode("</think>", add_special_tokens=False)
+#     ret = []
+#     for rollout in rollouts:
+#         clipped_outputs = []
+
+#         for output in rollout.outputs:
+#             # Find the position where </think> tokens start
+#             think_token_len = len(think_tokens)
+#             clip_index = None
+
+#             # Search for the think tokens sequence in token_ids
+#             for i in range(len(output.token_ids) - think_token_len + 1):
+#                 if output.token_ids[i : i + think_token_len] == think_tokens:
+#                     # Clip to include everything up to and including </think> plus num_tokens after
+#                     clip_index = i + think_token_len + num_tokens
+#                     break
+
+#             # If </think> not found, clip to just the first num_tokens
+#             if clip_index is None:
+#                 clip_index = num_tokens
+
+#             # Clip token_ids and logprobs
+#             clipped_token_ids = output.token_ids[:clip_index]
+#             clipped_logprobs = output.logprobs[:clip_index]
+
+#             # Recompute text from clipped tokens
+#             clipped_text = tokenizer.decode(clipped_token_ids)
+
+#             # Recalculate cumulative_logprob from clipped logprobs
+#             cumulative_logprob = 0.0
+#             for logprob_dict in clipped_logprobs:
+#                 # Get the first token's logprob (the selected token)
+#                 first_token_id = list(logprob_dict.keys())[0]
+#                 cumulative_logprob += logprob_dict[first_token_id].logprob
+
+#             # Create new CompletionOutput with clipped data
+#             clipped_output = type(output)(
+#                 index=output.index,
+#                 text=clipped_text,
+#                 token_ids=clipped_token_ids,
+#                 cumulative_logprob=cumulative_logprob,
+#                 logprobs=clipped_logprobs,
+#                 finish_reason=output.finish_reason,
+#                 stop_reason=output.stop_reason,
+#             )
+#             clipped_outputs.append(clipped_output)
+
+#         # Create new rollout object with clipped outputs
+#         clipped_rollout = type(rollout)(
+#             outputs=clipped_outputs,
+#             # Copy other attributes from original rollout
+#             **{k: v for k, v in vars(rollout).items() if k != "outputs"},
+#         )
+#         ret.append(clipped_rollout)
+
+#     return ret
+
+
+def clip_outputs_after_think_token(rollouts, tokenizer, num_tokens):
     """
     Clip token_ids and logprobs to keep only num_tokens after the </think> token sequence ends.
     If </think> is not found, clip to just the first num_tokens.
@@ -94,16 +167,24 @@ def clip_outputs_after_think_token(rollouts, tokenizer, think_tokens, num_tokens
         clipped_outputs = []
 
         for output in rollout.outputs:
-            # Find the position where </think> tokens start
-            think_token_len = len(think_tokens)
             clip_index = None
 
-            # Search for the think tokens sequence in token_ids
-            for i in range(len(output.token_ids) - think_token_len + 1):
-                if output.token_ids[i : i + think_token_len] == think_tokens:
-                    # Clip to include everything up to and including </think> plus num_tokens after
-                    clip_index = i + think_token_len + num_tokens
-                    break
+            # Search for "</think>" in the decoded text to avoid tokenizer context issues
+            text = output.text
+            think_end_pos = text.find("</think>")
+
+            if think_end_pos != -1:
+                # Found "</think>" in text - now find the corresponding token position
+                # The end of "</think>" in the text
+                think_end_pos += len("</think>")
+
+                # Decode tokens incrementally to find where we pass the think_end_pos
+                for i in range(1, len(output.token_ids) + 1):
+                    decoded_text = tokenizer.decode(output.token_ids[:i])
+                    if len(decoded_text) >= think_end_pos:
+                        # We've passed the end of "</think>", now add num_tokens more
+                        clip_index = min(i + num_tokens, len(output.token_ids))
+                        break
 
             # If </think> not found, clip to just the first num_tokens
             if clip_index is None:
@@ -370,13 +451,9 @@ def __call__(
                     ]
                     for text in prompt_batch.meta_info.get("suffix")
                 ]
-                think_tokens = self._rollout_tokenizer.encode(
-                    "</think>", add_special_tokens=False
-                )
                 rollouts = clip_outputs_after_think_token(
                     rollouts,
                     self._rollout_tokenizer,
-                    think_tokens,
                     self._config.clip_rollout_after_think,
                 )
             if self._config.loss_config.log_rollouts: