better

rwkv_pip_package/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rwkv"
-version = "0.8.23"
+version = "0.8.24"
 authors = [
   { name="Bo PENG" },
 ]

rwkv_pip_package/src/rwkv/utils.py

@@ -53,6 +53,9 @@ def decode(self, x):
         return self.tokenizer.decode(x)
 
     def sample_logits(self, logits, temperature=1.0, top_p=0.85, top_k=0):
+        if temperature == 0:
+            temperature = 1.0
+            top_p = 0
         probs = F.softmax(logits.float(), dim=-1)
         top_k = int(top_k)
         # 'privateuseone' is the type of custom devices like `torch_directml.device()`
@@ -109,11 +112,17 @@ def generate(self, ctx, token_count=100, args=PIPELINE_ARGS(), callback=None, st
             all_tokens += [token]
             for xxx in occurrence:
                 occurrence[xxx] *= args.alpha_decay
-            if self.decode([token]) not in ' \r\n\t,.;?!"\':0123456789+-*/=#@$%^&_`~|<>\\()[]{}，。；“”：？！（）【】':
-                if token not in occurrence:
-                    occurrence[token] = 1
-                else:
-                    occurrence[token] += 1
+
+            ttt = self.decode([token])
+            www = 1
+            if ttt in ' \t0123456789':
+                www = 0
+            # elif ttt in '\r\n,.;?!"\':+-*/=#@$%^&_`~|<>\\()[]{}，。；“”：？！（）【】':
+            #     www = 0.5
+            if token not in occurrence:
+                occurrence[token] = www
+            else:
+                occurrence[token] += www
             # print(occurrence) # debug
 
             # output