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)}"