From abbec970af811a4a169ed9043c62f310bca26993 Mon Sep 17 00:00:00 2001 From: Alex Mallen Date: Wed, 4 Oct 2023 20:41:10 +0000 Subject: [PATCH] fix save_logprobs --- elk/__init__.py | 4 ++++ elk/evaluation/evaluate.py | 3 +++ .../templates/_default/templates.yaml | 2 +- .../templates/_no_suffix/templates.yaml | 15 +++++++++++++++ elk/run.py | 13 ++++++------- elk/training/supervised.py | 2 +- elk/training/train.py | 7 +++++-- 7 files changed, 35 insertions(+), 11 deletions(-) create mode 100644 elk/promptsource/templates/_no_suffix/templates.yaml diff --git a/elk/__init__.py b/elk/__init__.py index cb6e8e5b..f6b13201 100644 --- a/elk/__init__.py +++ b/elk/__init__.py @@ -1,5 +1,9 @@ +from .evaluation import Eval from .extraction import Extract +from .training.train import Elicit __all__ = [ "Extract", + "Elicit", + "Eval", ] diff --git a/elk/evaluation/evaluate.py b/elk/evaluation/evaluate.py index 4f034970..a855e8d8 100644 --- a/elk/evaluation/evaluate.py +++ b/elk/evaluation/evaluate.py @@ -73,6 +73,9 @@ def apply_to_layer( } ) + if self.save_logprobs: + out_logprobs[ds_name]["lr"][mode] = dict() + for i, model in enumerate(lr_models): model.eval() val_log_odds = model(val_data.hiddens) diff --git a/elk/promptsource/templates/_default/templates.yaml b/elk/promptsource/templates/_default/templates.yaml index 7e9297f7..3ca19558 100644 --- a/elk/promptsource/templates/_default/templates.yaml +++ b/elk/promptsource/templates/_default/templates.yaml @@ -2,7 +2,7 @@ dataset: None templates: 7eab7254-bd71-4b1d-9f8a-0fc7110f8371: !Template answer_choices: False ||| True - id: 7eab7254-bd41-4b1d-9f8a-0fc7110f8371 + id: 7eab7254-bd71-4b1d-9f8a-0fc7110f8371 jinja: "{{ statement }}" metadata: !TemplateMetadata choices_in_prompt: true diff --git a/elk/promptsource/templates/_no_suffix/templates.yaml b/elk/promptsource/templates/_no_suffix/templates.yaml new file mode 100644 index 00000000..7187f83c --- /dev/null +++ b/elk/promptsource/templates/_no_suffix/templates.yaml @@ -0,0 +1,15 @@ +dataset: None +templates: + 8eab7252-bd71-4b2d-9f8a-0fc7260f8371: !Template + answer_choices: False ||| True + id: 8eab7252-bd71-4b2d-9f8a-0fc7260f8371 + jinja: "{{ statement }}" + metadata: !TemplateMetadata + choices_in_prompt: true + languages: + - en + metrics: + - Accuracy + original_task: true + name: _no_suffix + suffix: "" diff --git a/elk/run.py b/elk/run.py index 28ac1d4f..76bdd598 100644 --- a/elk/run.py +++ b/elk/run.py @@ -68,7 +68,7 @@ class Run(ABC, Serializable): concatenated_layer_offset: int = 0 debug: bool = False num_gpus: int = -1 - out_dir: Path | None = None + min_gpu_mem: int = 0 disable_cache: bool = field(default=False, to_dict=False) def execute( @@ -113,7 +113,7 @@ def execute( meta_f, ) - devices = select_usable_devices(self.num_gpus) + devices = select_usable_devices(self.num_gpus, min_memory=self.min_gpu_mem) num_devices = len(devices) func: Callable[[int], tuple[dict[str, pd.DataFrame], dict]] = partial( self.apply_to_layer, devices=devices, world_size=num_devices @@ -225,18 +225,17 @@ def apply_to_layers( if self.save_logprobs: save_dict = defaultdict(dict) for ds_name, logprobs_dict in logprobs_dicts.items(): + save_dict[ds_name]["row_ids"] = logprobs_dict[layers[0]][ + "row_ids" + ] save_dict[ds_name]["texts"] = logprobs_dict[layers[0]]["texts"] save_dict[ds_name]["labels"] = logprobs_dict[layers[0]][ "labels" ] save_dict[ds_name]["lm"] = logprobs_dict[layers[0]]["lm"] - save_dict[ds_name]["reporter"] = dict() save_dict[ds_name]["lr"] = dict() for layer, logprobs_dict_by_mode in logprobs_dict.items(): - save_dict[ds_name]["reporter"][ - layer - ] = logprobs_dict_by_mode["reporter"] save_dict[ds_name]["lr"][layer] = logprobs_dict_by_mode[ "lr" ] - torch.save(save_dict, self.out_dir / "logprobs.pt") + torch.save(dict(save_dict), self.out_dir / "logprobs.pt") diff --git a/elk/training/supervised.py b/elk/training/supervised.py index a4a53c50..b3f10064 100644 --- a/elk/training/supervised.py +++ b/elk/training/supervised.py @@ -24,7 +24,7 @@ def train_supervised( (n, v, d) = train_data.hiddens.shape train_h = rearrange(train_data.hiddens, "n v d -> (n v) d") - if erase_paraphrases: + if erase_paraphrases and v > 1: if leace is None: leace = LeaceFitter( d, diff --git a/elk/training/train.py b/elk/training/train.py index 31ca4f23..6129d9c9 100644 --- a/elk/training/train.py +++ b/elk/training/train.py @@ -26,7 +26,7 @@ class Elicit(Run): cross-validation. Defaults to "single", which means to train a single classifier on the training data. "cv" means to use cross-validation.""" - erase_paraphrases: bool = True + erase_paraphrases: bool = False """Whether to use LEACE to erase the paraphrase dimensions before training the classifier.""" @@ -41,7 +41,7 @@ def default(): datasets=("",), ) ) - + def create_models_dir(self, out_dir: Path): lr_dir = out_dir / "lr_models" @@ -115,6 +115,9 @@ def apply_to_layer( } ) + if self.save_logprobs: + out_logprobs[ds_name]["lr"][mode] = dict() + for i, model in enumerate(lr_models): model.eval() val_log_odds = model(val.hiddens)