Roll back last token optimization

lm_eval/models/gpt2.py

@@ -41,12 +41,13 @@ class GPT2LM(LM):
             # TODO: automatic batch size detection for vectorization
 
             def _collate(x):
-                toks = self.tokenizer.encode(x[0] + x[1])[:-1]
-                return (len(toks), self.tokenizer.decode(toks))
+                toks = self.tokenizer.encode(x[0] + x[1])
+                return (len(toks), x)
             
             reord = utils.Reorderer(requests, _collate)
             for context, continuation in tqdm(reord.get_reordered()):
                 # when too long to fit in context, truncate from the left
+                combined_toks = self.tokenizer.encode(context + continuation)
 
                 if context == "":
                     # end of text as context
@@ -68,12 +69,9 @@ class GPT2LM(LM):
 
                 logits = torch.gather(logits, 2, cont_toks.unsqueeze(-1)).squeeze(-1) # [batch, seq]
 
-                res.append((float(logits[:, :-1].sum() if logits.shape[-1] > 1 else 0), last_token_slice, bool(max_equal)))
+                res.append((float(logits.sum()), bool(max_equal)))
 
-        # optimization: if two requests have everything the same except the last token, use 
-        # last token distribution to save compute
-        lasttoks = [self.tokenizer.encode(x[1])[-1] for x in requests]
-        return [(l + lts[lasttok], m) for (l, lts, m), lasttok in zip(reord.get_original(res), lasttoks)]
+        return reord.get_original(res)
     
     def greedy_until(self, requests):
         # TODO: implement fully general `until` that handles untils that are