From 9ed35477387ffb8f1103a3645a5d4fb226801881 Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Wed, 28 Jun 2023 19:36:15 +0800
Subject: [PATCH] rwkv pip 0.8.0

---
 backend-python/utils/rwkv.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/backend-python/utils/rwkv.py b/backend-python/utils/rwkv.py
index 0f98175..603cbb1 100644
--- a/backend-python/utils/rwkv.py
+++ b/backend-python/utils/rwkv.py
@@ -12,7 +12,6 @@ from routes import state_cache
 
 
 END_OF_TEXT = 0
-END_OF_LINE = 187
 END_OF_LINE_DOUBLE = 535
 
 
@@ -39,12 +38,14 @@ class RWKV:
         self.penalty_alpha_frequency = 0.4
 
         self.interface = ":"
-        if "rwkv_vocab" in tokens_path:
+        if "world" in self.name.lower():
             self.user = "Question"
             self.bot = "Answer"
+            self.END_OF_LINE = 11
         else:
             self.user = "Bob"
             self.bot = "Alice"
+            self.END_OF_LINE = 187
 
         self.AVOID_REPEAT_TOKENS = []
         AVOID_REPEAT = "，：？！"
@@ -86,8 +87,10 @@ The following is a coherent verbose detailed conversation between a girl named {
 
     # Model only saw '\n\n' as [187, 187] before, but the tokenizer outputs [535] for it at the end
     def fix_tokens(self, tokens):
+        if "world" in self.name.lower():
+            return tokens
         if len(tokens) > 0 and tokens[-1] == END_OF_LINE_DOUBLE:
-            tokens = tokens[:-1] + [END_OF_LINE, END_OF_LINE]
+            tokens = tokens[:-1] + [self.END_OF_LINE, self.END_OF_LINE]
         return tokens
 
     def run_rnn(self, _tokens: List[str], newline_adj: int = 0):
@@ -101,7 +104,7 @@ The following is a coherent verbose detailed conversation between a girl named {
             )
             tokens = tokens[self.CHUNK_LEN :]
 
-        out[END_OF_LINE] += newline_adj  # adjust \n probability
+        out[self.END_OF_LINE] += newline_adj  # adjust \n probability
 
         if self.model_tokens[-1] in self.AVOID_REPEAT_TOKENS:
             out[self.model_tokens[-1]] = -999999999
@@ -313,6 +316,8 @@ The following is a coherent verbose detailed conversation between a girl named {
             if token == END_OF_TEXT:
                 yield response, "", prompt_token_len, completion_token_len
                 break
+            for xxx in occurrence:
+                occurrence[xxx] *= 0.996
             if token not in occurrence:
                 occurrence[token] = 1
             else: