check is streamable

elk/extraction/__init__.py

@@ -1,4 +1,16 @@
 from .balanced_sampler import BalancedSampler, FewShotSampler
-from .extraction import Extract, extract_hiddens, extract
-from .generator import _GeneratorConfig, _GeneratorBuilder
+from .extraction import Extract, extract, extract_hiddens
+from .generator import _GeneratorBuilder, _GeneratorConfig
 from .prompt_loading import PromptConfig, yield_prompts
+
+__all__ = [
+ "_GeneratorBuilder",
+ "_GeneratorConfig",
+ "BalancedSampler",
+ "extract_hiddens",
+ "extract",
+ "Extract",
+ "FewShotSampler",
+ "PromptConfig",
+ "yield_prompts",
+]

elk/extraction/balanced_sampler.py

@@ -9,12 +9,6 @@
 from ..math_util import stochastic_round_constrained
 from ..utils import infer_label_column
 from ..utils.typing import assert_type
-from collections import deque
-from datasets import IterableDataset, Features
-from itertools import cycle, islice
-from random import Random
-from torch.utils.data import IterableDataset as TorchIterableDataset
-from typing import Iterator, Optional, Iterable
 
 
 class BalancedSampler(TorchIterableDataset):

elk/extraction/extraction.py

@@ -2,7 +2,6 @@
 import logging
 import os
 from dataclasses import InitVar, dataclass
-from itertools import islice
 from typing import Iterable, Literal, Optional, Union
 
 import torch

elk/extraction/prompt_loading.py

@@ -19,16 +19,6 @@
  select_train_val_splits,
 )
 from .balanced_sampler import FewShotSampler
-from dataclasses import dataclass
-from datasets import (
- load_dataset,
- Dataset,
- Features,
-)
-from datasets.distributed import split_dataset_by_node
-from random import Random
-from simple_parsing.helpers import field, Serializable
-from typing import Any, Iterator, Literal, Optional
 
 
 @dataclass

elk/utils/__init__.py

@@ -14,6 +14,7 @@
  "get_columns_all_equal",
  "infer_label_column",
  "infer_num_classes",
+ "is_streamable",
  "float32_to_int16",
  "int16_to_float32",
  "select_train_val_splits",

elk/utils/data_utils.py

@@ -6,14 +6,12 @@
  ClassLabel,
  DatasetDict,
  Features,
- IterableDataset,
  Split,
  Value,
 )
-from random import Random
-from itertools import islice
-from typing import Iterable, List, Any
-import copy
+
+from ..promptsource.templates import Template
+from .typing import assert_type
 
 
 def get_columns_all_equal(dataset: DatasetDict) -> list[str]:

tests/test_load_prompts.py

@@ -1,11 +1,9 @@
-from elk.extraction import yield_prompts, PromptConfig
-from elk.promptsource.templates import DatasetTemplates
 from itertools import cycle, islice
 from typing import Literal
 
 import pytest
 
-from elk.extraction import PromptConfig, load_prompts
+from elk.extraction import PromptConfig, yield_prompts
 from elk.promptsource.templates import DatasetTemplates