batchify processing of multiple answers for speedup

Author: gianlucadetommaso

benchmarks/hallucination/mmlu/run.py

@@ -1,14 +1,14 @@
+import os
 import pickle
-from string import ascii_uppercase as auc
 
 from datasets import (
     get_dataset_config_names,
     load_dataset,
 )
 import numpy as np
 from transformers import (
-    GPT2LMHeadModel,
-    GPT2TokenizerFast,
+    AutoModelForCausalLM,
+    AutoTokenizer,
 )
 
 from fortuna.hallucination import HallucinationMulticalibrator
@@ -20,9 +20,9 @@
 
 if __name__ == "__main__":
     device = "cuda"
-    model_id = "gpt2-large"
-    model = GPT2LMHeadModel.from_pretrained(model_id).to(device)
-    tokenizer = GPT2TokenizerFast.from_pretrained(model_id)
+    model_id = "gpt2"
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
 
     # download and prepare data
     task_list = get_dataset_config_names("lukaemon/mmlu")
@@ -31,21 +31,24 @@
             load_dataset(
                 "lukaemon/mmlu",
                 task,
+                cache_dir=".cache/huggingface/datasets/"
+                if os.path.isdir(".cache/huggingface/datasets/")
+                else None,
             ),
             task,
         )
         for task in task_list
     ]
 
-    answer_map = {a: i for i, a in enumerate(auc)}
+    answer_map = dict(zip(["A", "B", "C", "D"], [0, 1, 2, 3]))
     samples = []
     for datasets, task in dataset_list:
         for dataset_key, dataset in datasets.items():
             for sample in dataset:
                 samples.append(
                     dict(
                         question=string_cleaner(sample["input"]),
-                        choices=[sample[letter] for letter in ["A", "B", "C", "D"]],
+                        choices=[sample[letter] for letter in answer_map.keys()],
                         targets=answer_map[sample["target"]],
                     )
                 )

fortuna/hallucination/base.py

