diff options
| -rw-r--r-- | bicleaner_ai/models.py | 7 |
1 files changed, 2 insertions, 5 deletions
diff --git a/bicleaner_ai/models.py b/bicleaner_ai/models.py index 75dd80a..6c3db5d 100644 --- a/bicleaner_ai/models.py +++ b/bicleaner_ai/models.py @@ -486,8 +486,6 @@ class BCXLMRoberta(BaseModel): self.tokenizer = None self.settings = { - "model_file": "model.tf", - "vocab_file": "vocab", "model": 'jplu/tf-xlm-roberta-base', "batch_size": 16, "maxlen": 150, @@ -537,9 +535,8 @@ class BCXLMRoberta(BaseModel): def load(self): ''' Load fine-tuned model ''' - vocab_file = self.dir + '/' + self.settings["vocab_file"] - self.tokenizer = XLMRobertaTokenizerFast.from_pretrained(vocab_file) - self.model = self.load_model(self.dir+'/'+self.settings["model_file"]) + self.tokenizer = XLMRobertaTokenizerFast.from_pretrained(self.dir) + self.model = self.load_model(self.dir) def softmax_pos_prob(self, x): # Compute softmax probability of the second (positive) class |