Annorita's picture
add model option
ae7ad9f
raw
history blame
719 Bytes
from transformers import AutoTokenizer
import itertools
def get_color():
colors = ['#df7b55', '#2c7482', '#2c8234', '#5581df', '#822c63','#b355df']
return itertools.cycle(colors)
def get_res(model_name, input_sentence, single_print=True):
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
color_iterator = get_color()
out = tokenizer.encode(input_sentence, add_special_tokens=False)
token_num = len(out)
w = [ f'<span style="font-size:1.25em;background-color:{next(color_iterator)}">{tokenizer.decode(x)}</span>' for x in out ]
res = ''.join(w)
if single_print:
print(res + str(token_num))
else:
return res, token_num