Chris4K commited on
Commit
82863a2
1 Parent(s): 4127d50

Update token_counter.py

Browse files
Files changed (1) hide show
  1. token_counter.py +5 -7
token_counter.py CHANGED
@@ -1,15 +1,13 @@
1
  from transformers import AutoTokenizer
2
  from transformers import Tool
3
 
4
- class TokenCounterTool(Tool):
5
- name = "token_counter"
6
- description = "This is a tool for counting token used by a prompt. It takes a prompt as input and returns the generated text."
7
  inputs = ["text"]
8
  outputs = ["text"]
9
 
10
  def __call__(self, prompt: str):
11
- # token = os.environ['hf']
12
- tokenizer = AutoTokenizer.from_pretrained("lgaalves/gpt2-dolly")
13
- tokens = tokenizer(prompt)["input_ids"]
14
- return f"{len(tokens)}"
15
 
 
1
  from transformers import AutoTokenizer
2
  from transformers import Tool
3
 
4
+ class WordCounterTool(Tool):
5
+ name = "word_counter"
6
+ description = "This is a tool for counting word of an input."
7
  inputs = ["text"]
8
  outputs = ["text"]
9
 
10
  def __call__(self, prompt: str):
11
+ words = str.split(" ")
12
+ return len(words)
 
 
13