@@ -1,3 +1,4 @@
+import logging
 from typing import (
     Callable,
     Dict,
@@ -57,6 +58,9 @@ def __init__(
         """
         self.generative_model = generative_model
         self.tokenizer = tokenizer
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            logging.info("`tokenizer.pad_token` is None. Set to `tokenizer.eos_token`.")
         self.embedding_reduction_fn = (
             embedding_reduction_fn or locally_linear_embedding_fn
         )
@@ -248,26 +252,15 @@ def _compute_scores_embeddings_which_choices(
         which_choices = []
 
         for text, context in tqdm(zip(texts, contexts)):
-            context_inputs = self.tokenizer(context, return_tensors="pt").to(
-                self.generative_model.device
-            )
+            _logits, _scores = self._get_logits_scores(text, context)
+            _embeddings = _logits.mean(1)
             if isinstance(text, list):
-                _scores = []
-                _embeddings = []
-
-                for _text in text:
-                    __logits, __scores = self._get_logits_scores(_text, context_inputs)
-                    _embeddings.append(__logits.mean(1))
-                    _scores.append(__scores)
-
                 which_choice = np.argmax(_scores)
                 which_choices.append(which_choice)
                 scores.append(_scores[which_choice])
                 embeddings.append(_embeddings[which_choice])
-
             elif isinstance(text, str):
-                _logits, _scores = self._get_logits_scores(text, context_inputs)
-                embeddings.append(_logits.mean(1))
+                embeddings.append(_embeddings)
                 scores.append(_scores)
 
         return (
@@ -277,28 +270,29 @@ def _compute_scores_embeddings_which_choices(
         )
 
     def _get_logits_scores(
-        self, _text: str, context_inputs
+        self, text: str, context: str
     ) -> Tuple[np.ndarray, np.ndarray]:
-        _text_inputs = self.tokenizer(_text, return_tensors="pt").to(
+        context_inputs = self.tokenizer(context, return_tensors="pt", padding=True).to(
             self.generative_model.device
         )
-        _inputs = {
-            k: torch.cat((context_inputs[k], v), dim=1) for k, v in _text_inputs.items()
+        text_inputs = self.tokenizer(text, return_tensors="pt", padding=True).to(
+            self.generative_model.device
+        )
+        inputs = {
+            k: torch.cat((context_inputs[k].repeat((v.shape[0], 1)), v), dim=1)
+            for k, v in text_inputs.items()
         }
 
         with torch.no_grad():
-            __logits = self.generative_model(
-                input_ids=_inputs["input_ids"],
-                attention_mask=_inputs["attention_mask"],
-            ).logits
-
-        __scores = self.scoring_fn(
-            logits=__logits,
-            labels=_inputs["input_ids"],
+            _logits = self.generative_model(**inputs).logits
+
+        _scores = self.scoring_fn(
+            logits=_logits,
+            labels=inputs["input_ids"],
             init_pos=len(context_inputs),
         )
 
-        return __logits.cpu().numpy(), __scores.cpu().numpy()
+        return _logits.cpu().numpy(), _scores.cpu().numpy()
 
 
 def locally_linear_embedding_fn(x: np.ndarray) -> np.ndarray:

fortuna/hallucination/scoring/inv_perplexity.py

@@ -8,10 +8,9 @@ def perplexity(logits: torch.Tensor, labels: torch.Tensor, init_pos: int = 0):
     shift_logits = logits[..., :-1, :].contiguous()
     shift_labels = labels[..., 1:].contiguous()
 
-    perplexities = torch.exp(
-        loss_fct(shift_logits.transpose(1, 2), shift_labels)[:, init_pos:].mean()
+    return torch.exp(
+        loss_fct(shift_logits.transpose(1, 2), shift_labels)[:, init_pos:].mean(1)
     )
-    return torch.mean(perplexities)
 
 
 @torch.no_grad()

tests/fortuna/hallucination/embeddings.py

@@ -4,8 +4,6 @@
 import numpy as np
 from sklearn.mixture import GaussianMixture
 
-from fortuna.data import InputsLoader
-from fortuna.hallucination.embedding import EmbeddingManager
 from fortuna.hallucination.grouping.clustering.base import GroupingModel
 
 
@@ -15,61 +13,44 @@ def __init__(self, *args, **kwargs):
         self.n_inputs = 10
         self.n_features = 4
         self.n_reduced_features = 3
-        self.n_extra_features = 5
-        self.inputs_loader = InputsLoader.from_array_inputs(
-            random.normal(random.PRNGKey(0), shape=(self.n_inputs, self.n_features)),
-            batch_size=2,
-        )
-        self.grouping_model = GroupingModel(
-            embedding_manager=EmbeddingManager(
-                encoding_fn=lambda x: 1 - x,
-                reduction_fn=lambda x: x[:, : self.n_reduced_features],
-            )
+        self.embeddings = random.normal(
+            random.PRNGKey(0), shape=(self.n_inputs, self.n_features)
         )
+        self.grouping_model = GroupingModel()
         self.extra_embeddings = random.normal(
             random.PRNGKey(0), shape=(self.n_inputs, self.n_extra_features)
         )
         self.clustering_models = [GaussianMixture(n_components=i) for i in range(2, 4)]
 
     def test_all(self):
         self.grouping_model.fit(
-            inputs_loader=self.inputs_loader,
-            extra_embeddings=None,
+            embeddings=self.embeddings,
             clustering_models=self.clustering_models,
         )
-        self._check_shape_types(extra_embeddings=None)
+        self._check_shape_types()
 
         self.grouping_model.fit(
-            inputs_loader=self.inputs_loader,
-            extra_embeddings=self.extra_embeddings,
+            embeddings=self.embeddings,
             clustering_models=self.clustering_models,
         )
-        self._check_shape_types(extra_embeddings=self.extra_embeddings)
+        self._check_shape_types()
 
         with self.assertRaises(ValueError):
             self.grouping_model.fit(
-                inputs_loader=self.inputs_loader,
-                extra_embeddings=None,
+                embeddings=self.embeddings,
                 clustering_models=[],
             )
 
         with self.assertRaises(ValueError):
             self.grouping_model.fit(
-                inputs_loader=self.inputs_loader,
-                extra_embeddings=np.zeros((self.n_inputs + 1, 2)),
+                embeddings=self.embeddings,
                 clustering_models=[],
             )
 
-    def _check_shape_types(self, extra_embeddings):
-        probs = self.grouping_model.predict_proba(
-            inputs_loader=self.inputs_loader, extra_embeddings=extra_embeddings
-        )
-        hard_preds = self.grouping_model.hard_predict(
-            inputs_loader=self.inputs_loader, extra_embeddings=extra_embeddings
-        )
-        soft_preds = self.grouping_model.hard_predict(
-            inputs_loader=self.inputs_loader, extra_embeddings=extra_embeddings
-        )
+    def _check_shape_types(self):
+        probs = self.grouping_model.predict_proba(embeddings=self.embeddings)
+        hard_preds = self.grouping_model.hard_predict(embeddings=self.embeddings)
+        soft_preds = self.grouping_model.hard_predict(embeddings=self.embeddings)
         assert probs.shape == (
             self.n_inputs,
             self.grouping_model._clustering_model.n_components,