Spaces:
Runtime error
Runtime error
File size: 536 Bytes
b4a8aea e03f966 b4a8aea a77a127 b4a8aea e03f966 7846d6e 8f5eeb9 b4a8aea b28599f ccf9460 b4a8aea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
from transformers import AutoTokenizer
from transformers import Tool
class TokenCounterTool(Tool):
name = "token_counter"
description = "This is a tool for counting token used by a prompt. It takes a prompt as input and returns the generated text."
inputs = ["text"]
outputs = ["text"]
def __call__(self, prompt: str):
# token = os.environ['hf']
tokenizer = AutoTokenizer.from_pretrained("lgaalves/gpt2-dolly")
tokens = tokenizer(prompt)["input_ids"]
return f"{len(tokens)}"
|