Merge pull request #163 from zalandoresearch/GH-162-memory-optimizations

GH-162: integration tests

flair/embeddings.py

@@ -462,13 +462,16 @@ def embedding_length(self) -> int:
 
     def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
 
-        cache_path = '{}-tmp-cache.sqllite'.format(self.name) if self.cache_directory is None else os.path.join(
-            self.cache_directory, '{}-tmp-cache.sqllite'.format(os.path.basename(self.name)))
-
-        # by default, use_cache is false (for older pre-trained models TODO: remove in version 0.4)
-        if 'cache' not in self.__dict__ or 'cache_directory' not in self.__dict__ or not os.path.exists(cache_path):
+        # this whole block is for compatibility with older serialized models  TODO: remove in version 0.4
+        if 'cache' not in self.__dict__ or 'cache_directory' not in self.__dict__:
             self.use_cache = False
             self.cache_directory = None
+        else:
+            cache_path = '{}-tmp-cache.sqllite'.format(self.name) if not self.cache_directory else os.path.join(
+                self.cache_directory, '{}-tmp-cache.sqllite'.format(os.path.basename(self.name)))
+            if not os.path.exists(cache_path):
+                self.use_cache = False
+                self.cache_directory = None
 
         # if cache is used, try setting embeddings from cache first
         if self.use_cache:
@@ -553,15 +556,14 @@ def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
 
 class DocumentMeanEmbeddings(DocumentEmbeddings):
 
-    def __init__(self, word_embeddings: List[TokenEmbeddings]):
+    def __init__(self, token_embeddings: List[TokenEmbeddings]):
         """The constructor takes a list of embeddings to be combined."""
         super().__init__()
 
-        self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=word_embeddings)
+        self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=token_embeddings)
         self.name: str = 'document_mean'
 
-        self.__embedding_length: int = 0
-        self.__embedding_length = self.embeddings.embedding_length
+        self.__embedding_length: int = self.embeddings.embedding_length
 
         if torch.cuda.is_available():
             self.cuda()
