hotfix

NeuralTextGenerator.py

@@ -20,7 +20,7 @@ DEFAULT_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 class BertTextGenerator:
-    def __init__(self, model_version, device=DEFAULT_DEVICE, use_apex=APEX_AVAILABLE, use_fast=True,
+    def __init__(self, model_version, tokenizer, device=DEFAULT_DEVICE, use_apex=APEX_AVAILABLE, use_fast=True,
                  do_basic_tokenize=True):
         """
         Wrapper of a BERT model from AutoModelForMaskedLM from huggingfaces.
@@ -47,7 +47,7 @@ class BertTextGenerator:
             self.model, optimizer = amp.initialize(self.model, optimizer, opt_level="O2", keep_batchnorm_fp32=True,
                                                    loss_scale="dynamic")
 
-        self.tokenizer = AutoTokenizer.from_pretrained(model_version, do_lower_case="uncased" in model_version,
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer, do_lower_case="uncased" in model_version,
                                                        use_fast=use_fast,
                                                        do_basic_tokenize=do_basic_tokenize)  # added to avoid splitting of unused tokens
         self.num_attention_masks = len(self.model.base_model.base_model.encoder.layer)

app.py

@@ -2,13 +2,13 @@ import gradio as gr
 from NeuralTextGenerator import BertTextGenerator
 
 model_name = "cardiffnlp/twitter-xlm-roberta-base" #"dbmdz/bert-base-italian-uncased"
-en_model = BertTextGenerator(model_name)
+en_model = BertTextGenerator(model_name, tokenizer='xlm-roberta')
 
 finetunned_BERT_model_name = "JuanJoseMV/BERT_text_gen"
-finetunned_BERT_en_model = BertTextGenerator(finetunned_BERT_model_name)
+finetunned_BERT_en_model = BertTextGenerator(finetunned_BERT_model_name, tokenizer='bert-base-uncased')
 
 finetunned_RoBERTa_model_name = "JuanJoseMV/XLM_RoBERTa_text_gen"
-finetunned_RoBERTa_en_model = BertTextGenerator(finetunned_RoBERTa_model_name)
+finetunned_RoBERTa_en_model = BertTextGenerator(finetunned_RoBERTa_model_name, tokenizer='xlm-roberta')
 
 special_tokens = [
     '[POSITIVE-0]',