tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Linting fix

Browse Source
This commit is contained in:
mhenrhcsen
2025-05-12 22:53:30 +02:00
parent 5dd8f0b2b8
commit f6623c34cc
1 changed files with 27 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/test_trainer_utils.py
View File

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