opt reward model

src/rlhf/reward.py

@@ -31,11 +31,7 @@ def loss_function(prefer_reward, alter_reward):
 
 @beartype
 class RewardModel(pl.LightningModule):
-    def __init__(
-        self,
-        args,
-        rwkv: RWKV
-    ):
+    def __init__(self, args):
         super().__init__()
 
         # 加载 RWKV 模型

train_rm.py

@@ -230,7 +230,6 @@ if __name__ == "__main__":
     from src.trainer import train_callback
 
     from src.rlhf.reward import RewardModel
-    from src.model import RWKV
     from src.dataset import RMDataset
 
     # 读入训练数据