Update models/transformer/text_generator.py

models/transformer/text_generator.py

@@ -1,7 +1,7 @@
 from transformers import GPT2LMHeadModel
 from pathlib import Path
 from .utils import modified_tokenizer
-from .constants import CHECKPOINT_PATH
+from .constants import CHECKPOINT_PATH, HF_TOKEN
 
 
 class TextGenerator:
@@ -12,7 +12,7 @@ class TextGenerator:
         """
         model_path = Path(data_path) / model_name
         self.tokenizer = modified_tokenizer(model_path, None, data_path)
-        self.model = GPT2LMHeadModel.from_pretrained(str(model_path), device_map="auto")
+        self.model = GPT2LMHeadModel.from_pretrained(str(model_path), device_map="auto", token=HF_TOKEN)
         self.model.eval()
 
     def generate_text(self,