# 基于目前最新的vllm 0.7.2进行验证pip install vllm -U
# 基于目前最新的trl 0.15.1进行验证pip install trl -U

import re
import torchfrom modelscope import AutoTokenizer, AutoModelForCausalLMfrom modelscope.msdatasets import MsDatasetfrom trl import GRPOConfig, GRPOTrainerSYSTEM_PROMPT = """You need to answer in XML format, include <reasoning> and <answer>, respond in the following format:<reasoning>...</reasoning><answer>...</answer>"""XML_COT_FORMAT = """\<reasoning>{reasoning}</reasoning><answer>{answer}</answer>"""def extract_xml_answer(text: str) -> str:    answer = text.split("<answer>")[-1]    answer = answer.split("</answer>")[0]    return answer.strip()def extract_hash_answer(text: str) -> str | None:    if "####" not in text:        return None    return text.split("####")[1].strip()def get_gsm8k_questions(split="train") -> MsDataset:    data = MsDataset.load('modelscope/gsm8k', subset_name='main', split=split)    data = data.map(lambda x: {        'prompt': [            {'role': 'system', 'content': SYSTEM_PROMPT},            {'role': 'user', 'content': x['question']}        ],        'answer': extract_hash_answer(x['answer'])    })    return datadataset = get_gsm8k_questions()# Reward functionsdef correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:    responses = [completion[0]['content'] for completion in completions]    q = prompts[0][-1]['content']    extracted_responses = [extract_xml_answer(r) for r in responses]    print('-' * 20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}",          f"\nExtracted:\n{extracted_responses[0]}")    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]def int_reward_func(completions, **kwargs) -> list[float]:    responses = [completion[0]['content'] for completion in completions]    extracted_responses = [extract_xml_answer(r) for r in responses]    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]# def strict_format_reward_func(completions, **kwargs) -> list[float]:#     pattern = r"\n<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"#     responses = [completion[0]["content"] for completion in completions]#     matches = [re.fullmatch(pattern, r, re.DOTALL) for r in responses]#     return [0.5 if match else 0.0 for match in matches]def strict_format_reward_func(completions, **kwargs) -> list[float]:    pattern = r"<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>"    responses = [completion[0]["content"] for completion in completions]    # 新增调试日志    matches = []    for idx, r in enumerate(responses):        print(f"\n--- Processing response {idx} ---")        print("Raw content:", repr(r))  # 使用 repr() 显示转义字符        match = re.fullmatch(pattern, r, re.DOTALL)        print("Match result:", bool(match))        matches.append(match)    return [0.5 if match else 0.0 for match in matches]def soft_format_reward_func(completions, **kwargs) -> list[float]:    pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"    responses = [completion[0]["content"] for completion in completions]    matches = [re.fullmatch(pattern, r, re.DOTALL) for r in responses]    return [0.5 if match else 0.0 for match in matches]def count_xml(text) -> float:    count = 0.0    if text.count("<reasoning>\n") == 1:        count += 0.125    if text.count("\n</reasoning>\n") == 1:        count += 0.125    if text.count("\n<answer>\n") == 1:        count += 0.125        count -= len(text.split("\n</answer>\n")[-1]) * 0.001    if text.count("\n</answer>") == 1:        count += 0.125        count -= (len(text.split("\n</answer>")[-1]) - 1) * 0.001    return countdef xmlcount_reward_func(completions, **kwargs) -> list[float]:    contents = [completion[0]["content"] for completion in completions]    return [count_xml(c) for c in contents]model_name = "Qwen/Qwen2.5-0.5B-Instruct"output_dir = "outputs/Qwen-0.5B-GRPO"run_name = "Qwen-0.5B-GRPO-gsm8k"training_args = GRPOConfig(    output_dir=output_dir,    run_name=run_name,    learning_rate=5e-6,    adam_beta1=0.9,    adam_beta2=0.99,    weight_decay=0.1,    warmup_ratio=0.1,    lr_scheduler_type='cosine',    logging_steps=1,    bf16=True,    per_device_train_batch_size=8,    gradient_accumulation_steps=4,    num_generations=8,    max_prompt_length=256,    max_completion_length=200,    num_train_epochs=1,    save_steps=100,    max_grad_norm=0.1,    log_on_each_node=False,    use_vllm=True,    vllm_gpu_memory_utilization=.2,    vllm_device="cuda:0",    report_to="none")model = AutoModelForCausalLM.from_pretrained(    model_name,    torch_dtype=torch.bfloat16,    device_map=None).to("cuda")tokenizer = AutoTokenizer.from_pretrained(model_name)tokenizer.pad_token = tokenizer.eos_tokentrainer = GRPOTrainer(    model=model,    processing_class=tokenizer,    reward_funcs=[        xmlcount_reward_func,        soft_format_reward_func,        strict_format_reward_func,        int_reward_func,        correctness_reward_func],    args=training_args,    train_dataset=dataset,)trainer.train()