54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import asyncio
|
|
from typing import Annotated
|
|
|
|
import ray
|
|
import torch
|
|
import typer
|
|
|
|
from slime.utils.misc import load_function
|
|
from slime.utils.types import Sample
|
|
|
|
|
|
def _truncate(text, max_len=200):
|
|
"""Truncate text and add ellipsis if too long."""
|
|
if text is None:
|
|
return None
|
|
text = str(text).replace("\n", "\\n")
|
|
if len(text) > max_len:
|
|
return text[:max_len] + "..."
|
|
return text
|
|
|
|
|
|
def main(
|
|
rollout_data_path: Annotated[str, typer.Option()],
|
|
custom_rm_path: Annotated[str, typer.Option()],
|
|
):
|
|
if not ray.is_initialized():
|
|
ray.init()
|
|
|
|
pack = torch.load(rollout_data_path)
|
|
samples = [Sample.from_dict(s) for s in pack["samples"]]
|
|
asyncio.run(_main_async(samples=samples, custom_rm_path=custom_rm_path))
|
|
|
|
|
|
async def _main_async(samples, custom_rm_path):
|
|
rm_function = load_function(custom_rm_path)
|
|
rewards = await asyncio.gather(*[rm_function(None, sample) for sample in samples])
|
|
|
|
for i, (sample, reward) in enumerate(zip(samples, rewards, strict=True)):
|
|
print("-" * 60)
|
|
print(f"Sample {i + 1}/{len(samples)}")
|
|
print(f" Index: {sample.index}")
|
|
print(f" Status: {sample.status}")
|
|
print(f" Reward: {reward}")
|
|
print(f" Prompt: {_truncate(sample.prompt, 200)}")
|
|
print(f" Response: {_truncate(sample.response, 200)}")
|
|
print("-" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|