Fix kto (modelscope#1737)

* fix * fix
kronoszhang · Aug 17, 2024 · 3081069 · 3081069
1 parent 250d0d8
commit 3081069
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/swift/trainers/kto_trainer.py b/swift/trainers/kto_trainer.py
@@ -75,7 +75,8 @@ def encode_batch(batch: Dict[str, List[Any]], template: Template):
                                       response=r,
                                       round0=i)
     template._concat_context_list(template.prompt, res_context_list, compute_loss_idx, query=query, round0=len(history))
-    res_context_list, compute_loss_idx = template._simplify_context_list(res_context_list, compute_loss_idx)
+    res_context_list, compute_loss_idx = template._simplify_context_list(
+        res_context_list, compute_loss_idx, example=batch)
     prompt = ''.join(res_context_list)
 
     return {'prompt': prompt, 'completion': batch['response'], 'label': batch['label']}