text-generation-webui/extensions/openai/tokens.py

27 lines
482 B
Python
Raw Normal View History

2023-09-16 05:11:16 +02:00
from modules.text_generation import decode, encode
2023-07-12 20:33:25 +02:00
def token_count(prompt):
tokens = encode(prompt)[0]
return {
'length': len(tokens)
}
def token_encode(input):
tokens = encode(input)[0]
if tokens.__class__.__name__ in ['Tensor', 'ndarray']:
tokens = tokens.tolist()
return {
'tokens': tokens,
'length': len(tokens),
}
def token_decode(tokens):
output = decode(tokens)
return {
'text': output
}