Add transformer import to count_occurrences.py

google-research · Mar 11, 2022 · 67351ad · 67351ad
1 parent 50b25c4
commit 67351ad
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/scripts/count_occurrences.py b/scripts/count_occurrences.py
@@ -27,8 +27,10 @@
 
 if args.tokenize:
  if args.tokenizer == 'gpt2':
+ from transformers import GPT2Tokenizer
  tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
  elif args.tokenizer == 't5':
+ from transformers import T5Tokenizer
  tokenizer = T5Tokenizer.from_pretrained('t5-small')
  else:
  raise