From 636bb9c1e65e72c3a27049dacb3200234d1c2782 Mon Sep 17 00:00:00 2001
From: hiyouga <467089858@qq.com>
Date: Thu, 4 Jul 2024 04:00:12 +0800
Subject: [PATCH] update tests

---
 tests/model/model_utils/test_packing.py | 75 +++++++++++++++----------
 1 file changed, 45 insertions(+), 30 deletions(-)

diff --git a/tests/model/model_utils/test_packing.py b/tests/model/model_utils/test_packing.py
index 8056099f..bee21691 100644
--- a/tests/model/model_utils/test_packing.py
+++ b/tests/model/model_utils/test_packing.py
@@ -12,42 +12,57 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 import torch
 
 from llamafactory.model.model_utils.packing import get_seqlens_in_batch, get_unpad_data
 
 
-def test_get_seqlens_in_batch():
-    attention_mask_with_indices = torch.tensor(
-        [
-            [1, 1, 2, 2, 2, 0],
-            [1, 2, 2, 3, 3, 3],
-        ]
-    )
+@pytest.mark.parametrize(
+    "attention_mask,golden_seq_lens",
+    [
+        (
+            [
+                [1, 1, 2, 2, 2, 0],
+                [1, 2, 2, 3, 3, 3],
+            ],
+            [2, 3, 1, 2, 3],
+        ),
+        (
+            [[1]],
+            [1],
+        ),
+    ],
+)
+def test_get_seqlens_in_batch(attention_mask, golden_seq_lens):
+    attention_mask_with_indices = torch.tensor(attention_mask)
     seqlens_in_batch = get_seqlens_in_batch(attention_mask_with_indices)
-    assert list(seqlens_in_batch.size()) == [5]
-    assert torch.all(seqlens_in_batch == torch.tensor([2, 3, 1, 2, 3]))
-
-    attention_mask_with_indices = torch.tensor([[1, 1, 1]])
-    seqlens_in_batch = get_seqlens_in_batch(attention_mask_with_indices)
-    assert list(seqlens_in_batch.size()) == [1]
-    assert torch.all(seqlens_in_batch == torch.tensor([3]))
+    assert torch.all(seqlens_in_batch == torch.tensor(golden_seq_lens))
 
 
-def test_get_unpad_data():
-    attention_mask_with_indices = torch.tensor(
-        [
-            [1, 1, 2, 2, 2, 0],
-            [1, 2, 2, 3, 3, 3],
-        ]
-    )
+@pytest.mark.parametrize(
+    "attention_mask,golden_indices,golden_cu_seqlens,golden_max_seqlen",
+    [
+        (
+            [
+                [1, 1, 2, 2, 2, 0],
+                [1, 2, 2, 3, 3, 3],
+            ],
+            [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11],
+            [0, 2, 5, 6, 8, 11],
+            3,
+        ),
+        (
+            [[1]],
+            [0],
+            [0, 1],
+            1,
+        ),
+    ],
+)
+def test_get_unpad_data(attention_mask, golden_indices, golden_cu_seqlens, golden_max_seqlen):
+    attention_mask_with_indices = torch.tensor(attention_mask)
     indices, cu_seqlens, max_seqlen_in_batch = get_unpad_data(attention_mask_with_indices)
-    assert torch.all(indices == torch.tensor([0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11]))
-    assert torch.all(cu_seqlens == torch.tensor([0, 2, 5, 6, 8, 11], dtype=torch.int32))
-    assert max_seqlen_in_batch == 3
-
-    attention_mask_with_indices = torch.tensor([[1, 1, 1]])
-    indices, cu_seqlens, max_seqlen_in_batch = get_unpad_data(attention_mask_with_indices)
-    assert torch.all(indices == torch.tensor([0, 1, 2]))
-    assert torch.all(cu_seqlens == torch.tensor([0, 3], dtype=torch.int32))
-    assert max_seqlen_in_batch == 3
+    assert torch.all(indices == torch.tensor(golden_indices))
+    assert torch.all(cu_seqlens == torch.tensor(golden_cu_seqlens, dtype=torch.int32))
+    assert max_seqlen_in_batch == golden_max_seqlen