Spaces:
Running
Running
File size: 1,755 Bytes
ac819bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import asyncio
import json
import time
import ray
from datasets import load_dataset
from lagent.distributed.ray_serve import AsyncAgentRayActor
from lagent.llms import INTERNLM2_META
from lagent.llms.lmdeploy_wrapper import AsyncLMDeployPipeline
ray.init()
# set up the loop
# initialize the model
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
model = dict(
type=AsyncLMDeployPipeline,
path='internlm/internlm2_5-7b-chat',
meta_template=INTERNLM2_META,
tp=1,
top_k=1,
temperature=1.0,
stop_words=['<|im_end|>', '<|action_end|>'],
max_new_tokens=1024,
)
# ----------------------- interpreter -----------------------
print('-' * 80, 'interpreter', '-' * 80)
ds = load_dataset('lighteval/MATH', split='test')
problems = [item['problem'] for item in ds.select(range(5000))]
coder = dict(
type='lagent.agents.stream.AsyncMathCoder',
llm=model,
interpreter=dict(type='AsyncIPythonInterpreter', max_kernels=300),
)
tic = time.time()
actor1 = AsyncAgentRayActor(coder.copy(), num_gpus=1)
actor2 = AsyncAgentRayActor(coder.copy(), num_gpus=1)
corots = [
actor1(query, session_id=i)
for i, query in enumerate(problems[:len(problems) // 2])
]
corots += [
actor2(query, session_id=i)
for i, query in enumerate(problems[len(problems) // 2:])
]
results = loop.run_until_complete(asyncio.gather(*corots))
print('-' * 120)
print(f'time elapsed: {time.time() - tic}')
all_step = ray.get([
actor1.agent_actor.get_steps.remote(i) for i in range(len(problems) // 2)
])
all_step += ray.get([
actor2.agent_actor.get_steps.remote(i)
for i in range(len(problems[len(problems) // 2:]))
])
with open('./tmp_1.json', 'w') as f:
json.dump(all_step, f, ensure_ascii=False, indent=4)
|