Update word_language_model/generate.py to remove duplicates, use abc order

Co-authored-by: Dmitry Rogozhkin <[email protected]>

Author: 2 people

word_language_model/generate.py

@@ -55,6 +55,7 @@
 with open(args.checkpoint, 'rb') as f:
     safe_globals = [
         PositionalEncoding,
+        RNNModel,
         TransformerModel,
         torch.nn.functional.relu,
         torch.nn.modules.activation.MultiheadAttention,
@@ -64,15 +65,11 @@
         torch.nn.modules.linear.NonDynamicallyQuantizableLinear,
         torch.nn.modules.normalization.LayerNorm,
         torch.nn.modules.sparse.Embedding,
-        torch.nn.modules.transformer.TransformerEncoder,
-        torch.nn.modules.transformer.TransformerEncoderLayer,
-        RNNModel,
-        torch.nn.modules.dropout.Dropout,
-        torch.nn.modules.linear.Linear,
         torch.nn.modules.rnn.GRU,
         torch.nn.modules.rnn.LSTM,
         torch.nn.modules.rnn.RNN,
-        torch.nn.modules.sparse.Embedding,
+        torch.nn.modules.transformer.TransformerEncoder,
+        torch.nn.modules.transformer.TransformerEncoderLayer,
     ]
     with torch.serialization.safe_globals(safe_globals):
         model = torch.load(f, map_location=device)