初始化项目,由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
53
slime/utils/debug_utils/replay_reward_fn.py
Normal file
53
slime/utils/debug_utils/replay_reward_fn.py
Normal file
@@ -0,0 +1,53 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user