Add world support

app.py

@@ -13,6 +13,7 @@ desc = f'''链接：<a href='https://colab.research.google.com/drive/1J1gLMMMA8G
 
 parser = argparse.ArgumentParser(prog = 'ChatGal RWKV')
 parser.add_argument('--share',action='store_true') 
+parser.add_argument("--world",type=bool, default=False)
 parser.add_argument('--ckpt',type=str,default="rwkv-loramerge-0426-v2-4096-epoch11.pth") 
 parser.add_argument('--model_path',type=str,default=None,help="local model path") 
 parser.add_argument('--lora', type=str, default=None, help='lora checkpoint path')
@@ -40,7 +41,8 @@ if torch.cuda.is_available() and torch.cuda.device_count()>0:
 else:
     model = RWKV(model=model_path, strategy='cpu bf16',**lora_kwargs)
 from utils import PIPELINE, PIPELINE_ARGS
-pipeline = PIPELINE(model, "20B_tokenizer.json")
+tokenizer_file = "rwkv_vocab_v20230424" if args.world else "20B_tokenizer.json"
+pipeline = PIPELINE(model, tokenizer_file)
 
 def infer(
         ctx,

rwkv_tokenizer.py

@@ -0,0 +1,103 @@
+########################################################################################################
+# The RWKV Language Model - https://github.com/BlinkDL/RWKV-LM
+########################################################################################################
+
+class TRIE:
+    __slots__ = tuple("ch,to,values,front".split(","))
+    to:list
+    values:set
+    def __init__(self, front=None, ch=None):
+        self.ch = ch
+        self.to = [None for ch in range(256)]
+        self.values = set()
+        self.front = front
+
+    def __repr__(self):
+        fr = self
+        ret = []
+        while(fr!=None):
+            if(fr.ch!=None):
+                ret.append(fr.ch)
+            fr = fr.front
+        return "<TRIE %s %s>"%(ret[::-1], self.values)
+    
+    def add(self, key:bytes, idx:int=0, val=None):
+        if(idx == len(key)):
+            if(val is None):
+                val = key
+            self.values.add(val)
+            return self
+        ch = key[idx]
+        if(self.to[ch] is None):
+            self.to[ch] = TRIE(front=self, ch=ch)
+        return self.to[ch].add(key, idx=idx+1, val=val)
+    
+    def find_longest(self, key:bytes, idx:int=0):
+        u:TRIE = self
+        ch:int = key[idx]
+        
+        while(u.to[ch] is not None):
+            u = u.to[ch]
+            idx += 1
+            if(u.values):
+                ret = idx, u, u.values
+            if(idx==len(key)):
+                break
+            ch = key[idx]
+        return ret
+
+class TRIE_TOKENIZER():
+    def __init__(self, file_name):
+        self.idx2token = {}
+        sorted = [] # must be already sorted
+        with open(file_name, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+        for l in lines:
+            idx = int(l[:l.index(' ')])
+            x = eval(l[l.index(' '):l.rindex(' ')])
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+            assert len(x) == int(l[l.rindex(' '):])
+            sorted += [x]
+            self.idx2token[idx] = x
+
+        self.token2idx = {}
+        for k,v in self.idx2token.items():
+            self.token2idx[v] = int(k)
+
+        self.root = TRIE()
+        for t, i in self.token2idx.items():
+            _ = self.root.add(t, val=(t, i))
+
+    def encodeBytes(self, src:bytes):
+        idx:int = 0
+        tokens = []
+        while (idx < len(src)):
+            _idx:int = idx
+            idx, _, values = self.root.find_longest(src, idx)
+            assert(idx != _idx)
+            _, token = next(iter(values))            
+            tokens.append(token)
+        return tokens
+
+    def decodeBytes(self, tokens):
+        return b''.join(map(lambda i: self.idx2token[i], tokens))
+
+    def encode(self, src):
+        return self.encodeBytes(src.encode("utf-8"))
+
+    def decode(self, tokens):
+        try:
+            return self.decodeBytes(tokens).decode('utf-8')
+        except:
+            return '\ufffd' # bad utf-8
+
+    def printTokens(self, tokens):
+        for i in tokens:
+            s = self.idx2token[i]
+            try:
+                s = s.decode('utf-8')
+            except:
+                pass
+            print(f'{repr(s)}{i}', end=' ')
+        print()

utils.py

@@ -22,6 +22,9 @@ class PIPELINE():
         if WORD_NAME == 'cl100k_base':
             import tiktoken
             self.tokenizer = tiktoken.get_encoding(WORD_NAME)
+        elif WORD_NAME == 'rwkv_vocab_v20230424':
+            from rwkv_tokenizer import TRIE_TOKENIZER
+            self.tokenizer = TRIE_TOKENIZER(f'./{WORD_NAME}.txt')        
         else:
             from tokenizers import Tokenizer
             self.tokenizer = Tokenizer.from_file(WORD_NAME)
@@ -37,10 +40,14 @@ class PIPELINE():
         return context
 
     def encode(self, x):
-        if 'tiktoken' in str(type(self.tokenizer)):
-            return self.tokenizer.encode(x)
-        else:
-            return self.tokenizer.encode(x).ids
+        # if 'tiktoken' in str(type(self.tokenizer)):
+        #     return self.tokenizer.encode(x)
+        # else:
+        #     return self.tokenizer.encode(x).ids
+        encoded = self.tokenizer.encode(x)
+        if hasattr(encoded,"ids"):
+            encoded = encoded.ids
+        return encoded
     
     def decode(self, x):
         return self.tokenizer.decode(x)