ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

PatchRL

Browse source
This commit is contained in:
Daniel Han 2025-02-05 05:17:38 -08:00
commit eea3fe103d
2 changed files with 6 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/models/__init__.py
View file

@ -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

9
unsloth/models/rl.py
View file

@ -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)