fix: address PR comments - category filter, unused imports

- Fix category filter to handle list[str] in setup_*_dataset functions
- Remove unused BenchmarkInfo imports in test files
- Align with PR #502 fixes

Author: davidberenstein1957

src/pruna/data/datasets/prompt.py

@@ -72,7 +72,14 @@ def setup_parti_prompts_dataset(
     ds = load_dataset("nateraw/parti-prompts")["train"]  # type: ignore[index]
 
     if category is not None:
-        ds = ds.filter(lambda x: x["Category"] == category or x["Challenge"] == category)
+        if isinstance(category, list):
+            ds = ds.filter(
+                lambda x: x["Category"] in category or x["Challenge"] in category
+            )
+        else:
+            ds = ds.filter(
+                lambda x: x["Category"] == category or x["Challenge"] == category
+            )
 
     ds = ds.shuffle(seed=seed)
 

tests/data/test_datamodule.py

@@ -4,7 +4,6 @@
 import torch
 from transformers import AutoTokenizer
 
-from pruna.data import BenchmarkInfo, benchmark_info
 from pruna.data.datasets.image import setup_imagenet_dataset
 from pruna.data.pruna_datamodule import PrunaDataModule