adopt transformers 4.x API

makcedward · Dec 11, 2020 · 935e68a · 935e68a
1 parent f40cbf3
commit 935e68a
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 5 deletions.
diff --git a/nlpaug/model/lang_models/bert.py b/nlpaug/model/lang_models/bert.py
@@ -49,7 +49,13 @@ def is_skip_candidate(self, candidate):
  return candidate.startswith(self.SUBWORD_PREFIX)
 
  def token2id(self, token):
- return self.tokenizer._convert_token_to_id(token)
+ # Iseue 181: TokenizerFast have convert_tokens_to_ids but not convert_tokens_to_id
+ if 'TokenizerFast' in self.tokenizer.__class__.__name__:
+ # New transformers API
+ return self.tokenizer.convert_tokens_to_ids(token)
+ else:
+ # Old transformers API
+ return self.tokenizer._convert_token_to_id(token)
 
  def id2token(self, _id):
  return self.tokenizer._convert_id_to_token(_id)

diff --git a/nlpaug/model/lang_models/distilbert.py b/nlpaug/model/lang_models/distilbert.py
@@ -51,7 +51,13 @@ def is_skip_candidate(self, candidate):
  return candidate[:2] == self.SUBWORD_PREFIX
 
  def token2id(self, token):
- return self.tokenizer._convert_token_to_id(token)
+ # Iseue 181: TokenizerFast have convert_tokens_to_ids but not convert_tokens_to_id
+ if 'TokenizerFast' in self.tokenizer.__class__.__name__:
+ # New transformers API
+ return self.tokenizer.convert_tokens_to_ids(token)
+ else:
+ # Old transformers API
+ return self.tokenizer._convert_token_to_id(token)
 
  def id2token(self, _id):
  return self.tokenizer._convert_id_to_token(_id)

diff --git a/nlpaug/model/lang_models/roberta.py b/nlpaug/model/lang_models/roberta.py
@@ -16,7 +16,7 @@ class Roberta(LanguageModels):
  START_TOKEN = '<s>'
  SEPARATOR_TOKEN = '</s>'
  MASK_TOKEN = '<mask>'
- PAD_TOKEN = '<pad>',
+ PAD_TOKEN = '<pad>'
  UNKNOWN_TOKEN = '<unk>'
  SUBWORD_PREFIX = 'Ġ'
 
@@ -48,7 +48,13 @@ def get_max_num_token(self):
  return self.model.config.max_position_embeddings - 2 * 5
 
  def token2id(self, token):
- return self.tokenizer._convert_token_to_id(token)
+ # Iseue 181: TokenizerFast have convert_tokens_to_ids but not convert_tokens_to_id
+ if 'TokenizerFast' in self.tokenizer.__class__.__name__:
+ # New transformers API
+ return self.tokenizer.convert_tokens_to_ids(token)
+ else:
+ # Old transformers API
+ return self.tokenizer._convert_token_to_id(token)
 
  def id2token(self, _id):
  return self.tokenizer._convert_id_to_token(_id)

diff --git a/nlpaug/model/lang_models/xlnet.py b/nlpaug/model/lang_models/xlnet.py
@@ -67,7 +67,13 @@ def get_max_num_token(self):
  return 500
 
  def token2id(self, token):
- return self.tokenizer._convert_token_to_id(token)
+ # Iseue 181: TokenizerFast have convert_tokens_to_ids but not convert_tokens_to_id
+ if 'TokenizerFast' in self.tokenizer.__class__.__name__:
+ # New transformers API
+ return self.tokenizer.convert_tokens_to_ids(token)
+ else:
+ # Old transformers API
+ return self.tokenizer._convert_token_to_id(token)
 
  def id2token(self, _id):
  return self.tokenizer._convert_id_to_token(_id)