From 657ce5cab9246be8f11dba1b89efca442856ffad Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 11 Feb 2025 03:23:05 -0800
Subject: [PATCH] Update rl.py

---
 unsloth/models/rl.py | 48 ++++++++++++++++++++++----------------------
 1 file changed, 24 insertions(+), 24 deletions(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 225e0e48f1..9b8b410f46 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -349,31 +349,31 @@ def patch_vllm(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports):
         init,
         flags = re.MULTILINE | re.DOTALL,
     )
-    if len(vllm_part) != 1: return None
+    if len(vllm_part) == 1:
+        vllm_part, args = vllm_part[0][0], vllm_part[0][1]
+        # Strip all comments
+        new_vllm_part = re.sub(r"\#[^\n]{1,}\n", "", vllm_part)
 
-    vllm_part, args = vllm_part[0][0], vllm_part[0][1]
-    # Strip all comments
-    new_vllm_part = re.sub(r"\#[^\n]{1,}\n", "", vllm_part)
-
-    # Get SamplingParams
-    sampling_params = re.findall(
-        r"\n[\s]{4,}(self\.[^\s]{1,}[\s]{0,}\=[\s]{0,}"\
-        r"SamplingParams\(.+?\))",
-        new_vllm_part,
-        flags = re.MULTILINE | re.DOTALL,
-    )
-    print(len(sampling_params), RLTrainer_name)
-    if len(sampling_params) == 1:
-        sampling_params = sampling_params[0]
-        # Replace with our vLLM engine
-        sampling_params = \
-            " "*12 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \
-            sampling_params # Add spaces
-        new_vllm_part = \
-            f"\n{' '*8}if {args}.use_vllm:\n{sampling_params} "\
-            f"if getattr(args, 'sampling_params', None) is None else "\
-            f"getattr(args, 'sampling_params', None)\n{' '*8}else:\n"
-        init = init.replace(vllm_part, new_vllm_part)
+        # Get SamplingParams
+        sampling_params = re.findall(
+            r"\n[\s]{4,}(self\.[^\s]{1,}[\s]{0,}\=[\s]{0,}"\
+            r"SamplingParams\(.+?\))",
+            new_vllm_part,
+            flags = re.MULTILINE | re.DOTALL,
+        )
+        print(sampling_params)
+        if len(sampling_params) == 1:
+            sampling_params = sampling_params[0]
+            # Replace with our vLLM engine
+            sampling_params = \
+                " "*12 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \
+                sampling_params # Add spaces
+            new_vllm_part = \
+                f"\n{' '*8}if {args}.use_vllm:\n{sampling_params} "\
+                f"if getattr(args, 'sampling_params', None) is None else "\
+                f"getattr(args, 'sampling_params', None)\n{' '*8}else:\n"
+            init = init.replace(vllm_part, new_vllm_part)
+        pass
     pass
 
     # Search for vLLM calling in all child functions