Update rl_replacements.py

This commit is contained in:
Daniel Han 2025-07-17 06:47:23 -07:00
commit fd1547cd7b

View file

@ -247,7 +247,7 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
# Always between max_prompt_length and use_vllm
found = re.findall(
r"(\n([\s]{1,})if self\.max_prompt_length is not None:.*?"\
r"\n(([\s]{1,})if self\.max_prompt_length is not None:.*?"\
r"\2if self\.use_vllm:)",
function,
flags = re.DOTALL | re.MULTILINE,
@ -256,14 +256,11 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
replace_part, spacing = found[0]
removed_comments = re.sub(r"\#[^\n]{1,}", "", replace_part)
splits = removed_comments.split("\n")
print("##########")
print("#", spacing, "#")
print(splits)
print(sum(re.match(rf"{spacing}[^\s]", x) is not None for x in splits))
if sum(re.match(rf"^{spacing}[^\s]", x) is not None for x in splits) == 2 and len(spacing) == 8:
if sum(re.match(rf"{spacing}[^\s]", x) is not None for x in splits) == 2 and len(spacing) == 8:
new_replacement = spacing + \
"""if self.max_prompt_length is not None:
new_replacement = \
f"""\n{spacing}if self.max_prompt_length is not None:
# If max_prompt_length is set, we trim the prompt to keep only the last `max_prompt_length` tokens.
# Then we decode those tokens back into text. We manually remove leading pad tokens from the decoded text,
# because we can't use `skip_special_tokens=True` (some special tokens are still needed for generation).