add kto

Former-commit-id: ec51986cf70b0bdd79b8141e45916670fb97a08e
2024-05-17 13:09:17 +08:00
parent 92b3697e2c
commit 66b5634ebf
12 changed files with 452 additions and 8 deletions
--- a/src/llamafactory/data/init.py
+++ b/src/llamafactory/data/init.py
@@ -1,4 +1,4 @@
-from .collator import PairwiseDataCollatorWithPadding
+from .collator import PairwiseDataCollatorWithPadding,KTODataCollatorWithPadding
 from .loader import get_dataset
 from .template import Template, get_template_and_fix_tokenizer, templates
 from .utils import Role, split_dataset
@@ -6,6 +6,7 @@ from .utils import Role, split_dataset

 __all__ = [
    "PairwiseDataCollatorWithPadding",
+    "KTODataCollatorWithPadding",
    "get_dataset",
    "Template",
    "get_template_and_fix_tokenizer",