Fix prm in sampler (#3184)

modelscope · Feb 19, 2025 · 70e19b4 · 70e19b4
1 parent 0d0aa03
commit 70e19b4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/swift/llm/infer/infer_engine/infer_engine.py b/swift/llm/infer/infer_engine/infer_engine.py
@@ -122,6 +122,9 @@ def _update_metrics(result, metrics: Optional[List[Metric]] = None):
                 metric.update(response)
         return result_origin
 
+    def __call__(self, *args, **kwargs):
+        return self.infer(*args, **kwargs)
+
     def infer(self,
               infer_requests: List[InferRequest],
               request_config: Optional[RequestConfig] = None,

diff --git a/swift/llm/sampling/utils.py b/swift/llm/sampling/utils.py
@@ -43,6 +43,9 @@ def get_reward(model: Any,
     if 'ground_truths' in parameters:
         gt_param = {'ground_truths': ground_truths}
     rewards = model(infer_requests, request_config=request_config, **gt_param)
+    from swift.llm.infer.protocol import ChatCompletionResponse
+    if isinstance(rewards[0], ChatCompletionResponse):
+        rewards = [float(r.choices[0].message.content) for r in rewards]
     arr = []
     for reward in rewards:
         if isinstance(reward, (list, tuple)):