From 974075827aa8115b0341d10923ae9b5ebb33a181 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CKlaudiaTH=E2=80=9D?= Date: Mon, 16 Oct 2023 18:45:12 +0200 Subject: [PATCH] MegatronLM Client: Truncate to max_length and not max_length+1 in _loglikelihood_tokens --- lm_eval/models/megatronlm.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/lm_eval/models/megatronlm.py b/lm_eval/models/megatronlm.py index ba8675b98f..89d5bd87cc 100644 --- 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) )