Fix unnatural tokenizations if possible

OpenGPTX · Nov 8, 2023 · 3e5377f · 3e5377f
1 parent b05c53d
commit 3e5377f
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/lm_eval/base.py b/lm_eval/base.py
@@ -179,7 +179,20 @@ def loglikelihood(self, requests):
                 context_enc = self.tok_encode(context)
 
             continuation_enc = self.tok_encode(continuation)
-            # continuation_enc = self.tok_encode(continuation, is_continuation=True)
+            ctx_cont_enc = self.tok_encode(context + continuation)
+
+            if context_enc + continuation_enc != ctx_cont_enc:
+                if ctx_cont_enc[: len(context_enc)] == context_enc:
+                    # continuation_enc is incorrect and context_enc is correct
+                    continuation_enc = ctx_cont_enc[len(context_enc) :]
+                elif ctx_cont_enc[-len(continuation_enc) :] == continuation_enc:
+                    # continuation_enc is correct and context_enc is incorrect
+                    context_enc = ctx_cont_enc[: -len(continuation_enc)]
+                else:
+                    # Both are incorrect
+                    print(
+                        f"WARNING: Unnatural tokenization of concatenated context ...{repr(context[-20:])} and continuation {repr(continuation)}"
+                    )
 
             new_reqs.append(((context, continuation), context_enc, continuation_enc))