fix import error

GRAAL-Research · Nov 14, 2023 · 217abbc · 217abbc
1 parent 03f4629
commit 217abbc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/training/metrics/model_salience.py b/src/training/metrics/model_salience.py
@@ -2,7 +2,7 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-import utils_masking
+from utils_masking import NonStopMasker
 
 
 def unfold(sent_toks, make_tensor=True):
@@ -27,7 +27,7 @@ def __init__(
         if self.tokenizer.eos_token_id is None:
             self.eos_token_id = 0
 
-        self.masking_model = utils_masking.NonStopMasker()
+        self.masking_model = NonStopMasker()
         self.masking_model.register_tokenizer(self.tokenizer)
 
         self.vocab_size = self.tokenizer.vocab_size