From 407d707ced89b66825ef8039ae5b5c527933aa88 Mon Sep 17 00:00:00 2001 From: Edd <68678137+Erland366@users.noreply.github.com> Date: Wed, 27 Nov 2024 02:22:02 +0400 Subject: [PATCH] Feat/kto (#1316) * Add PatchKTOTrainer and update model imports * Update dpo.py * Update __init__.py * Delete unsloth/models/kto.py --------- Co-authored-by: Daniel Han --- unsloth/models/__init__.py | 2 +- unsloth/models/dpo.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index 3230cdc207..9c032b28bd 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -16,5 +16,5 @@ from .loader import FastLanguageModel, FastVisionModel from .llama import FastLlamaModel from .mistral import FastMistralModel from .qwen2 import FastQwen2Model -from .dpo import PatchDPOTrainer +from .dpo import PatchDPOTrainer, PatchKTOTrainer from ._utils import is_bfloat16_supported diff --git a/unsloth/models/dpo.py b/unsloth/models/dpo.py index e7074350c3..5dc71f920a 100644 --- a/unsloth/models/dpo.py +++ b/unsloth/models/dpo.py @@ -14,6 +14,7 @@ __all__ = [ "PatchDPOTrainer", + "PatchKTOTrainer", ] try: @@ -127,4 +128,4 @@ def PatchDPOTrainer(): pass pass pass - +PatchKTOTrainer = PatchDPOTrainer