This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from swift.utils.import_utils import _LazyModule
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .orm import ORM, AsyncORM, orms
|
||||
from .prm import PRM, prms
|
||||
from .rm_plugin import rm_plugins
|
||||
|
||||
else:
|
||||
_import_structure = {
|
||||
'prm': ['prms', 'PRM'],
|
||||
'orm': ['orms', 'ORM', 'AsyncORM'],
|
||||
'rm_plugin': ['rm_plugins'],
|
||||
}
|
||||
|
||||
import sys
|
||||
|
||||
sys.modules[__name__] = _LazyModule(
|
||||
__name__,
|
||||
globals()['__file__'],
|
||||
_import_structure,
|
||||
module_spec=__spec__,
|
||||
extra_objects={},
|
||||
)
|
||||
@@ -0,0 +1,464 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
# Outcome Reward Model (ORM) implementations for GRPO training.
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Union
|
||||
|
||||
from swift.infer_engine import InferRequest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from swift.megatron.arguments import MegatronArguments
|
||||
from swift.rlhf_trainers import GRPOConfig
|
||||
|
||||
|
||||
class ORM:
|
||||
"""Base class for synchronous outcome reward models (ORM).
|
||||
|
||||
Subclasses should implement the __call__ method to compute rewards.
|
||||
|
||||
Example:
|
||||
class MyReward(ORM):
|
||||
def __call__(self, completions, **kwargs) -> List[float]:
|
||||
return [1.0 if len(c) > 100 else 0.0 for c in completions]
|
||||
"""
|
||||
|
||||
def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
|
||||
self.args = args
|
||||
|
||||
def __call__(self, **kwargs) -> List[float]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AsyncORM:
|
||||
"""Base class for asynchronous outcome reward models (ORM).
|
||||
|
||||
Use this for reward functions that involve I/O operations (e.g., API calls,
|
||||
database queries) that can benefit from async execution.
|
||||
|
||||
Async reward functions are executed in parallel using asyncio.gather,
|
||||
which can significantly speed up reward computation when multiple async
|
||||
reward functions are used or when the reward function involves network calls.
|
||||
|
||||
Example:
|
||||
class MyAsyncReward(AsyncORM):
|
||||
async def __call__(self, completions, **kwargs) -> List[float]:
|
||||
# Use asyncio.gather for parallel execution of all API calls
|
||||
import asyncio
|
||||
import aiohttp
|
||||
|
||||
async def score_single(session, text):
|
||||
async with session.post(api_url, json={'text': text}) as resp:
|
||||
result = await resp.json()
|
||||
return result['score']
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [score_single(session, c) for c in completions]
|
||||
rewards = await asyncio.gather(*tasks)
|
||||
return list(rewards)
|
||||
"""
|
||||
|
||||
def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
|
||||
self.args = args
|
||||
|
||||
async def __call__(self, **kwargs) -> List[float]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class MathAccuracy(ORM):
|
||||
|
||||
def __init__(self, args=None, **kwargs):
|
||||
super().__init__(args, **kwargs)
|
||||
import importlib.util
|
||||
assert importlib.util.find_spec('math_verify') is not None, (
|
||||
'The math_verify package is required but not installed. '
|
||||
"Please install it using 'pip install math_verify'.")
|
||||
|
||||
def __call__(self, completions, solution, **kwargs) -> List[float]:
|
||||
from latex2sympy2_extended import NormalizationConfig
|
||||
from math_verify import LatexExtractionConfig, parse, verify
|
||||
rewards = []
|
||||
for content, sol in zip(completions, solution):
|
||||
content_match = re.search(r'<answer>(.*?)</answer>', content, re.DOTALL)
|
||||
content_to_parse = content_match.group(1).strip() if content_match else content
|
||||
has_answer_tag = content_match is not None
|
||||
|
||||
sol_match = re.search(r'<answer>(.*?)</answer>', sol, re.DOTALL)
|
||||
sol_to_parse = sol_match.group(1).strip() if sol_match else sol
|
||||
|
||||
gold_parsed = parse(sol_to_parse, extraction_mode='first_match')
|
||||
if len(gold_parsed) != 0:
|
||||
if has_answer_tag:
|
||||
answer_parsed = parse(content_to_parse, extraction_mode='first_match')
|
||||
else:
|
||||
answer_parsed = parse(
|
||||
content_to_parse,
|
||||
extraction_config=[
|
||||
LatexExtractionConfig(
|
||||
normalization_config=NormalizationConfig(
|
||||
nits=False,
|
||||
malformed_operators=False,
|
||||
basic_latex=True,
|
||||
boxed=True,
|
||||
units=True,
|
||||
),
|
||||
boxed_match_priority=0,
|
||||
try_extract_without_anchor=False,
|
||||
)
|
||||
],
|
||||
extraction_mode='first_match',
|
||||
)
|
||||
try:
|
||||
reward = float(verify(gold_parsed, answer_parsed))
|
||||
except Exception:
|
||||
reward = 0.0
|
||||
else:
|
||||
# If the gold solution is not parseable, we reward 0 to skip this example
|
||||
reward = 0.0
|
||||
rewards.append(reward)
|
||||
return rewards
|
||||
|
||||
|
||||
class Format(ORM):
|
||||
|
||||
def __call__(self, completions, **kwargs) -> List[float]:
|
||||
"""Reward function that checks if the completion has a specific format."""
|
||||
pattern = r'^<think>.*?</think>\s*<answer>.*?</answer>(?![\s\S])'
|
||||
matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE) for content in completions]
|
||||
return [1.0 if match else 0.0 for match in matches]
|
||||
|
||||
|
||||
class ReActFormat(ORM):
|
||||
|
||||
def __call__(self, completions, **kwargs) -> List[float]:
|
||||
"""Reward function that checks if the completion has a specific format."""
|
||||
pattern = r'^<think>.*?</think>\s*Action:.*?Action Input:.*?$'
|
||||
matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE) for content in completions]
|
||||
return [1.0 if match else 0.0 for match in matches]
|
||||
|
||||
|
||||
class CosineReward(ORM):
|
||||
# https://arxiv.org/abs/2502.03373
|
||||
def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, accuracy_orm=None):
|
||||
super().__init__(args)
|
||||
self.min_len_value_wrong = args.cosine_min_len_value_wrong
|
||||
self.max_len_value_wrong = args.cosine_max_len_value_wrong
|
||||
self.min_len_value_correct = args.cosine_min_len_value_correct
|
||||
self.max_len_value_correct = args.cosine_max_len_value_correct
|
||||
self.max_len = args.cosine_max_len
|
||||
self.accuracy_orm = accuracy_orm or MathAccuracy()
|
||||
|
||||
@staticmethod
|
||||
def cosfn(t, T, min_value, max_value):
|
||||
import math
|
||||
return max_value - (max_value - min_value) * (1 - math.cos(t * math.pi / T)) / 2
|
||||
|
||||
def __call__(self, completions, solution, **kwargs) -> List[float]:
|
||||
acc_rewards = self.accuracy_orm(completions, solution, **kwargs)
|
||||
response_token_ids = kwargs.get('response_token_ids')
|
||||
rewards = []
|
||||
for ids, acc_reward in zip(response_token_ids, acc_rewards):
|
||||
is_correct = acc_reward >= 1.
|
||||
if is_correct:
|
||||
# Swap min/max for correct answers
|
||||
min_value = self.max_len_value_correct
|
||||
max_value = self.min_len_value_correct
|
||||
else:
|
||||
min_value = self.max_len_value_wrong
|
||||
max_value = self.min_len_value_wrong
|
||||
gen_len = len(ids)
|
||||
reward = self.cosfn(gen_len, self.max_len, min_value, max_value)
|
||||
rewards.append(reward)
|
||||
return rewards
|
||||
|
||||
|
||||
class RepetitionPenalty(ORM):
|
||||
# https://arxiv.org/abs/2502.03373
|
||||
def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
|
||||
super().__init__(args)
|
||||
self.ngram_size = args.repetition_n_grams
|
||||
self.max_penalty = args.repetition_max_penalty
|
||||
|
||||
@staticmethod
|
||||
def zipngram(text: str, ngram_size: int):
|
||||
words = text.lower().split()
|
||||
return zip(*[words[i:] for i in range(ngram_size)])
|
||||
|
||||
def __call__(self, completions, **kwargs) -> List[float]:
|
||||
"""
|
||||
reward function the penalizes repetitions
|
||||
|
||||
Args:
|
||||
completions: List of model completions
|
||||
"""
|
||||
rewards = []
|
||||
for completion in completions:
|
||||
if completion == '':
|
||||
rewards.append(0.0)
|
||||
continue
|
||||
if len(completion.split()) < self.ngram_size:
|
||||
rewards.append(0.0)
|
||||
continue
|
||||
|
||||
ngrams = set()
|
||||
total = 0
|
||||
for ng in self.zipngram(completion, self.ngram_size):
|
||||
ngrams.add(ng)
|
||||
total += 1
|
||||
|
||||
scaling = 1 - len(ngrams) / total
|
||||
reward = scaling * self.max_penalty
|
||||
rewards.append(reward)
|
||||
return rewards
|
||||
|
||||
|
||||
class SoftOverlong(ORM):
|
||||
|
||||
def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
|
||||
super().__init__(args)
|
||||
assert args.soft_cache_length < args.soft_max_length
|
||||
self.soft_max_length = args.soft_max_length
|
||||
self.soft_cache_length = args.soft_cache_length
|
||||
|
||||
def __call__(self, completions, **kwargs) -> List[float]:
|
||||
rewards = []
|
||||
response_token_ids = kwargs.get('response_token_ids')
|
||||
for ids in response_token_ids:
|
||||
completion_length = len(ids)
|
||||
expected_len = self.soft_max_length - self.soft_cache_length
|
||||
exceed_len = completion_length - expected_len
|
||||
rewards.append(min(-exceed_len / self.soft_cache_length, 0))
|
||||
return rewards
|
||||
|
||||
|
||||
class ReactORM(ORM):
|
||||
|
||||
@staticmethod
|
||||
def evaluate_action_reward(action_pred: list, action_ref: list, cand_list: list, ref_list: list):
|
||||
f1 = []
|
||||
for i in range(len(action_pred)):
|
||||
ref_action = action_ref[i]
|
||||
pred_action = action_pred[i]
|
||||
|
||||
ref_input = ref_list[i]
|
||||
cand_input = cand_list[i]
|
||||
|
||||
ref_is_json = False
|
||||
try:
|
||||
ref_input_json = json.loads(ref_input)
|
||||
ref_is_json = True
|
||||
except Exception:
|
||||
ref_input_json = ref_input
|
||||
|
||||
cand_is_json = False
|
||||
try:
|
||||
cand_input_json = json.loads(cand_input)
|
||||
cand_is_json = True
|
||||
except Exception:
|
||||
cand_input_json = cand_input
|
||||
|
||||
if ref_action != pred_action or (ref_is_json ^ cand_is_json):
|
||||
f1.append(0)
|
||||
elif not ref_is_json and not cand_is_json:
|
||||
rougel = ReactORM.evaluate_rougel([ref_input_json], [cand_input_json])
|
||||
if rougel is None or rougel < 10:
|
||||
f1.append(0)
|
||||
elif 10 <= rougel < 20:
|
||||
f1.append(0.1)
|
||||
else:
|
||||
f1.append(1)
|
||||
else:
|
||||
if not isinstance(ref_input_json, dict) or not isinstance(cand_input_json, dict):
|
||||
# This cannot be happen, but:
|
||||
# line 62, in evaluate_action_reward
|
||||
# for k, v in ref_input_json.items():
|
||||
# AttributeError: 'str' object has no attribute 'items'
|
||||
# print(f'>>>>>>ref_input_json: {ref_input_json}, cand_input_json: {cand_input_json}')
|
||||
f1.append(0)
|
||||
continue
|
||||
|
||||
half_match = 0
|
||||
full_match = 0
|
||||
if ref_input_json == {}:
|
||||
if cand_input_json == {}:
|
||||
f1.append(1)
|
||||
else:
|
||||
f1.append(0)
|
||||
else:
|
||||
for k, v in ref_input_json.items():
|
||||
if k in cand_input_json.keys():
|
||||
if cand_input_json[k] == v:
|
||||
full_match += 1
|
||||
else:
|
||||
half_match += 1
|
||||
|
||||
recall = (0.5 * half_match + full_match) / (len(ref_input_json) + 1e-30)
|
||||
precision = (0.5 * half_match + full_match) / (len(cand_input_json) + 1e-30)
|
||||
try:
|
||||
f1.append((2 * recall * precision) / (recall + precision))
|
||||
except Exception:
|
||||
f1.append(0.0)
|
||||
|
||||
if f1[0] == 1.0:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def parse_action(text):
|
||||
if 'Action Input:' in text:
|
||||
input_idx = text.rindex('Action Input:')
|
||||
action_input = text[input_idx + len('Action Input:'):].strip()
|
||||
else:
|
||||
action_input = '{}'
|
||||
|
||||
if 'Action:' in text:
|
||||
action_idx = text.rindex('Action:')
|
||||
action = text[action_idx + len('Action:'):].strip()
|
||||
if 'Action Input:' in action:
|
||||
input_idx = action.index('Action Input:')
|
||||
action = action[:input_idx].strip()
|
||||
else:
|
||||
action = 'none'
|
||||
return action, action_input
|
||||
|
||||
@staticmethod
|
||||
def parse_output(text):
|
||||
action, action_input = ReactORM.parse_action(text)
|
||||
return action, action_input
|
||||
|
||||
def __call__(self, infer_requests: List[Union['InferRequest', Dict]], solution: List[str], **kwargs) -> List[float]:
|
||||
rewards = []
|
||||
if not isinstance(infer_requests[0], str):
|
||||
predictions = [request['messages'][-1]['content'] for request in infer_requests]
|
||||
else:
|
||||
predictions = infer_requests
|
||||
for prediction, ground_truth in zip(predictions, solution):
|
||||
if prediction.endswith('Observation:'):
|
||||
prediction = prediction[:prediction.index('Observation:')].strip()
|
||||
action_ref = []
|
||||
action_input_ref = []
|
||||
action_pred = []
|
||||
action_input_pred = []
|
||||
reference = ground_truth
|
||||
prediction = prediction.replace('<|endoftext|>', '').replace('<|im_end|>', '').strip()
|
||||
ref_action, ref_input = ReactORM.parse_output(reference)
|
||||
pred_action, pred_input = ReactORM.parse_output(prediction)
|
||||
action_ref.append(ref_action)
|
||||
action_input_ref.append(ref_input)
|
||||
if pred_action is None:
|
||||
action_pred.append('none')
|
||||
else:
|
||||
action_pred.append(pred_action)
|
||||
|
||||
if pred_input is None:
|
||||
action_input_pred.append('{}')
|
||||
else:
|
||||
action_input_pred.append(pred_input)
|
||||
|
||||
reward = ReactORM.evaluate_action_reward(action_pred, action_ref, action_input_pred, action_input_ref)
|
||||
rewards.append(float(reward))
|
||||
return rewards
|
||||
|
||||
@staticmethod
|
||||
def evaluate_rougel(cand_list: list, ref_list: list):
|
||||
if len(ref_list) == 0:
|
||||
return None
|
||||
try:
|
||||
from rouge import Rouge
|
||||
rouge = Rouge()
|
||||
rouge_score = rouge.get_scores(hyps=cand_list, refs=ref_list, avg=True)
|
||||
rougel = rouge_score['rouge-l']['f']
|
||||
return rougel
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
class MathORM(ORM):
|
||||
|
||||
def __init__(self, args=None, **kwargs):
|
||||
super().__init__(args)
|
||||
from transformers.utils import strtobool
|
||||
self.use_opencompass = strtobool(os.environ.get('USE_OPENCOMPASS_EVALUATOR', 'False'))
|
||||
if self.use_opencompass:
|
||||
from opencompass.datasets.math import MATHEvaluator
|
||||
self.evaluator = MATHEvaluator()
|
||||
|
||||
@staticmethod
|
||||
def check_terminate(answers: Union[str, List[str]]) -> List[bool]:
|
||||
if isinstance(answers, str):
|
||||
answers = [answers]
|
||||
results = []
|
||||
for answer in answers:
|
||||
results.append('\\boxed' in answer)
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def extract_boxed_result(text):
|
||||
pattern = r'\\boxed{([^}]*)}'
|
||||
match = re.search(pattern, text)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
else:
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def clean_latex(latex_str):
|
||||
latex_str = re.sub(r'\\\(|\\\)|\\\[|\\]', '', latex_str)
|
||||
latex_str = latex_str.replace('}}', '}').replace('{', '').replace('}', '')
|
||||
return latex_str.strip()
|
||||
|
||||
@staticmethod
|
||||
def parse_expression(latex_str):
|
||||
from sympy import simplify
|
||||
from sympy.parsing.latex import parse_latex
|
||||
try:
|
||||
expr = parse_latex(latex_str)
|
||||
return simplify(expr)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def compare_consecutive(first, second):
|
||||
cleaned_list = [MathORM.clean_latex(latex) for latex in [first, second]]
|
||||
parsed_exprs = [MathORM.parse_expression(latex) for latex in cleaned_list]
|
||||
if hasattr(parsed_exprs[0], 'equals') and hasattr(parsed_exprs[1], 'equals'):
|
||||
value = parsed_exprs[0].equals(parsed_exprs[1])
|
||||
else:
|
||||
value = parsed_exprs[0] == parsed_exprs[1]
|
||||
if value is None:
|
||||
value = False
|
||||
return value
|
||||
|
||||
def __call__(self, infer_requests: List[Union['InferRequest', Dict]], ground_truths: List[str],
|
||||
**kwargs) -> List[float]:
|
||||
rewards = []
|
||||
predictions = [request.messages[-1]['content'] for request in infer_requests]
|
||||
for prediction, ground_truth in zip(predictions, ground_truths):
|
||||
if '# Answer' in prediction:
|
||||
prediction = prediction.split('# Answer')[1]
|
||||
if '# Answer' in ground_truth:
|
||||
ground_truth = ground_truth.split('# Answer')[1]
|
||||
prediction = prediction.strip()
|
||||
ground_truth = ground_truth.strip()
|
||||
prediction = MathORM.extract_boxed_result(prediction)
|
||||
ground_truth = MathORM.extract_boxed_result(ground_truth)
|
||||
if self.use_opencompass:
|
||||
reward = self.evaluator.is_equiv(prediction, ground_truth)
|
||||
else:
|
||||
reward = MathORM.compare_consecutive(prediction, ground_truth)
|
||||
rewards.append(float(reward))
|
||||
return rewards
|
||||
|
||||
|
||||
orms = {
|
||||
'toolbench': ReactORM,
|
||||
'math': MathORM,
|
||||
'accuracy': MathAccuracy,
|
||||
'format': Format,
|
||||
'react_format': ReActFormat,
|
||||
'cosine': CosineReward,
|
||||
'repetition': RepetitionPenalty,
|
||||
'soft_overlong': SoftOverlong,
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
# Process Reward Model (PRM) implementations for sampling.
|
||||
# GRPO training is not currently supported.
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
from swift.infer_engine import InferClient, InferRequest
|
||||
|
||||
|
||||
class PRM:
|
||||
|
||||
def __call__(self, **kwargs) -> List[Any]:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
SYSTEM = """
|
||||
You are a process reward model, give the reward value of the answer, you must follow the instructions below:
|
||||
|
||||
1. Output a float reward value between -1.0 and 1.0, -1.0 means the worst answer, 1.0 means the best answer, please think step by step to give your reasons and thoughts, but the reward must appare at the end with this format: **Reward: your-reward-value**.
|
||||
|
||||
2. The answer may be incomplete, you must give the reward by the existing part of the answer, taking into account semantic coherence, logical correctness, and clarity.
|
||||
|
||||
3. A ground truth answer will be given to you, it may be not the best one, consider it as a reference example.
|
||||
|
||||
Begin!
|
||||
""" # noqa
|
||||
|
||||
QUERY = """
|
||||
The original question or the previous conversation:
|
||||
|
||||
#query#
|
||||
|
||||
Here is the ground truth as the reference:
|
||||
|
||||
#ground_truth#
|
||||
|
||||
Given the upper information, give your reward(-1.0~1.0) of the following answer:
|
||||
|
||||
#response#
|
||||
"""
|
||||
|
||||
|
||||
class QwenMaxPRM(PRM):
|
||||
|
||||
def __call__(self, infer_requests: List[Union['InferRequest', Dict]], ground_truths: List[str],
|
||||
**kwargs) -> List[float]:
|
||||
# TODO: check request_config
|
||||
rewards = []
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv('DASHSCOPE_API_KEY'),
|
||||
base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
|
||||
)
|
||||
|
||||
for request, ground_truth in zip(infer_requests, ground_truths):
|
||||
previous = request.messages[:-1]
|
||||
if previous[0]['role'] == 'system':
|
||||
previous = previous[1:]
|
||||
|
||||
assert request.messages[-1]['role'] == 'assistant'
|
||||
query = QUERY.replace('#query#', json.dumps(previous))
|
||||
query = query.replace('#ground_truth#', ground_truth)
|
||||
query = query.replace('#response#', request.messages[-1]['content'])
|
||||
messages = [
|
||||
{
|
||||
'role': 'system',
|
||||
'content': SYSTEM
|
||||
},
|
||||
{
|
||||
'role': 'user',
|
||||
'content': query
|
||||
},
|
||||
]
|
||||
completion = client.chat.completions.create(
|
||||
model='qwen-max',
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
content = completion.choices[0].message.content
|
||||
if 'Reward:' not in content:
|
||||
rewards.append(0.)
|
||||
else:
|
||||
try:
|
||||
reward = float(content.split('Reward:')[1].strip().replace('*', ''))
|
||||
rewards.append(reward)
|
||||
except Exception:
|
||||
rewards.append(0.)
|
||||
|
||||
return rewards
|
||||
|
||||
|
||||
class ClientPRM(PRM):
|
||||
|
||||
def __init__(self, api_key=None, base_url=None, model=None):
|
||||
import os
|
||||
if api_key is None:
|
||||
api_key = os.getenv('DASHSCOPE_API_KEY')
|
||||
if base_url is None:
|
||||
base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
|
||||
if model is None:
|
||||
model = 'qwen-plus'
|
||||
self.infer_engine = InferClient(base_url=base_url, api_key=api_key)
|
||||
self.infer_engine.strict = False
|
||||
self.infer_kwargs = {
|
||||
'model': model,
|
||||
}
|
||||
|
||||
def __call__(self, infer_requests: List[Union['InferRequest', Dict]], ground_truths: List[str],
|
||||
**kwargs) -> List[float]:
|
||||
prm_infer_requests = []
|
||||
request_config = kwargs.get('request_config')
|
||||
for request, ground_truth in zip(infer_requests, ground_truths):
|
||||
previous = request['messages'][:-1]
|
||||
if previous[0]['role'] == 'system':
|
||||
previous = previous[1:]
|
||||
|
||||
assert request['messages'][-1]['role'] == 'assistant'
|
||||
query = QUERY.replace('#query#', json.dumps(previous))
|
||||
query = query.replace('#ground_truth#', ground_truth)
|
||||
query = query.replace('#response#', request['messages'][-1]['content'])
|
||||
messages = [
|
||||
{
|
||||
'role': 'system',
|
||||
'content': SYSTEM
|
||||
},
|
||||
{
|
||||
'role': 'user',
|
||||
'content': query
|
||||
},
|
||||
]
|
||||
|
||||
prm_infer_requests.append(InferRequest(messages=messages))
|
||||
|
||||
responses = self.infer_engine.infer(prm_infer_requests, request_config=request_config, **self.infer_kwargs)
|
||||
rewards = []
|
||||
for response in responses:
|
||||
content = response.choices[0].message.content
|
||||
if 'Reward:' not in content:
|
||||
rewards.append(0.)
|
||||
else:
|
||||
try:
|
||||
reward = float(content.split('Reward:')[1].strip().replace('*', ''))
|
||||
rewards.append(reward)
|
||||
except Exception:
|
||||
rewards.append(0.)
|
||||
return rewards
|
||||
|
||||
|
||||
prms = {
|
||||
'qwen_max': QwenMaxPRM,
|
||||
'client': ClientPRM,
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
# Reward Model Plugin implementations for GRPO training.
|
||||
# This module provides plugins for integrating external reward models,
|
||||
# including both discriminative reward models (with value heads) and
|
||||
# generative reward models (LLM-as-judge style).
|
||||
|
||||
import re
|
||||
import textwrap
|
||||
import torch
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Dict, List
|
||||
|
||||
from swift.infer_engine import ChatCompletionResponse, RequestConfig, TransformersEngine
|
||||
from swift.template import Template
|
||||
from swift.utils import get_logger, to_device
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class DefaultRMPlugin:
|
||||
"""
|
||||
Default Reward Model Plugin
|
||||
|
||||
This class implements the default processing logic for reward models.
|
||||
It assumes that `self.model` is a classification model with a value head(output dimmension 1).
|
||||
The first logits value from the model's output is used as the reward score.
|
||||
"""
|
||||
|
||||
def __init__(self, model, template):
|
||||
self.model = model
|
||||
self.template: Template = template
|
||||
|
||||
def __call__(self, inputs, **kwargs):
|
||||
batched_inputs = [self.template.encode(deepcopy(infer_request)) for infer_request in inputs]
|
||||
reward_inputs = to_device(self.template.data_collator(batched_inputs), self.model.device)
|
||||
|
||||
with torch.inference_mode():
|
||||
return self.model(**reward_inputs).logits[:, 0]
|
||||
|
||||
|
||||
class GenRMPlugin(DefaultRMPlugin):
|
||||
|
||||
def __init__(self, model, template):
|
||||
"""
|
||||
Generative Reward Model Plugin Example.
|
||||
|
||||
This method sets up the reward model plugin by initializing the TransformersEngine for efficient inference,
|
||||
configuring the request parameters, and defining the system prompt that guides the reward model in
|
||||
evaluating responses.
|
||||
|
||||
Args:
|
||||
model (torch.nn.Module): The generative reward model.
|
||||
template (Template): The template used for encoding input data.
|
||||
"""
|
||||
|
||||
super().__init__(model, template)
|
||||
# initilize TransformersEngine to infer
|
||||
self.engine = TransformersEngine(self.model, template=self.template, max_batch_size=0) # 0: no limit
|
||||
self.request_config = RequestConfig() # customise your request config here
|
||||
self.system = textwrap.dedent("""
|
||||
Based on the dialogue history, analyze in detail whether the model's response is accurate, complete, and relevant.
|
||||
Assign a reward score between 0 and 1, where 0 indicates completely incorrect and 1 indicates fully correct.
|
||||
Before finishing your response, please assign a reward using the following format:
|
||||
|
||||
Reward: {reward}
|
||||
|
||||
For example:
|
||||
Reward: 0.85
|
||||
""") # noqa
|
||||
|
||||
def __call__(self, inputs, **kwargs):
|
||||
"""
|
||||
Compute reward scores for the provided inputs.
|
||||
|
||||
This method processes each input by converting dialogue messages into a query, sending the query to the
|
||||
reward model for inference, and extracting the reward scores from the model's responses. The final reward
|
||||
for each input is the average of all extracted scores.
|
||||
Args:
|
||||
inputs (List[Dict]): A list of input requests. Each input request is a dictionary containing:
|
||||
- 'messages' (List[Dict]): messages from the training model. Each message dictionary includes:
|
||||
- 'role' (str): The role of the speaker (e.g., 'user', 'assistant').
|
||||
- 'content' (str): The content of the message.
|
||||
- Additional dataset columns as key-value pairs (e.g., 'solutions', 'images').
|
||||
Returns:
|
||||
torch.Tensor: A tensor containing the average reward scores for each input. The tensor has a shape of (N,),
|
||||
where N is the number of input requests.
|
||||
"""
|
||||
|
||||
rm_inputs = self.prepare_rm_inputs(inputs)
|
||||
results = self.engine.infer(rm_inputs, self.request_config, use_tqdm=False)
|
||||
rewards = self.compute_rewards(results)
|
||||
return torch.tensor(rewards, dtype=torch.float32)
|
||||
|
||||
def prepare_rm_inputs(self, inputs: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
Prepare inputs for the reward model by converting messages into queries.
|
||||
|
||||
Args:
|
||||
inputs (List[Dict]): A list of input requests.
|
||||
|
||||
Returns:
|
||||
List[Dict]: Processed inputs for the reward model.
|
||||
"""
|
||||
rm_inputs = []
|
||||
for idx, infer_request in enumerate(inputs):
|
||||
# Deep copy to prevent modification of original input
|
||||
rm_infer_request = deepcopy(infer_request)
|
||||
|
||||
# Extract and convert messages to a single query string
|
||||
messages = rm_infer_request.get('messages')
|
||||
query = self.messages_to_query(messages)
|
||||
|
||||
# Construct new messages tailored for the reward model
|
||||
rm_messages = [{'role': 'system', 'content': self.system}, {'role': 'user', 'content': query}]
|
||||
|
||||
# Update the messages in the reward infer request
|
||||
rm_infer_request['messages'] = rm_messages
|
||||
rm_inputs.append(rm_infer_request)
|
||||
return rm_inputs
|
||||
|
||||
@staticmethod
|
||||
def extract_reward(model_output: str) -> float:
|
||||
"""
|
||||
Extract the reward score from the model's output.
|
||||
|
||||
Args:
|
||||
model_output (str): The model's output string, expected to follow the format "Reward: {reward}".
|
||||
|
||||
Returns:
|
||||
float: The extracted reward score.
|
||||
|
||||
Raises:
|
||||
ValueError: If the reward score cannot be extracted or the format is incorrect.
|
||||
"""
|
||||
match = re.search(r'Reward:\s*([0-1](?:\.\d+)?)', model_output)
|
||||
if match:
|
||||
return float(match.group(1))
|
||||
else:
|
||||
logger.warning("Unable to extract reward score from the model's output, set reward to 0")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def messages_to_query(messages):
|
||||
"""
|
||||
Compress a list of message dictionaries into a single query string.
|
||||
|
||||
Args:
|
||||
messages (list[dict]): A list of message dictionaries, each containing:
|
||||
- 'role' (str): The role of the speaker (e.g., 'user', 'assistant').
|
||||
- 'content' (str): The content of the message.
|
||||
|
||||
Returns:
|
||||
str: A single string that concatenates all messages in a formatted manner.
|
||||
|
||||
Example:
|
||||
>>> messages = [
|
||||
... {'role': 'user', 'content': 'Hello, how are you?'},
|
||||
... {'role': 'assistant', 'content': 'I am fine, thank you! How can I assist you today?'},
|
||||
... {'role': 'user', 'content': 'Can you help me with my homework?'}
|
||||
... ]
|
||||
>>> print(messages_to_query(messages))
|
||||
User: Hello, how are you?
|
||||
Assistant: I am fine, thank you! How can I assist you today?
|
||||
User: Can you help me with my homework?
|
||||
"""
|
||||
# Initialize an empty list to hold formatted messages
|
||||
formatted_messages = []
|
||||
|
||||
# Define a mapping for role capitalization if needed
|
||||
role_mapping = {
|
||||
'user': 'User',
|
||||
'assistant': 'Assistant',
|
||||
'system': 'System'
|
||||
# Add more roles here as needed
|
||||
}
|
||||
|
||||
for idx, message in enumerate(messages):
|
||||
if not isinstance(message, dict):
|
||||
raise TypeError(f'Each message must be a dictionary. Found {type(message)} at index {idx}.')
|
||||
|
||||
# Extract 'role' and 'content' from each message
|
||||
role = message.get('role')
|
||||
content = message.get('content')
|
||||
if not content:
|
||||
continue
|
||||
|
||||
# Capitalize the role using the mapping, default to capitalized original role
|
||||
role_formatted = role_mapping.get(role.lower(), role.capitalize())
|
||||
|
||||
# Append the formatted message to the list
|
||||
formatted_messages.append(f'{role_formatted}: {content}')
|
||||
|
||||
# Join all formatted messages with newline characters
|
||||
query = '\n'.join(formatted_messages)
|
||||
|
||||
return query
|
||||
|
||||
def compute_rewards(self, results: List['ChatCompletionResponse']) -> List[float]:
|
||||
"""
|
||||
Compute average reward scores from the reward model's outputs.
|
||||
|
||||
Args:
|
||||
results (List['ChatCompletionResponse']): A list of results from the reward model.
|
||||
|
||||
Returns:
|
||||
List[float]: A list of average reward scores.
|
||||
"""
|
||||
rewards = []
|
||||
for idx, output in enumerate(results):
|
||||
try:
|
||||
cur_rewards = []
|
||||
for choice in output.choices:
|
||||
response = choice.message.content
|
||||
reward = self.extract_reward(response)
|
||||
cur_rewards.append(reward)
|
||||
cur_rewards = [r for r in cur_rewards if r is not None]
|
||||
if cur_rewards:
|
||||
average_reward = sum(cur_rewards) / len(cur_rewards)
|
||||
else:
|
||||
average_reward = 0.0
|
||||
logger.warning('No valid rewards extracted. Assigning reward score of 0.0.')
|
||||
|
||||
rewards.append(average_reward)
|
||||
except Exception as e:
|
||||
logger.error(f'Error computing reward: {e}')
|
||||
rewards.append(0.0) # Assign default reward score on failure
|
||||
return rewards
|
||||
|
||||
|
||||
rm_plugins = {
|
||||
'default': DefaultRMPlugin,
|
||||
'genrm': GenRMPlugin,
|
||||
}
|
||||
Reference in New Issue
Block a user