Multi datasets

elk/evaluation/evaluate.py

@@ -99,7 +99,9 @@ def evaluate_reporter(
 
  def evaluate(self):
  """Evaluate the reporter on all layers."""
- devices = select_usable_devices(self.cfg.num_gpus, min_memory=cfg.target.min_gpu_mem)
+ devices = select_usable_devices(
+ self.cfg.num_gpus, min_memory=cfg.target.min_gpu_mem
+ )
 
  num_devices = len(devices)
  func: Callable[[int], EvalLog] = partial(

elk/extraction/__init__.py

@@ -1,4 +1,4 @@
 from .balanced_sampler import BalancedSampler, FewShotSampler
 from .extraction import Extract, extract_hiddens, extract
 from .generator import _GeneratorConfig, _GeneratorBuilder
-from .prompt_loading import PromptConfig, load_prompts
+from .prompt_loading import PromptConfig, load_prompts

elk/extraction/extraction.py

@@ -9,12 +9,25 @@
 from simple_parsing import Serializable, field
 from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedModel
 
-from datasets import (Array3D, ClassLabel, DatasetDict, Features, Sequence,
- SplitDict, SplitInfo, Value, get_dataset_config_info)
+from datasets import (
+ Array3D,
+ ClassLabel,
+ DatasetDict,
+ Features,
+ Sequence,
+ SplitDict,
+ SplitInfo,
+ Value,
+ get_dataset_config_info,
+)
 from elk.utils.typing import float32_to_int16
 
-from ..utils import (assert_type, infer_label_column, select_train_val_splits,
- select_usable_devices)
+from ..utils import (
+ assert_type,
+ infer_label_column,
+ select_train_val_splits,
+ select_usable_devices,
+)
 from .balanced_sampler import BalancedSampler
 from .generator import _GeneratorBuilder
 from .prompt_loading import PromptConfig, load_prompts

elk/run.py

@@ -93,7 +93,9 @@ def prepare_data(
  def concatenate(self, layers):
  """Concatenate hidden states from a previous layer."""
  for layer in range(self.cfg.concatenated_layer_offset, len(layers)):
- layers[layer] = layers[layer] + [layers[layer][0] - cfg.concatenated_layer_offset]
+ layers[layer] = layers[layer] + [
+ layers[layer][0] - cfg.concatenated_layer_offset
+ ]
  return layers
 
  def apply_to_layers(