Update README.md

README.md

@@ -62,6 +62,56 @@ We apply tailored prompts for coding and math task:
 {question} + "\n\nPresent the answer in LaTex format: \\boxed{Your answer}"
 ```
 
+
+```python
+import os
+from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer
+from vllm import LLM, SamplingParams
+os.environ["NCCL_IGNORE_DISABLED_P2P"] = "1"
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+def generate(question_list,model_path):
+    llm = LLM(
+        model=model_path,
+        trust_remote_code=True,
+        tensor_parallel_size=torch.cuda.device_count(),
+        gpu_memory_utilization=0.90,
+    )
+    sampling_params = SamplingParams(max_tokens=8192,
+                                    temperature=0.0,
+                                    n=1)
+    outputs = llm.generate(question_list, sampling_params, use_tqdm=True)
+    completions = [[output.text for output in output_item.outputs] for output_item in outputs]
+    return completions
+
+def make_conv_hf(question, tokenizer):
+    # for math problem
+    content = question + "\n\nPresent the answer in LaTex format: \\boxed{Your answer}"
+    # for code problem
+    # content = question + "\n\nWrite Python code to solve the problem. Present the code in \n```python\nYour code\n```\nat the end." 
+    msg = [
+        {"role": "user", "content": content}
+    ]
+    chat = tokenizer.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
+    return chat
+    
+def run():
+    model_path = "PRIME-RL/Eurus-2-7B-PRIME"
+    all_problems = [
+        "which number is larger? 9.11 or 9.9?"
+    ]
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    completions = generate([make_conv_hf(problem_data, tokenizer) for problem_data in all_problems],model_path)
+    print(completions)
+    # [['[ASSESS]\n\n# The problem asks us to compare two decimal numbers, 9.11 and 9.9, to determine which one is larger.\n# We need to compare the whole parts and the decimal parts of the numbers.\n\nNext action: [ADVANCE]\n\n# Compare the whole parts of the numbers: both 9.11 and 9.9 have the same whole part, which is 9.\n# Compare the decimal parts of the numbers: 0.11 (from 9.11) is less than 0.9 (from 9.9).\n\nNext action: [ADVANCE]\n\n# Since the whole parts are the same and the decimal part of 9.9 is greater than the decimal part of 9.11, we can conclude that 9.9 is larger than 9.11.\n\nNext action: [OUTPUT]\n\nThe final answer is $\\boxed{9.9}$.\n\n']]
+if __name__ == "__main__":
+    run()
+```
+
+
+
 ## Evaluation
 
 Through PRIME, we successfully achieve substantial improvements on key reasoning benchmarks over our SFT version of the model, leading to **16.7%** improvement on average, and over **20%** on AMC&AIME competitions. Our final model Eurus-2-7B-PRIME, based on Qwen-2.5-Math-7B-Base, surpassed its instruct version on 5 key reasoning benchmarks.