Spaces:
Sleeping
Sleeping
File size: 1,986 Bytes
787c99c |
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 67 68 69 70 71 72 |
#!/usr/bin/env python
import argparse
import json
from pathlib import Path
ROLE_TO_REWARD = {
"expert": 1.0,
"medium": 0.0,
"bad": -0.5,
}
def rewrite_file(input_path: Path, output_path: Path) -> None:
"""Rewrite rewards in a Battlegrounds RLAIF JSONL file based on candidate role.
Mapping:
- expert -> 1.0
- medium -> 0.0
- bad -> -0.5
Other roles keep their original reward.
"""
with input_path.open("r", encoding="utf-8") as fin, output_path.open(
"w", encoding="utf-8"
) as fout:
for line in fin:
line = line.strip()
if not line:
continue
obj = json.loads(line)
candidates = obj.get("candidates", [])
for cand in candidates:
role = cand.get("role")
if role in ROLE_TO_REWARD:
cand["reward"] = ROLE_TO_REWARD[role]
fout.write(json.dumps(obj, ensure_ascii=False) + "\n")
def main() -> None:
parser = argparse.ArgumentParser(
description="Rewrite rewards in Battlegrounds RLAIF JSONL dataset based on candidate role."
)
parser.add_argument(
"--input",
type=str,
required=True,
help="Path to input JSONL file (original dataset).",
)
parser.add_argument(
"--output",
type=str,
required=True,
help="Path to output JSONL file (rewritten dataset).",
)
args = parser.parse_args()
input_path = Path(args.input)
output_path = Path(args.output)
if not input_path.exists():
raise SystemExit(f"Input file not found: {input_path}")
if output_path.exists() and input_path.resolve() == output_path.resolve():
raise SystemExit("Refusing to overwrite input file in-place. Use a different --output path.")
rewrite_file(input_path, output_path)
print(f"Rewritten dataset saved to: {output_path}")
if __name__ == "__main__":
main()
|