MegatronLM Client: Truncate to max_length and not max_length+1 in _lo…

…glikelihood_tokens
OpenGPTX · Nov 5, 2023 · 9740758 · 9740758
1 parent ac8192f
commit 9740758
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/lm_eval/models/megatronlm.py b/lm_eval/models/megatronlm.py
@@ -176,9 +176,8 @@ def _collate(x):
             inps = []
             ctxlens = []
             for cache_key, context_enc, continuation_enc in chunk:
-                # max_length+1 because the API takes up to 2049 tokens, including the first context token
-                inp = (context_enc + continuation_enc)[-(self.max_length + 1) :]
-                # TODO: the logic is much simpler if we just look at the length of continuation tokens
+                inp = (context_enc + continuation_enc)[-self.max_length :]
+
                 ctxlen = len(context_enc) - max(
                     0, len(context_enc) + len(continuation_enc) - (self.max_length + 1)
                 )