Linting fix

tests/test_trainer_utils.py

@@ -1,9 +1,6 @@
 import unittest
 from functools import partial
 
-import pytest
-
-# Assuming the function is in axolotl.utils.trainer
 from axolotl.utils.trainer import truncate_or_drop_long_seq
 
 
@@ -69,9 +66,11 @@ class TestTruncateOrDropLongSeq(unittest.TestCase):
         self.assertEqual(len(result_long["labels"]), self.sequence_len)
         self.assertEqual(result_long["labels"], list(range(self.sequence_len)))
 
-
         # Just right
-        sample_ok = {"input_ids": list(range(self.min_sequence_len)), "labels": list(range(self.min_sequence_len))}
+        sample_ok = {
+            "input_ids": list(range(self.min_sequence_len)),
+            "labels": list(range(self.min_sequence_len)),
+        }
         result_ok = handler(sample_ok)
         self.assertEqual(len(result_ok["input_ids"]), self.min_sequence_len)
         self.assertEqual(result_ok, sample_ok)  # Should be unchanged
@@ -81,7 +80,6 @@ class TestTruncateOrDropLongSeq(unittest.TestCase):
         result_empty = handler(sample_empty)
         self.assertEqual(result_empty, sample_empty)  # Unchanged
 
-
     def test_drop_mode_batched(self):
         """Test drop mode with batched examples."""
         handler = partial(
@@ -102,7 +100,6 @@ class TestTruncateOrDropLongSeq(unittest.TestCase):
         expected = [False, False, True, True, False]
         self.assertEqual(handler(sample), expected)
 
-
     def test_truncate_mode_batched(self):
         """Test truncate mode with batched examples."""
         handler = partial(
@@ -146,7 +143,6 @@ class TestTruncateOrDropLongSeq(unittest.TestCase):
             [],  # Unchanged (Empty)
         ]
 
-
         self.assertEqual(result["input_ids"], expected_input_ids)
         self.assertEqual(result["labels"], expected_labels)