@@ -631,18 +633,12 @@ def __init__(self,
         """
         super().__init__()
 
-        self.embeddings: List[TokenEmbeddings] = token_embeddings
-
-        # IMPORTANT: add embeddings as torch modules
-        for i, embedding in enumerate(self.embeddings):
-            self.add_module('token_embedding_{}'.format(i), embedding)
+        self.embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=token_embeddings)
 
         self.reproject_words = reproject_words
         self.bidirectional = bidirectional
 
-        self.length_of_all_token_embeddings = 0
-        for token_embedding in self.embeddings:
-            self.length_of_all_token_embeddings += token_embedding.embedding_length
+        self.length_of_all_token_embeddings: int = self.embeddings.embedding_length
 
         self.name = 'document_lstm'
         self.static_embeddings = False
@@ -691,8 +687,7 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
 
         sentences.sort(key=lambda x: len(x), reverse=True)
 
-        for token_embedding in self.embeddings:
-            token_embedding.embed(sentences)
+        self.embeddings.embed(sentences)
 
         # first, sort sentences by number of tokens
         longest_token_sequence_in_batch: int = len(sentences[0])

tests/test_model_integration.py

@@ -0,0 +1,252 @@
+import os
+import shutil
+
+from flair.data import Sentence
+from flair.data_fetcher import NLPTaskDataFetcher, NLPTask
+from flair.embeddings import WordEmbeddings, CharLMEmbeddings, DocumentLSTMEmbeddings, TokenEmbeddings
+from flair.models import SequenceTagger, TextClassifier
+from flair.trainers import SequenceTaggerTrainer, TextClassifierTrainer
+
+
+def test_train_load_use_tagger():
+
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.FASHION)
+    tag_dictionary = corpus.make_tag_dictionary('ner')
+
+    embeddings = WordEmbeddings('glove')
+
+    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+                                            embeddings=embeddings,
+                                            tag_dictionary=tag_dictionary,
+                                            tag_type='ner',
+                                            use_crf=False)
+
+    # initialize trainer
+    trainer: SequenceTaggerTrainer = SequenceTaggerTrainer(tagger, corpus, test_mode=True)
+
+    trainer.train('./results', learning_rate=0.1, mini_batch_size=2, max_epochs=3)
+
+    loaded_model: SequenceTagger = SequenceTagger.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')
+
+
+def test_train_charlm_load_use_tagger():
+
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.FASHION)
+    tag_dictionary = corpus.make_tag_dictionary('ner')
+
+    embeddings = CharLMEmbeddings('news-forward-fast')
+
+    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+                                            embeddings=embeddings,
+                                            tag_dictionary=tag_dictionary,
+                                            tag_type='ner',
+                                            use_crf=False)
+
+    # initialize trainer
+    trainer: SequenceTaggerTrainer = SequenceTaggerTrainer(tagger, corpus, test_mode=True)
+
+    trainer.train('./results', learning_rate=0.1, mini_batch_size=2, max_epochs=3)
+
+    loaded_model: SequenceTagger = SequenceTagger.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')
+
+
+def test_train_charlm_changed_chache_load_use_tagger():
+
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.FASHION)
+    tag_dictionary = corpus.make_tag_dictionary('ner')
+
+    # make a temporary cache directory that we remove afterwards
+    os.makedirs('./results/cache/', exist_ok=True)
+    embeddings = CharLMEmbeddings('news-forward-fast', cache_directory='./results/cache/')
+
+    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+                                            embeddings=embeddings,
+                                            tag_dictionary=tag_dictionary,
+                                            tag_type='ner',
+                                            use_crf=False)
+
+    # initialize trainer
+    trainer: SequenceTaggerTrainer = SequenceTaggerTrainer(tagger, corpus, test_mode=True)
+
+    trainer.train('./results', learning_rate=0.1, mini_batch_size=2, max_epochs=3)
+
+    # remove the cache directory
+    shutil.rmtree('./results/cache')
+
+    loaded_model: SequenceTagger = SequenceTagger.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')
+
+
+def test_train_charlm_nochache_load_use_tagger():
+
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.FASHION)
+    tag_dictionary = corpus.make_tag_dictionary('ner')
+
+    embeddings = CharLMEmbeddings('news-forward-fast', use_cache=False)
+
+    tagger: SequenceTagger = SequenceTagger(hidden_size=256,
+                                            embeddings=embeddings,
+                                            tag_dictionary=tag_dictionary,
+                                            tag_type='ner',
+                                            use_crf=False)
+
+    # initialize trainer
+    trainer: SequenceTaggerTrainer = SequenceTaggerTrainer(tagger, corpus, test_mode=True)
+
+    trainer.train('./results', learning_rate=0.1, mini_batch_size=2, max_epochs=3)
+
+    loaded_model: SequenceTagger = SequenceTagger.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')
+
+
+def test_load_use_serialized_tagger():
+
+    loaded_model: SequenceTagger = SequenceTagger.load('ner')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+
+def test_train_load_use_classifier():
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.IMDB)
+    label_dict = corpus.make_label_dictionary()
+
+    glove_embedding: WordEmbeddings = WordEmbeddings('en-glove')
+    document_embeddings: DocumentLSTMEmbeddings = DocumentLSTMEmbeddings([glove_embedding], 128, 1, False, 64, False,
+                                                                         False)
+
+    model = TextClassifier(document_embeddings, label_dict, False)
+
+    trainer = TextClassifierTrainer(model, corpus, label_dict, False)
+    trainer.train('./results', max_epochs=2)
+
+    sentence = Sentence("Berlin is a really nice city.")
+
+    for s in model.predict(sentence):
+        for l in s.labels:
+            assert (l.value is not None)
+            assert (0.0 <= l.score <= 1.0)
+            assert (type(l.score) is float)
+
+        loaded_model = TextClassifier.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')
+
+
+def test_train_charlm_load_use_classifier():
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.IMDB)
+    label_dict = corpus.make_label_dictionary()
+
+    glove_embedding: TokenEmbeddings = CharLMEmbeddings('news-forward-fast')
+    document_embeddings: DocumentLSTMEmbeddings = DocumentLSTMEmbeddings([glove_embedding], 128, 1, False, 64, False,
+                                                                         False)
+
+    model = TextClassifier(document_embeddings, label_dict, False)
+
+    trainer = TextClassifierTrainer(model, corpus, label_dict, False)
+    trainer.train('./results', max_epochs=2)
+
+    sentence = Sentence("Berlin is a really nice city.")
+
+    for s in model.predict(sentence):
+        for l in s.labels:
+            assert (l.value is not None)
+            assert (0.0 <= l.score <= 1.0)
+            assert (type(l.score) is float)
+
+        loaded_model = TextClassifier.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')
+
+
+def test_train_charlm__nocache_load_use_classifier():
+    corpus = NLPTaskDataFetcher.fetch_data(NLPTask.IMDB)
+    label_dict = corpus.make_label_dictionary()
+
+    glove_embedding: TokenEmbeddings = CharLMEmbeddings('news-forward-fast', use_cache=False)
+    document_embeddings: DocumentLSTMEmbeddings = DocumentLSTMEmbeddings([glove_embedding], 128, 1, False, 64,
+                                                                         False,
+                                                                         False)
+
+    model = TextClassifier(document_embeddings, label_dict, False)
+
+    trainer = TextClassifierTrainer(model, corpus, label_dict, False)
+    trainer.train('./results', max_epochs=2)
+
+    sentence = Sentence("Berlin is a really nice city.")
+
+    for s in model.predict(sentence):
+        for l in s.labels:
+            assert (l.value is not None)
+            assert (0.0 <= l.score <= 1.0)
+            assert (type(l.score) is float)
+
+        loaded_model = TextClassifier.load_from_file('./results/final-model.pt')
+
+    sentence = Sentence('I love Berlin')
+    sentence_empty = Sentence('       ')
+
+    loaded_model.predict(sentence)
+    loaded_model.predict([sentence, sentence_empty])
+    loaded_model.predict([sentence_empty])
+
+    # clean up results directory
+    shutil.rmtree('./results')