From eea3fe103d4ed7cdec23f3f0ca853adec38a0de6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 5 Feb 2025 05:17:38 -0800 Subject: [PATCH] PatchRL --- unsloth/models/__init__.py | 1 + unsloth/models/rl.py | 9 +++++---- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index c52d14f402..3478dfc31a 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -20,3 +20,4 @@ from .mistral import FastMistralModel from .qwen2 import FastQwen2Model from .dpo import PatchDPOTrainer, PatchKTOTrainer from ._utils import is_bfloat16_supported +from .rl import PatchRL diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index efe2d33e01..2aa8f02659 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -13,14 +13,15 @@ # limitations under the License. __all__ = [ - "patch_rl", + "PatchRL", ] -from trl.models.utils import unwrap_model_for_generation -from contextlib import contextmanager +def PatchRL(FastLanguageModel): -def patch_rl(FastLanguageModel): + from trl.models.utils import unwrap_model_for_generation + from contextlib import contextmanager + @contextmanager def unsloth_unwrap_model_for_generation(model, *args, **kwargs): FastLanguageModel.for_inference(model)