batch surprisal computation, now GPU friendly

syntaxgym.py

@@ -174,68 +174,82 @@ class SyntaxGym(evaluate.EvaluationModule):
 
         tokenizer, tokenizer_kwargs = prepare_tokenizer(model, batch_size, add_start_token)
 
+        # Flatten sentences, enforcing that sentences are always ordered by the same condition.
+        condition_order = dataset[0]["conditions"]["condition_name"]
+        all_sentences = []
+        for item in dataset:
+            for condition_name in condition_order:
+                # Get idx of condition for this item.
+                condition_idx = item["conditions"]["condition_name"].index(condition_name)
+                all_sentences.append(item["conditions"]["content"][condition_idx])
+
+        # Tokenize sentences and split into batches.
+        all_tokenized_sentences = tokenizer(all_sentences, return_tensors="pt",
+                                            return_offsets_mapping=True,
+                                            **tokenizer_kwargs).to(device)
+        tokenized_batches = torch.split(all_tokenized_sentences["input_ids"], batch_size)
+
+        # Compute surprisal per-batch and combine into a single surprisal tensor.
+        n_sentences, n_timesteps = all_tokenized_sentences["input_ids"].shape
+        surprisals = torch.zeros(n_sentences, n_timesteps - 1).float().to(device)
+        for i, batch in enumerate(datasets.logging.tqdm(tokenized_batches)) :
+            batch = batch.to(device)
+            with torch.no_grad():
+                # logits are B * T * V
+                b_logits = model(batch)["logits"]
+                b_surprisals = -b_logits.log_softmax(dim=2) / np.log(2)
+
+            # Get surprisals of ground-truth words.
+            gt_idxs = batch[:, 1:]
+            # Reindexed surprisals: B * (T - 1)
+            b_surprisals_gt = torch.gather(b_surprisals[:, :-1, :], 2, gt_idxs.unsqueeze(2)).squeeze(2)
+
+            surprisals[i * batch_size : (i + 1) * batch_size] = b_surprisals_gt
+
+        # Reshape to intuitive axes n_items * n_conditions * ...
+        surprisals = surprisals.reshape((len(dataset), len(condition_order), -1))
+        offset_mapping = all_tokenized_sentences["offset_mapping"] \
+            .reshape((len(dataset), len(condition_order), -1, 2))
+
+        # Now evaluate per-item.
         results = {}
         result_keys = ["prediction_results", "region_totals"]
-        # TODO batch all items together
-        for item in datasets.logging.tqdm(dataset):
-            result_single = self._compute_single(item, tokenizer, tokenizer_kwargs,
-                                                 model, device)
+        for item, item_surprisals, item_offset_mapping in zip(datasets.logging.tqdm(dataset), surprisals, offset_mapping):
+            result_i = self._compute_item(item, item_surprisals, item_offset_mapping, condition_order)
 
             suite_name = item["suite_name"]
             if suite_name not in results:
                 results[suite_name] = SyntaxGymMetricSuiteResult(suite_name, [], [])
             for k in result_keys:
-                getattr(results[suite_name], k).append(result_single[k])
+                getattr(results[suite_name], k).append(result_i[k])
 
         return results
 
-    def _compute_single(self, item, tokenizer, tokenizer_kwargs, model, device):
-        tokenized = tokenizer(item["conditions"]["content"],
-                              return_tensors="pt",
-                              return_offsets_mapping=True,
-                              **tokenizer_kwargs).to(device)
-
-        # input_ids: B * T
-        input_ids = tokenized["input_ids"]
-        assert input_ids.ndim == 2
-
-        # Compute sentence level surprisals.
-        with torch.no_grad():
-            # Pre-softmax predictive distribution B * T * V
-            logits = model(input_ids).logits
-            surprisals = -logits.log_softmax(dim=2) / np.log(2)
-
-        # surprisals: B * T * V
-        assert surprisals.ndim == 3
-
-        # Get surprisals of expected words.
-        surps_shifted = surprisals[:, :-1, :]
-        expected_ids = input_ids[:, 1:]
-
-        # reindexed surprisals: B * (T - 1)
-        surprisals = torch.gather(surps_shifted, 2, expected_ids.unsqueeze(2)) \
-            .squeeze(2)
+    def _compute_item(self, item, item_surprisals, offset_mapping, condition_order):
+        """
+        Aggregate token-level surprisals to region-level surprisals for the given item,
+        and evaluate the item's predictions.
+        """
 
         #### aggregate
-        condition_names = item["conditions"]["condition_name"]
         region_totals = {condition_name: defaultdict(float)
-                         for condition_name in condition_names}
+                         for condition_name in condition_order}
         region2tokens = self.compute_region_token_mapping(
-            item, input_ids, tokenized["offset_mapping"])
+            item, condition_order, offset_mapping)
 
-        for i, (i_cond, i_inputs) in enumerate(zip(condition_names, input_ids)):
-            for region_number, region_tokens in region2tokens[i_cond].items():
+        for i, (cond_i, surprisals_i) in enumerate(zip(condition_order, item_surprisals)):
+            for region_number, region_tokens in region2tokens[cond_i].items():
                 for token in region_tokens:
                     if token == 0:
                         # surprisal not defined. pass.
                         continue
-                    elif token <= surprisals.shape[1]:
-                        region_totals[i_cond][region_number] += surprisals[i, token - 1]
+                    elif token <= item_surprisals.shape[1]:
+                        region_totals[cond_i][region_number] += surprisals_i[token - 1]
                     else:
                         # TODO don't think this is an issue, just should clean
                         # up the aggregation output
-                        assert token == surprisals.shape[1], \
-                            "%s %s" % (token, surprisals.shape[1])
+                        assert token == surprisals_i.shape[1], \
+                            "%s %s" % (token, surprisals_i.shape[1])
 
         region_totals = {(condition_name, region_number): float(total)
                          for condition_name, totals in region_totals.items()
@@ -275,23 +289,20 @@ class SyntaxGym(evaluate.EvaluationModule):
 
         return ret
 
-    def compute_region_token_mapping(self, item, input_ids: torch.LongTensor,
+    def compute_region_token_mapping(self, item, condition_order,
                                      offset_mapping: List[Tuple[int, int]]
                                      ) -> Dict[str, Dict[int, List[int]]]:
-        # input_ids: B * T
         # offset_mapping: B * T * 2
-        # assumes batch is sorted according to item's condition_name order
 
-        condition_names = item["conditions"]["condition_name"]
-        region2tokens = {cond: defaultdict(list) for cond in condition_names}
+        region2tokens = {cond: defaultdict(list) for cond in condition_order}
 
         max_long = torch.iinfo(torch.int64).max
 
-        for i_cond, (i_tokens, i_offsets) in enumerate(zip(input_ids, offset_mapping)):
+        for i_cond, i_offsets in enumerate(offset_mapping):
             region_edges = self.get_region_edges(item, i_cond)
 
             t_cursor, r_cursor = 0, 0
-            while t_cursor < i_tokens.shape[0]:
+            while t_cursor < i_offsets.shape[0]:
                 # token = i_tokens[t_cursor]
                 token_char_start, token_char_end = i_offsets[t_cursor]
 
@@ -310,7 +321,7 @@ class SyntaxGym(evaluate.EvaluationModule):
                     r_cursor += 1
                     continue
 
-                region2tokens[condition_names[i_cond]][r_cursor + 1].append(t_cursor)
+                region2tokens[condition_order[i_cond]][r_cursor + 1].append(t_cursor)
                 t_cursor += 1
 
         return region2tokens