Repository: flairNLP/flair

Issue:
Keeping the order of sentences in the batch in DocumentRNNEmbedding
In DocumentRNNEmbedding the sentences in the batch are permuted (sorted) to take advantage of PyTorch's `pack_padded_sequence` and `pad_packed_sequence` require sentences ordered by length. In the rest of the code sentences are kept in the permuted order, so the outputs will also be permuted. This is fine when the order of the sentences in the batch doesn't matter or when labels are assigned to Sentence objects (like in Flair). However, for some tasks, where association between sentence and label (or an object in general) is not 1-1, it is assumed that order of sentences in the batch does not change.

Therefore, it would be good to keep the order of sentences in the batch. For that we'd need to keep sorting indices so we can undo the permutation due to sorting, and we can do that just after the RNN embeddings are calculated.


Discussion:

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -2122,6 +2122,8 @@ def __init__(
 
         self.to(flair.device)
 
+        self.eval()
+
     @property
     def embedding_length(self) -> int:
         return self.__embedding_length
@@ -2135,11 +2137,17 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
 
         self.rnn.zero_grad()
 
-        sentences.sort(key=lambda x: len(x), reverse=True)
+        # the permutation that sorts the sentences by length, descending
+        sort_perm = np.argsort([len(s) for s in sentences])[::-1]
+
+        # the inverse permutation that restores the input order; it's an index tensor therefore LongTensor
+        sort_invperm = np.argsort(sort_perm)
+
+        # sort sentences by number of tokens
+        sentences = [sentences[i] for i in sort_perm]
 
         self.embeddings.embed(sentences)
 
-        # first, sort sentences by number of tokens
         longest_token_sequence_in_batch: int = len(sentences[0])
 
         all_sentence_tensors = []
@@ -2211,6 +2219,9 @@ def embed(self, sentences: Union[List[Sentence], Sentence]):
             sentence = sentences[sentence_no]
             sentence.set_embedding(self.name, embedding)
 
+        # restore original order of sentences in the batch
+        sentences = [sentences[i] for i in sort_invperm]
+
     def _add_embeddings_internal(self, sentences: List[Sentence]):
         pass
 


Test patch:
diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py
--- a/tests/test_embeddings.py
+++ b/tests/test_embeddings.py
@@ -1,4 +1,5 @@
 import pytest
+import torch
 
 from flair.embeddings import (
     WordEmbeddings,
@@ -25,6 +26,22 @@ def test_loading_not_existing_char_lm_embedding():
         FlairEmbeddings("other")
 
 
+def test_keep_batch_order():
+    sentence, glove, charlm = init_document_embeddings()
+    embeddings = DocumentRNNEmbeddings([glove])
+    sentences_1 = [Sentence("First sentence"), Sentence("This is second sentence")]
+    sentences_2 = [Sentence("This is second sentence"), Sentence("First sentence")]
+
+    embeddings.embed(sentences_1)
+    embeddings.embed(sentences_2)
+
+    assert sentences_1[0].to_original_text() == "First sentence"
+    assert sentences_1[1].to_original_text() == "This is second sentence"
+
+    assert torch.norm(sentences_1[0].embedding - sentences_2[1].embedding) == 0.0
+    assert torch.norm(sentences_1[0].embedding - sentences_2[1].embedding) == 0.0
+
+
 @pytest.mark.integration
 def test_stacked_embeddings():
     sentence, glove, charlm = init_document_embeddings()