Merge pull request #609 from duesenfranz/reduce-ram-usage

Only initialize counterfitted_GLOVE_embedding when needed, massively decreasing ram usage
QData · Mar 20, 2022 · 4c91157 · 4c91157
2 parents 11782c4 + 92374d3
commit 4c91157
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 10 deletions.
diff --git a/textattack/constraints/semantics/sentence_encoders/thought_vector.py b/textattack/constraints/semantics/sentence_encoders/thought_vector.py
@@ -19,9 +19,9 @@ class ThoughtVector(SentenceEncoder):
         word_embedding (textattack.shared.AbstractWordEmbedding): The word embedding to use
     """
 
-    def __init__(
-        self, embedding=WordEmbedding.counterfitted_GLOVE_embedding(), **kwargs
-    ):
+    def __init__(self, embedding=None, **kwargs):
+        if embedding is None:
+            embedding = WordEmbedding.counterfitted_GLOVE_embedding()
         if not isinstance(embedding, AbstractWordEmbedding):
             raise ValueError(
                 "`embedding` object must be of type `textattack.shared.AbstractWordEmbedding`."

diff --git a/textattack/constraints/semantics/word_embedding_distance.py b/textattack/constraints/semantics/word_embedding_distance.py
@@ -24,14 +24,16 @@ class WordEmbeddingDistance(Constraint):
 
     def __init__(
         self,
-        embedding=WordEmbedding.counterfitted_GLOVE_embedding(),
+        embedding=None,
         include_unknown_words=True,
         min_cos_sim=None,
         max_mse_dist=None,
         cased=False,
         compare_against_original=True,
     ):
         super().__init__(compare_against_original)
+        if embedding is None:
+            embedding = WordEmbedding.counterfitted_GLOVE_embedding()
         self.include_unknown_words = include_unknown_words
         self.cased = cased
 

diff --git a/textattack/transformations/word_swaps/word_swap_embedding.py b/textattack/transformations/word_swaps/word_swap_embedding.py
@@ -28,13 +28,10 @@ class WordSwapEmbedding(WordSwap):
     >>> augmenter.augment(s)
     """
 
-    def __init__(
-        self,
-        max_candidates=15,
-        embedding=WordEmbedding.counterfitted_GLOVE_embedding(),
-        **kwargs
-    ):
+    def __init__(self, max_candidates=15, embedding=None, **kwargs):
         super().__init__(**kwargs)
+        if embedding is None:
+            embedding = WordEmbedding.counterfitted_GLOVE_embedding()
         self.max_candidates = max_candidates
         if not isinstance(embedding, AbstractWordEmbedding):
             raise ValueError(