extend args accpeted by Embedding

libmultilabel/nn/networks/kim_cnn.py

@@ -28,7 +28,7 @@ def __init__(
         activation="relu",
     ):
         super(KimCNN, self).__init__()
-        self.embedding = Embedding(embed_vecs, embed_dropout)
+        self.embedding = Embedding(embed_vecs, dropout=embed_dropout)
         self.encoder = CNNEncoder(
             embed_vecs.shape[1], filter_sizes, num_filter_per_size, activation, post_encoder_dropout, num_pool=1
         )

libmultilabel/nn/networks/labelwise_attention_networks.py

@@ -27,7 +27,7 @@ class LabelwiseAttentionNetwork(ABC, nn.Module):
 
     def __init__(self, embed_vecs, num_classes, embed_dropout, encoder_dropout, post_encoder_dropout, hidden_dim):
         super(LabelwiseAttentionNetwork, self).__init__()
-        self.embedding = Embedding(embed_vecs, embed_dropout)
+        self.embedding = Embedding(embed_vecs, dropout=embed_dropout)
         self.encoder = self._get_encoder(embed_vecs.shape[1], encoder_dropout, post_encoder_dropout)
         self.attention = self._get_attention()
         self.output = LabelwiseLinearOutput(hidden_dim, num_classes)

libmultilabel/nn/networks/modules.py

@@ -11,12 +11,14 @@ class Embedding(nn.Module):
 
     Args:
         embed_vecs (torch.Tensor): The pre-trained word vectors of shape (vocab_size, embed_dim).
+        freeze (bool): If True, the tensor does not get updated in the learning process.
+            Equivalent to embedding.weight.requires_grad = False. Default: False.
         dropout (float): The dropout rate of the word embedding. Defaults to 0.2.
     """
 
-    def __init__(self, embed_vecs, dropout=0.2):
+    def __init__(self, embed_vecs, freeze=False, dropout=0.2):
         super(Embedding, self).__init__()
-        self.embedding = nn.Embedding.from_pretrained(embed_vecs, freeze=False, padding_idx=0)
+        self.embedding = nn.Embedding.from_pretrained(embed_vecs, freeze=freeze, padding_idx=0)
         self.dropout = nn.Dropout(dropout)
 
     def forward(self, input):

libmultilabel/nn/networks/xml_cnn.py

@@ -33,7 +33,7 @@ def __init__(
         activation="relu",
     ):
         super(XMLCNN, self).__init__()
-        self.embedding = Embedding(embed_vecs, embed_dropout)
+        self.embedding = Embedding(embed_vecs, dropout=embed_dropout)
         self.encoder = CNNEncoder(embed_vecs.shape[1], filter_sizes, num_filter_per_size, activation, num_pool=num_pool)
         total_output_size = len(filter_sizes) * num_filter_per_size * num_pool
         self.linear1 = nn.Linear(total_output_size, hidden_dim)