From e3baa0da861f2545734bcae21f989bb4d6dc3294 Mon Sep 17 00:00:00 2001 From: josc146 Date: Fri, 1 Mar 2024 13:18:03 +0800 Subject: [PATCH] improve occurrence[token] condition --- backend-python/utils/rwkv.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py index 957d536..6a0f64a 100644 --- a/backend-python/utils/rwkv.py +++ b/backend-python/utils/rwkv.py @@ -403,16 +403,12 @@ class TextRWKV(AbstractRWKV): + occurrence[n] * self.penalty_alpha_frequency ) + # comment the codes below to get the same generated results as the official RWKV Gradio if i == 0: for token in self.model_tokens: token = int(token) - for xxx in occurrence: - occurrence[xxx] *= self.penalty_decay if token not in self.AVOID_PENALTY_TOKENS: - if token not in occurrence: - occurrence[token] = 1 - else: - occurrence[token] += 1 + self.adjust_occurrence(occurrence, token) # Model only saw '\n\n' as [187, 187] before, but the tokenizer outputs [535] for it at the end def fix_tokens(self, tokens) -> List[int]: