LMartinezEXEX commited on
Commit
c05d0d5
·
1 Parent(s): 44f418e

Updated model and tokenizer HugingFace loading with Auto... clases

Browse files
Files changed (1) hide show
  1. modules/module_languageModel.py +4 -4
modules/module_languageModel.py CHANGED
@@ -1,5 +1,5 @@
1
  # --- Imports libs ---
2
- from transformers import BertForMaskedLM, BertTokenizer
3
 
4
  class LanguageModel:
5
  def __init__(
@@ -7,9 +7,9 @@ class LanguageModel:
7
  model_name: str
8
  ) -> None:
9
 
10
- print("Download language model...")
11
- self.__tokenizer = BertTokenizer.from_pretrained(model_name)
12
- self.__model = BertForMaskedLM.from_pretrained(model_name, return_dict=True)
13
 
14
  def initTokenizer(
15
  self
 
1
  # --- Imports libs ---
2
+ from transformers import AutoTokenizer, AutoModelForMaskedLM
3
 
4
  class LanguageModel:
5
  def __init__(
 
7
  model_name: str
8
  ) -> None:
9
 
10
+ print("Downloading language model...")
11
+ self.__tokenizer = AutoTokenizer.from_pretrained(model_name)
12
+ self.__model = AutoModelForMaskedLM.from_pretrained(model_name)
13
 
14
  def initTokenizer(
15
  self