first commit

jailbreak_toolbox/datasets/implementations/static_dataset.py

@@ -0,0 +1,26 @@
+from ..base_dataset import BaseDataset
+from ...core.registry import dataset_registry
+from typing import Iterator, Any, List
+
+@dataset_registry.register("static")
+class StaticDataset(BaseDataset):
+    """
+    一个简单的静态内存数据集。
+    主要用于快速测试和演示，数据直接在配置文件中定义。
+    """
+    def __init__(self, prompts: List[str], **kwargs):
+        super().__init__(**kwargs)
+        self.prompts = prompts
+
+    def __iter__(self) -> Iterator[Any]:
+        return iter(self.prompts)
+    
+    def __getitem__(self, idx):
+        return self.prompts[idx]
+
+    def __len__(self) -> int:
+        return len(self.prompts)
+
+    def __getitem__(self, index: int) -> Any:
+        """Support indexing for dataset access"""
+        return self.prompts[index]