token-counter-tool / token_counter.py
Chris4K's picture
Update token_counter.py
a77a127
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)}"