From 9899c924f97e4e79ef4de24fa9b54180b82df5fa Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Sun, 25 May 2025 17:00:12 -0400
Subject: [PATCH] suport sampling params/max new tokens

---
 src/axolotl/integrations/kd/collator_online_teacher.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/axolotl/integrations/kd/collator_online_teacher.py b/src/axolotl/integrations/kd/collator_online_teacher.py
index a43936335..4593401a3 100644
--- a/src/axolotl/integrations/kd/collator_online_teacher.py
+++ b/src/axolotl/integrations/kd/collator_online_teacher.py
@@ -231,6 +231,9 @@ class OnlineTeacherCollator(KDBatchSamplerDataCollatorForSeq2Seq):
             "max_new_tokens": 0,
             "skip_special_tokens": False,
             "temperature": self.kd_temperature,
+            "sampling_params": {
+                "max_tokens": 0,
+            },
         }
 
         # Initialize with empty lists, so if API call fails, these are returned.