Pass the label_decoder to the base class

stanfordnlp · Sep 18, 2024 · 3dfb5ad · 3dfb5ad
1 parent e474ee2
commit 3dfb5ad
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 7 deletions.
diff --git a/stanza/models/lemma_classifier/base_model.py b/stanza/models/lemma_classifier/base_model.py
@@ -19,9 +19,10 @@
 logger = logging.getLogger('stanza.lemmaclassifier')
 
 class LemmaClassifier(ABC, nn.Module):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, label_decoder, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+        self.label_decoder = label_decoder
         self.unsaved_modules = []
 
     def add_unsaved_module(self, name, module):

diff --git a/stanza/models/lemma_classifier/lstm_model.py b/stanza/models/lemma_classifier/lstm_model.py
@@ -41,7 +41,7 @@ def __init__(self, model_args, output_dim, pt_embedding, label_decoder, upos_to_
         Raises:
             FileNotFoundError: if the forward or backward charlm file cannot be found.
         """
-        super(LemmaClassifierLSTM, self).__init__()
+        super(LemmaClassifierLSTM, self).__init__(label_decoder)
         self.model_args = model_args
 
         self.hidden_dim = model_args['hidden_dim']
@@ -63,9 +63,6 @@ def __init__(self, model_args, output_dim, pt_embedding, label_decoder, upos_to_
 
         self.input_size += self.embedding_dim
 
-        # TODO: pass this up to the parent class
-        self.label_decoder = label_decoder
-
         # Optionally, include charlm embeddings
         self.use_charlm = use_charlm
 

diff --git a/stanza/models/lemma_classifier/transformer_model.py b/stanza/models/lemma_classifier/transformer_model.py
@@ -28,7 +28,7 @@ def __init__(self, model_args: dict, output_dim: int, transformer_name: str, lab
             transformer_name (str): name of the HF transformer to use
             label_decoder (dict): a map of the labels available to the model
         """
-        super(LemmaClassifierWithTransformer, self).__init__()
+        super(LemmaClassifierWithTransformer, self).__init__(label_decoder)
         self.model_args = model_args
 
         # Choose transformer
@@ -45,7 +45,6 @@ def __init__(self, model_args: dict, output_dim: int, transformer_name: str, lab
             nn.ReLU(),
             nn.Linear(64, output_dim)
         )
-        self.label_decoder = label_decoder
 
     def get_save_dict(self):
         save_dict = {