| import asyncio |
| import json |
| import time |
|
|
| from datasets import load_dataset |
|
|
| from lagent.agents import AsyncMathCoder |
| from lagent.agents.aggregator import InternLMToolAggregator |
| from lagent.llms import AsyncGPTAPI |
| from lagent.prompts.parsers import ToolParser |
|
|
| loop = asyncio.new_event_loop() |
| asyncio.set_event_loop(loop) |
|
|
| interpreter_prompt = ( |
| 'Below is a math problem. Please solve it step by step with the assistance of Python programming. Consider using Sympy or Numpy library ' |
| 'to facilitate your derivation, calculation and equation solving. Utilize the "pi" symbol and "Rational" from Sympy ' |
| 'for $$\pi$$ and fractions, and simplify all fractions and square roots without converting them to decimal values. ' |
| 'Please encapsulate each generated Jupyter Python code block with tags "{begin}" and "{end}". Conclude the ' |
| r'final answer when observations are sufficient and encapsulate the numerical result with LaTeX syntax "\boxed{{}}" ' |
| 'without any unit, and end your conclusion with the special token "[END]" to denote the completion of your response. ' |
| 'Keep the following points in mind:\n' |
| '- You must alternately use human and programming languages in the chain of thought;\n' |
| '- The number of your reasoning steps should not exceed **three**, which means you may merge some intermediate steps when the original answer is tedious.' |
| ) |
|
|
| async_llm = dict( |
| type=AsyncGPTAPI, |
| model='gpt-4o-2024-05-13', |
| retry=50, |
| key='', |
| max_new_tokens=2048, |
| stop_words=['</python'], |
| proxies=dict(), |
| ) |
| async_agent = AsyncMathCoder( |
| llm=async_llm, |
| output_format=ToolParser( |
| tool_type='interpreter', |
| template=interpreter_prompt, |
| begin='<python>\n', |
| end='\n</python>'), |
| aggregator=InternLMToolAggregator( |
| environment_role='system', |
| environment_begin='<output>\n', |
| environment_end='\n</output>'), |
| finish_condition=lambda m: '[END]' in m.content, |
| ) |
|
|
| ds = load_dataset('lighteval/MATH', split='train') |
| problems = [item['problem'] for item in ds.select(range(30))] |
|
|
| tic = time.time() |
| coros = [async_agent(q, session_id=i) for i, q in enumerate(problems)] |
| res = loop.run_until_complete(asyncio.gather(*coros)) |
| print(time.time() - tic) |
| with open('tmp_6.json', 'w') as f: |
| json.dump([async_agent.get_steps(i) for i in range(len(problems))], |
| f, |
| ensure_ascii=False, |
| indent=4) |
|
|
