Reset extraction

Author: Kyle1668

elk/extraction/extraction.py

@@ -5,12 +5,14 @@
 from dataclasses import InitVar, dataclass
 from itertools import islice
 from typing import Any, Iterable, Literal
+from warnings import filterwarnings
 
 import torch
 from datasets import (
+    Array2D,
     Array3D,
-    ClassLabel,
     DatasetDict,
+    DownloadMode,
     Features,
     Sequence,
     SplitDict,
@@ -20,23 +22,23 @@
 )
 from simple_parsing import Serializable, field
 from torch import Tensor
-from transformers import AutoConfig, AutoTokenizer, GPT2TokenizerFast
+from transformers import AutoConfig, PreTrainedModel
 from transformers.modeling_outputs import Seq2SeqLMOutput
 
 from ..promptsource import DatasetTemplates
 from ..utils import (
     assert_type,
-    convert_span,
     float32_to_int16,
+    infer_label_column,
+    infer_num_classes,
     instantiate_model,
+    instantiate_tokenizer,
     is_autoregressive,
     select_train_val_splits,
     select_usable_devices,
 )
-from .balanced_sampler import BalancedSampler
 from .generator import _GeneratorBuilder
 from .prompt_loading import PromptConfig, load_prompts
-from ..rwkv_lm.rwkv_hf import RWKVConfig
 
 
 @dataclass
@@ -58,6 +60,7 @@ class Extract(Serializable):
     layers: tuple[int, ...] = ()
     layer_stride: InitVar[int] = 1
     token_loc: Literal["first", "last", "mean"] = "last"
+    use_encoder_states: bool = False
 
     def __post_init__(self, layer_stride: int):
         if self.layers and layer_stride > 1:
@@ -85,7 +88,7 @@ def explode(self) -> list["Extract"]:
         return copies
 
 
