File size: 5,106 Bytes
7feac49 |
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
import os
from typing import Any, Dict, List, Union
import json
from swift.llm import 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):
from swift.llm import InferClient
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,
}
|