From 3cffe881bb49b67e5edc75ea7276261cf01878cc Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Wed, 28 May 2025 10:01:25 -0400
Subject: [PATCH] accept compressed responses for smaller wire payload

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

diff --git a/src/axolotl/integrations/kd/collator_online_teacher.py b/src/axolotl/integrations/kd/collator_online_teacher.py
index 6a6a4e2ee..ecd0a0c01 100644
--- a/src/axolotl/integrations/kd/collator_online_teacher.py
+++ b/src/axolotl/integrations/kd/collator_online_teacher.py
@@ -336,11 +336,12 @@ class OnlineTeacherCollator(KDBatchSamplerDataCollatorForSeq2Seq):
         ret_data_target_mask: List[List[List[int]]] = []
 
         try:
+            headers = {"Accept-Encoding": "deflate, gzip, br, zstd"}
             response = self.http_session.post(
                 api_endpoint,
                 json=payload,
+                headers=headers,
                 timeout=self.kd_online_timeout,
-                # json_decoder=orjson.loads,
             )
             response.raise_for_status()
             api_data: dict = orjson.loads(response.content)