From 02855031c5dc41d6b1efbb3b838e4f992ea8b50f Mon Sep 17 00:00:00 2001
From: Daniel King <daniel@mosaicml.com>
Date: Thu, 13 Jun 2024 16:04:45 -0700
Subject: [PATCH] pyright

---
 tests/data/test_packing.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/data/test_packing.py b/tests/data/test_packing.py
index 7af4666617..bb28ec40e7 100644
--- a/tests/data/test_packing.py
+++ b/tests/data/test_packing.py
@@ -16,6 +16,7 @@
 from llmfoundry.data.finetuning.dataloader import build_finetuning_dataloader
 from llmfoundry.data.packing import BinPackCollator, auto_packing_ratio
 from llmfoundry.utils.builders import build_tokenizer
+from llmfoundry.data.finetuning.tasks import StreamingFinetuningDataset
 
 
 def _data_to_batch(data: List[List[int]], max_seq_len: int,
@@ -205,7 +206,9 @@ def test_auto_packing_with_streaming_dataloader(tmp_path: Path):
         batch_ix += 1
         if batch_ix >= 3:
             break
-
+    
+    assert isinstance(loader, DataLoader)
+    assert isinstance(loader.dataset, StreamingFinetuningDataset)
     assert loader.dataset.packing_ratio is not None
     assert loader.dataset.packing_ratio == int(loader.batch_size / 6)