|
29 | 29 | maybe_extract_prompt, |
30 | 30 | maybe_unpair_preference_dataset, |
31 | 31 | pack_dataset, |
32 | | - pack_examples, |
33 | 32 | truncate_dataset, |
34 | 33 | unpair_preference_dataset, |
35 | 34 | ) |
@@ -424,48 +423,6 @@ def test_maybe_extract_prompt_standard_already_explicit(self): |
424 | 423 | ) |
425 | 424 |
|
426 | 425 |
|
427 | | -class TestPackExamples(unittest.TestCase): |
428 | | - def test_larger_chunks(self): |
429 | | - examples = { |
430 | | - "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]], |
431 | | - "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]], |
432 | | - } |
433 | | - seq_length = 5 |
434 | | - expected_output = { |
435 | | - "input_ids": [[1, 2, 3, 4, 5], [6, 7, 8]], |
436 | | - "attention_mask": [[0, 1, 1, 0, 0], [1, 1, 1]], |
437 | | - } |
438 | | - result = pack_examples(examples, seq_length) |
439 | | - self.assertEqual(result, expected_output) |
440 | | - |
441 | | - def test_smaller_chunks(self): |
442 | | - examples = { |
443 | | - "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]], |
444 | | - "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]], |
445 | | - } |
446 | | - seq_length = 2 |
447 | | - expected_output = { |
448 | | - "input_ids": [[1, 2], [3, 4], [5, 6], [7, 8]], |
449 | | - "attention_mask": [[0, 1], [1, 0], [0, 1], [1, 1]], |
450 | | - } |
451 | | - result = pack_examples(examples, seq_length) |
452 | | - self.assertEqual(result, expected_output) |
453 | | - |
454 | | - def test_with_dataset(self): |
455 | | - examples = { |
456 | | - "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]], |
457 | | - "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]], |
458 | | - } |
459 | | - dataset = Dataset.from_dict(examples) |
460 | | - seq_length = 3 |
461 | | - expected_output = { |
462 | | - "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8]], |
463 | | - "attention_mask": [[0, 1, 1], [0, 0, 1], [1, 1]], |
464 | | - } |
465 | | - dataset = dataset.map(pack_examples, batched=True, fn_kwargs={"seq_length": seq_length}) |
466 | | - self.assertEqual(dataset.to_dict(), expected_output) |
467 | | - |
468 | | - |
469 | 426 | class TestPackDatasetWrapped(unittest.TestCase): |
470 | 427 | def test_with_dataset(self): |
471 | 428 | examples = { |
|
0 commit comments