wehub-resource-sync
210 lines
7.6 KiB
Python
210 lines
7.6 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import gradio as gr
|
|
from functools import partial
|
|
from typing import Type
|
|
|
|
from swift.model import ModelType, get_model_list
|
|
from ..base import BaseUI
|
|
|
|
|
|
class RLHF(BaseUI):
|
|
|
|
group = 'llm_rlhf'
|
|
|
|
locale_dict = {
|
|
'rlhf_tab': {
|
|
'label': {
|
|
'zh': '对齐参数设置',
|
|
'en': 'Alignment params settings'
|
|
},
|
|
},
|
|
'ref_model': {
|
|
'label': {
|
|
'zh': 'Ref模型id或路径',
|
|
'en': 'Ref model id or path'
|
|
},
|
|
'info': {
|
|
'zh': '实际的模型id或路径',
|
|
'en': 'The actual model id or path'
|
|
}
|
|
},
|
|
'ref_model_type': {
|
|
'label': {
|
|
'zh': 'Ref模型类型',
|
|
'en': 'Ref model type'
|
|
},
|
|
'info': {
|
|
'zh': 'SWIFT已支持的模型类型',
|
|
'en': 'Model type supported by SWIFT'
|
|
}
|
|
},
|
|
'reward_model': {
|
|
'label': {
|
|
'zh': '奖励模型id或路径',
|
|
'en': 'Reward model id or path'
|
|
},
|
|
'info': {
|
|
'zh': '实际的模型id或路径',
|
|
'en': 'The actual model id or path'
|
|
}
|
|
},
|
|
'reward_model_type': {
|
|
'label': {
|
|
'zh': '奖励模型类型',
|
|
'en': 'Reward model type'
|
|
},
|
|
'info': {
|
|
'zh': 'SWIFT已支持的模型类型',
|
|
'en': 'Model type supported by SWIFT'
|
|
}
|
|
},
|
|
'teacher_model': {
|
|
'label': {
|
|
'zh': '教师模型id或路径',
|
|
'en': 'Teacher model id or path'
|
|
},
|
|
'info': {
|
|
'zh': '实际的模型id或路径',
|
|
'en': 'The actual model id or path'
|
|
}
|
|
},
|
|
'teacher_model_type': {
|
|
'label': {
|
|
'zh': '教师模型类型',
|
|
'en': 'Teacher model type'
|
|
},
|
|
'info': {
|
|
'zh': 'SWIFT已支持的模型类型',
|
|
'en': 'Model type supported by SWIFT'
|
|
}
|
|
},
|
|
'beta': {
|
|
'label': {
|
|
'zh': 'KL正则项系数',
|
|
'en': 'KL regression ratio'
|
|
},
|
|
},
|
|
'max_completion_length': {
|
|
'label': {
|
|
'zh': '最大生成长度',
|
|
'en': 'Max completion length'
|
|
},
|
|
},
|
|
'loss_scale': {
|
|
'label': {
|
|
'zh': '损失权重设置',
|
|
'en': 'Loss weights setting'
|
|
},
|
|
},
|
|
'lmbda': {
|
|
'label': {
|
|
'zh': 'GKD学生数据比例',
|
|
'en': 'GKD student data ratio'
|
|
},
|
|
},
|
|
'cpo_alpha': {
|
|
'label': {
|
|
'zh': 'CPO/SimPO中NLL损失系数',
|
|
'en': 'CPO/SimPO NLL loss coefficient'
|
|
},
|
|
},
|
|
'rpo_alpha': {
|
|
'label': {
|
|
'zh': 'DPO中混合sft交叉熵的系数',
|
|
'en': 'DPO Cross Entropy ratio'
|
|
},
|
|
},
|
|
'simpo_gamma': {
|
|
'label': {
|
|
'zh': 'SimPO reward margin',
|
|
'en': 'SimPO reward margin'
|
|
},
|
|
},
|
|
'desirable_weight': {
|
|
'label': {
|
|
'zh': 'KTO符合项系数',
|
|
'en': 'KTO desirable ratio'
|
|
},
|
|
},
|
|
'undesirable_weight': {
|
|
'label': {
|
|
'zh': 'KTO不符合项系数',
|
|
'en': 'KTO undesirable ratio'
|
|
},
|
|
}
|
|
}
|
|
|
|
rlhf_args_dict = {
|
|
'dpo': ['rpo_alpha', 'ref_model', 'ref_model_type'],
|
|
'cpo': ['cpo_alpha'],
|
|
'kto': ['desirable_weight', 'undesirable_weight', 'ref_model', 'ref_model_type'],
|
|
'simpo': ['simpo_gamma', 'cpo_alpha'],
|
|
'gkd': ['teacher_model', 'teacher_model_type', 'max_completion_length', 'lmbda'],
|
|
'ppo': ['reward_model', 'reward_model_type', 'max_completion_length', 'ref_model', 'ref_model_type']
|
|
}
|
|
|
|
@classmethod
|
|
def do_build_ui(cls, base_tab: Type['BaseUI']):
|
|
with gr.Accordion(elem_id='rlhf_tab', open=False):
|
|
with gr.Blocks():
|
|
with gr.Row():
|
|
gr.Slider(elem_id='beta', minimum=0., maximum=5.0, step=0.1, value=0.1, scale=10)
|
|
gr.Slider(elem_id='rpo_alpha', minimum=0., maximum=2, step=0.1, scale=10)
|
|
gr.Slider(elem_id='lmbda', minimum=0., maximum=1.0, step=0.1, scale=10)
|
|
gr.Slider(elem_id='simpo_gamma', minimum=0., maximum=2.0, step=0.1, scale=10)
|
|
gr.Slider(elem_id='desirable_weight', minimum=0., maximum=2.0, step=0.1, scale=10)
|
|
gr.Slider(elem_id='undesirable_weight', minimum=0., maximum=2.0, step=0.1, scale=10)
|
|
with gr.Row():
|
|
gr.Textbox(elem_id='max_completion_length', scale=10)
|
|
gr.Textbox(elem_id='loss_scale', scale=10)
|
|
gr.Slider(elem_id='cpo_alpha', minimum=0., maximum=1, step=0.1, scale=10)
|
|
gr.Dropdown(
|
|
elem_id='teacher_model',
|
|
scale=20,
|
|
value=None,
|
|
choices=get_model_list(),
|
|
allow_custom_value=True)
|
|
gr.Dropdown(
|
|
elem_id='teacher_model_type',
|
|
choices=ModelType.get_model_name_list(),
|
|
value=None,
|
|
scale=10,
|
|
allow_custom_value=True)
|
|
with gr.Row():
|
|
gr.Dropdown(
|
|
elem_id='ref_model', scale=20, value=None, choices=get_model_list(), allow_custom_value=True)
|
|
gr.Dropdown(
|
|
elem_id='ref_model_type',
|
|
choices=ModelType.get_model_name_list(),
|
|
value=None,
|
|
scale=10,
|
|
allow_custom_value=True)
|
|
gr.Dropdown(
|
|
elem_id='reward_model', scale=20, value=None, choices=get_model_list(), allow_custom_value=True)
|
|
gr.Dropdown(
|
|
elem_id='reward_model_type',
|
|
choices=ModelType.get_model_name_list(),
|
|
value=None,
|
|
scale=10,
|
|
allow_custom_value=True)
|
|
|
|
@classmethod
|
|
def after_build_ui(cls, base_tab: Type['BaseUI']):
|
|
cls.element('ref_model').change(
|
|
partial(cls.update_input_model, allow_keys=['ref_model_type'], has_record=False, is_ref_model=True),
|
|
inputs=[cls.element('ref_model')],
|
|
outputs=[cls.element('ref_model_type')])
|
|
cls.element('reward_model').change(
|
|
partial(cls.update_input_model, allow_keys=['reward_model_type'], has_record=False, is_ref_model=True),
|
|
inputs=[cls.element('reward_model')],
|
|
outputs=[cls.element('reward_model_type')])
|
|
cls.element('teacher_model').change(
|
|
partial(cls.update_input_model, allow_keys=['teacher_model_type'], has_record=False, is_ref_model=True),
|
|
inputs=[cls.element('teacher_model')],
|
|
outputs=[cls.element('teacher_model_type')])
|
|
|
|
@staticmethod
|
|
def update_beta(rlhf_type):
|
|
beta_value_dict = {'simpo': 2., 'gkd': 0.5, 'grpo': 0.04}
|
|
return beta_value_dict.get(rlhf_type, 0.1) if rlhf_type else 0.1
|