mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Tokenization improvements

Browse source
This commit is contained in:
oobabooga 2023-09-17 07:01:34 -07:00
parent cd08eb0753
commit ad8ac545a5
5 changed files with 19 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

15
modules/text_generation.py
View file

@ -39,8 +39,7 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False, escap
if generate_func is None:
if shared.model_name == 'None' or shared.model is None:
logger.error("No model is loaded! Select one in the Model tab.")
yield ''
return
raise ValueError('No model is loaded! Select one in the Model tab.')
if shared.model.__class__.__name__ in ['LlamaCppModel', 'RWKVModel', 'ExllamaModel', 'Exllamav2Model', 'CtransformersModel']:
generate_func = generate_reply_custom
@ -106,6 +105,10 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False, escap
def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_length=None):
if shared.tokenizer is None:
logger.error('No tokenizer is loaded')
raise ValueError('No tokenizer is loaded')
if shared.model.__class__.__name__ in ['LlamaCppModel', 'RWKVModel', 'CtransformersModel', 'Exllamav2Model']:
input_ids = shared.tokenizer.encode(str(prompt))
if shared.model.__class__.__name__ not in ['Exllamav2Model']:
@ -133,6 +136,10 @@ def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_lengt
def decode(output_ids, skip_special_tokens=True):
if shared.tokenizer is None:
logger.error('No tokenizer is loaded')
raise ValueError('No tokenizer is loaded')
return shared.tokenizer.decode(output_ids, skip_special_tokens)
@ -146,11 +153,11 @@ def get_encoded_length(prompt):
def get_token_ids(prompt):
tokens = encode(prompt)[0]
decoded_tokens = [shared.tokenizer.decode(i) for i in tokens]
decoded_tokens = [shared.tokenizer.decode([i]) for i in tokens]
output = ''
for row in list(zip(tokens, decoded_tokens)):
output += f"{str(int(row[0])).ljust(5)} - {repr(row[1])[1:-1]}\n"
output += f"{str(int(row[0])).ljust(5)} - {repr(row[1])}\n"
return output