-@torch.no_grad()
+@torch.inference_mode()
 def extract_hiddens(
     cfg: "Extract",
     *,
@@ -99,135 +102,135 @@ def extract_hiddens(
 
     # Silence datasets logging messages from all but the first process
     if rank != 0:
+        filterwarnings("ignore")
         logging.disable(logging.CRITICAL)
 
-    ds_names = cfg.prompts.datasets
+    p_cfg = cfg.prompts
+    ds_names = p_cfg.datasets
     assert len(ds_names) == 1, "Can only extract hiddens from one dataset at a time."
 
-    prompt_ds = load_prompts(
-        ds_names[0],
-        split_type=split_type,
-        stream=cfg.prompts.stream,
-        rank=rank,
-        world_size=world_size,
-    )  # this dataset is already sharded, but hasn't been truncated to max_examples
-
     model = instantiate_model(
         cfg.model, torch_dtype="auto" if device != "cpu" else torch.float32
     ).to(device)
-    tokenizer = None
+    tokenizer = instantiate_tokenizer(
+        cfg.model, truncation_side="left", verbose=rank == 0
+    )
 
-    if cfg.model.startswith("RWKV"):
-        tokenizer = GPT2TokenizerFast(tokenizer_file='/home/kyle/repos/elk/elk/rwkv_lm/20B_tokenizer.json')
-    else:
-        tokenizer = AutoTokenizer.from_pretrained(
-            cfg.model, truncation_side="left", verbose=False
-        )
+    is_enc_dec = model.config.is_encoder_decoder
+    if is_enc_dec and cfg.use_encoder_states:
+        assert hasattr(model, "get_encoder") and callable(model.get_encoder)
+        model = assert_type(PreTrainedModel, model.get_encoder())
+        is_enc_dec = False
 
-    has_lm_preds = is_autoregressive(model.config)
+    has_lm_preds = is_autoregressive(model.config, not cfg.use_encoder_states)
     if has_lm_preds and rank == 0:
         print("Model has language model head, will store predictions.")
 
+    prompt_ds = load_prompts(
+        ds_names[0],
+        label_column=p_cfg.label_columns[0] if p_cfg.label_columns else None,
+        num_classes=p_cfg.num_classes,
+        split_type=split_type,
+        stream=p_cfg.stream,
+        rank=rank,
+        world_size=world_size,
+    )
+
     # Iterating over questions
     layer_indices = cfg.layers or tuple(range(model.config.num_hidden_layers))
 
-    global_max_examples = cfg.prompts.max_examples[0 if split_type == "train" else 1]
+    global_max_examples = p_cfg.max_examples[0 if split_type == "train" else 1]
     # break `max_examples` among the processes roughly equally
     max_examples = global_max_examples // world_size
     # the last process gets the remainder (which is usually small)
     if rank == world_size - 1:
         max_examples += global_max_examples % world_size
 
-    for example in islice(BalancedSampler(prompt_ds), max_examples):
+    for example in islice(prompt_ds, max_examples):
         num_variants = len(example["prompts"])
+        num_choices = len(example["prompts"][0])
+
         hidden_dict = {
             f"hidden_{layer_idx}": torch.empty(
                 num_variants,
-                2,  # contrast pair
+                num_choices,
                 model.config.hidden_size,
                 device=device,
                 dtype=torch.int16,
             )
             for layer_idx in layer_indices
         }
-        lm_preds = torch.empty(
+        lm_logits = torch.empty(
             num_variants,
-            2,  # contrast pair
+            num_choices,
             device=device,
             dtype=torch.float32,
         )
-        text_inputs = []
+        text_questions = []
 
         # Iterate over variants
         for i, record in enumerate(example["prompts"]):
-            variant_inputs = []
+            variant_questions = []
 
             # Iterate over answers
             for j, choice in enumerate(record):
-                text = choice["text"]
-
-                # TODO: Do something smarter than "rindex" here. Really we want to
-                # get the span of the answer directly from Jinja, but that doesn't
-                # seem possible. This approach may fail for complex templates.
-                answer_start = text.rindex(choice["answer"])
+                text = choice["question"]
 
                 # Only feed question, not the answer, to the encoder for enc-dec models
-                if model.config.is_encoder_decoder:
-                    # TODO: Maybe make this more generic for complex templates?
-                    text = text[:answer_start].rstrip()
-                    target = choice["answer"]
-                else:
-                    target = None
-
-                # Record the EXACT string we fed to the model
-                variant_inputs.append(text)
-                # inputs = None
-                # if cfg.model.startswith("RWKV"):
-                #     inputs = tokenizer(
-                #         text,
-                #         return_offsets_mapping=True,
-                #         text_target=target,  # type: ignore[arg-type]
-                #         truncation=True,
-                #     )
-                # else:
-                inputs = tokenizer(
+                target = choice["answer"] if is_enc_dec else None
+
+                # Record the EXACT question we fed to the model
+                variant_questions.append(text)
+                encoding = tokenizer(
                     text,
-                    return_offsets_mapping=True,
+                    add_special_tokens=False,
                     return_tensors="pt",
                     text_target=target,  # type: ignore[arg-type]
                     truncation=True,
-                )
+                ).to(device)
+                input_ids = assert_type(Tensor, encoding.input_ids)
+
+                if is_enc_dec:
+                    answer = assert_type(Tensor, encoding.labels)
+                else:
+                    encoding2 = tokenizer(
+                        choice["answer"],
+                        add_special_tokens=False,
+                        return_tensors="pt",
+                    ).to(device)
+                    answer = assert_type(Tensor, encoding2.input_ids)
 
-                # The offset_mapping is a sorted list of (start, end) tuples. We locate
-                # the start of the answer in the tokenized sequence with binary search.
-                offsets = inputs.pop("offset_mapping") if cfg.model.startswith("RWKV") else inputs.pop("offset_mapping").squeeze().tolist()
-                inputs = inputs if cfg.model.startswith("RWKV") else  inputs.to(device)
+                    input_ids = torch.cat([input_ids, answer], dim=-1)
+                    if max_len := tokenizer.model_max_length:
+                        input_ids = input_ids[..., -max_len:]
 
-                # Run the forward pass
-                outputs = model(**inputs) if cfg.model.startswith("RWKV") else model(**inputs, output_hidden_states=True)
+                # Make sure we only pass the arguments that the model expects
+                inputs = dict(input_ids=input_ids)
+                if is_enc_dec:
+                    inputs["labels"] = answer
+
+                with torch.autocast("cuda", enabled=torch.cuda.is_available()):
+                    outputs = model(**inputs, output_hidden_states=True)
 
                 # Compute the log probability of the answer tokens if available
                 if has_lm_preds:
-                    start, end = convert_span(
-                        offsets, (answer_start, answer_start + len(choice["answer"]))
-                    )
-                    log_p = outputs.logits[..., start - 1 : end - 1, :].log_softmax(
-                        dim=-1
-                    )
-                    tokens = inputs.input_ids[..., start:end, None]
-                    lm_preds[i, j] = log_p.gather(-1, tokens).sum()
+                    answer_len = answer.shape[-1]
+
+                    log_p = outputs.logits[..., -answer_len:, :].log_softmax(dim=-1)
+                    tokens = answer[..., None]
+                    lm_logits[i, j] = log_p.gather(-1, tokens).sum()
 
                 elif isinstance(outputs, Seq2SeqLMOutput):
                     # The cross entropy loss is averaged over tokens, so we need to
                     # multiply by the length to get the total log probability.
-                    length = inputs.labels.shape[-1]
-                    lm_preds[i, j] = -assert_type(Tensor, outputs.loss) * length
+                    length = encoding.labels.shape[-1]
+                    lm_logits[i, j] = -assert_type(Tensor, outputs.loss) * length
 
-                hiddens = outputs if cfg.model.startswith("RWKV") else (
+                hiddens = (
                     outputs.get("decoder_hidden_states") or outputs["hidden_states"]
                 )
                 # First element of list is the input embeddings
-                hiddens = hiddens if cfg.model.startswith("RWKV") else hiddens[1:]
+                hiddens = hiddens[1:]
 
                 # Throw out layers we don't care about
                 hiddens = [hiddens[i] for i in layer_indices]
@@ -245,17 +248,16 @@ def extract_hiddens(
                 for layer_idx, hidden in zip(layer_indices, hiddens):
                     hidden_dict[f"hidden_{layer_idx}"][i, j] = float32_to_int16(hidden)
 
-            text_inputs.append(variant_inputs)
+            text_questions.append(variant_questions)
 
         out_record: dict[str, Any] = dict(
             label=example["label"],
             variant_ids=example["template_names"],
-            text_inputs=text_inputs,
+            text_questions=text_questions,
             **hidden_dict,
         )
         if has_lm_preds:
-            # We only need the probability of the positive example since this is binary
-            out_record["model_preds"] = lm_preds.softmax(dim=-1)[..., 1]
+            out_record["model_logits"] = lm_logits
 
         yield out_record
 
@@ -266,7 +268,11 @@ def _extraction_worker(**kwargs):
 
 
 def extract(
-    cfg: "Extract", num_gpus: int = -1, min_gpu_mem: int | None = None
+    cfg: "Extract",
+    *,
+    disable_cache: bool = False,
+    num_gpus: int = -1,
+    min_gpu_mem: int | None = None,
 ) -> DatasetDict:
     """Extract hidden states from a model and return a `DatasetDict` containing them."""
 
@@ -292,15 +298,18 @@ def get_splits() -> SplitDict:
             dataset_name=available_splits.dataset_name,
         )
 
-    model_cfg = None
-    if cfg.model.startswith("RWKV"):
-        model_cfg = RWKVConfig()
-    else:
-        model_cfg = AutoConfig.from_pretrained(cfg.model)
+    model_cfg = AutoConfig.from_pretrained(cfg.model)
 
     ds_name, _, config_name = cfg.prompts.datasets[0].partition(" ")
     info = get_dataset_config_info(ds_name, config_name or None)
 
+    ds_features = assert_type(Features, info.features)
+    label_col = (
+        cfg.prompts.label_columns[0]
+        if cfg.prompts.label_columns
+        else infer_label_column(ds_features)
+    )
+    num_classes = cfg.prompts.num_classes or infer_num_classes(ds_features[label_col])
     num_variants = cfg.prompts.num_variants
     if num_variants < 0:
         prompter = DatasetTemplates(ds_name, config_name)
@@ -309,7 +318,7 @@ def get_splits() -> SplitDict:
     layer_cols = {
         f"hidden_{layer}": Array3D(
             dtype="int16",
-            shape=(num_variants, 2, model_cfg.hidden_size),
+            shape=(num_variants, num_classes, model_cfg.hidden_size),
         )
         for layer in cfg.layers or range(model_cfg.num_hidden_layers)
     }
@@ -318,21 +327,20 @@ def get_splits() -> SplitDict:
             Value(dtype="string"),
             length=num_variants,
         ),
-        "label": ClassLabel(names=["neg", "pos"]),
-        "text_inputs": Sequence(
+        "label": Value(dtype="int64"),
+        "text_questions": Sequence(
             Sequence(
                 Value(dtype="string"),
-                length=2,
             ),
             length=num_variants,
         ),
     }
 
-    # Only add model_preds if the model is an autoregressive model
-    if is_autoregressive(model_cfg):
-        other_cols["model_preds"] = Sequence(
-            Value(dtype="float32"),
-            length=num_variants,
+    # Only add model_logits if the model is an autoregressive model
+    if is_autoregressive(model_cfg, not cfg.use_encoder_states):
+        other_cols["model_logits"] = Array2D(
+            shape=(num_variants, num_classes),
+            dtype="float32",
         )
 
     devices = select_usable_devices(num_gpus, min_memory=min_gpu_mem)
@@ -361,7 +369,10 @@ def get_splits() -> SplitDict:
 
     ds = dict()
     for split, builder in builders.items():
-        builder.download_and_prepare(num_proc=len(devices))
+        builder.download_and_prepare(
+            download_mode=DownloadMode.FORCE_REDOWNLOAD if disable_cache else None,
+            num_proc=len(devices),
+        )
         ds[split] = builder.as_dataset(split=split)
 
-    return DatasetDict(ds)
+    return DatasetDict(ds)