theluyuan/RWKV-Runner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

custom tokenizer .txt support

Browse Source
This commit is contained in:
josc146 2023-09-18 17:20:55 +08:00
parent f8388a0527
commit 5e5e1e9651
1 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
backend-python/rwkv_pip/utils.py vendored
View File

@ -33,7 +33,7 @@ class PIPELINE_ARGS:
class PIPELINE: class PIPELINE:
def __init__(self, model, WORD_NAME): def __init__(self, model, WORD_NAME: str):
self.model = model self.model = model
if WORD_NAME == "cl100k_base": if WORD_NAME == "cl100k_base":
import tiktoken import tiktoken
@ -46,6 +46,12 @@ class PIPELINE:
self.tokenizer = TRIE_TOKENIZER( self.tokenizer = TRIE_TOKENIZER(
os.path.dirname(os.path.abspath(__file__)) + "/rwkv_vocab_v20230424.txt" os.path.dirname(os.path.abspath(__file__)) + "/rwkv_vocab_v20230424.txt"
) )
else:
if WORD_NAME.endswith(".txt"):
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rwkv_tokenizer import TRIE_TOKENIZER
self.tokenizer = TRIE_TOKENIZER(WORD_NAME)
else: else:
from tokenizers import Tokenizer from tokenizers import Tokenizer