| from transformers import AutoTokenizer | |
| from transformers import Tool | |
| class TokenCounterTool(Tool): | |
| name = "text_generator" | |
| 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(input_text)["input_ids"] | |
| return f"{len(tokens)}" | |