Repository: flairNLP/flair

Issue:
Train FlairEmbedding with document boundaries separator token
Hello
I have tried to train new Flair Embeddings and as mentioned In https://github.com/flairNLP/flair/blob/master/resources/docs/TUTORIAL_9_TRAINING_LM_EMBEDDINGS.md , it is possible to use a document separator token like [SEP]. I use this token in my corpus to separate document boundaries, how should we introduce this token to LanguageModelTrainer ?!
Is there any docs to share more details to how sentences  feed to model in this way when shuffle_lines variabble is true?!


Discussion:
Hello @otahmasebi good points, thanks for noticing! Setting `shuffle_lines` to `True` simply shuffles a file by newlines and does not respect document boundaries. I think we need to change this.

Resolving patch:
diff --git a/flair/embeddings.py b/flair/embeddings.py
--- a/flair/embeddings.py
+++ b/flair/embeddings.py
@@ -2046,7 +2046,7 @@ def _add_embeddings_internal(self, sentences: List[Sentence]) -> List[Sentence]:
             # if this is not possible, use LM to generate embedding. First, get text sentences
             text_sentences = [sentence.to_tokenized_string() for sentence in sentences]
 
-            start_marker = "\n"
+            start_marker = self.lm.document_delimiter if "document_delimiter" in self.lm.__dict__ else '\n'
             end_marker = " "
 
             # get hidden states from language model
diff --git a/flair/models/language_model.py b/flair/models/language_model.py
--- a/flair/models/language_model.py
+++ b/flair/models/language_model.py
@@ -23,12 +23,14 @@ def __init__(
         nlayers: int,
         embedding_size: int = 100,
         nout=None,
+        document_delimiter: str = '\n',
         dropout=0.1,
     ):
 
         super(LanguageModel, self).__init__()
 
         self.dictionary = dictionary
+        self.document_delimiter = document_delimiter
         self.is_forward_lm: bool = is_forward_lm
 
         self.dropout = dropout
@@ -187,14 +189,17 @@ def load_language_model(cls, model_file: Union[Path, str]):
 
         state = torch.load(str(model_file), map_location=flair.device)
 
+        document_delimiter = state["document_delimiter"] if "document_delimiter" in state else '\n'
+
         model = LanguageModel(
-            state["dictionary"],
-            state["is_forward_lm"],
-            state["hidden_size"],
-            state["nlayers"],
-            state["embedding_size"],
-            state["nout"],
-            state["dropout"],
+            dictionary=state["dictionary"],
+            is_forward_lm=state["is_forward_lm"],
+            hidden_size=state["hidden_size"],
+            nlayers=state["nlayers"],
+            embedding_size=state["embedding_size"],
+            nout=state["nout"],
+            document_delimiter=document_delimiter,
+            dropout=state["dropout"],
         )
         model.load_state_dict(state["state_dict"])
         model.eval()
@@ -209,18 +214,21 @@ def load_checkpoint(cls, model_file: Path):
         epoch = state["epoch"] if "epoch" in state else None
         split = state["split"] if "split" in state else None
         loss = state["loss"] if "loss" in state else None
+        document_delimiter = state["document_delimiter"] if "document_delimiter" in state else '\n'
+
         optimizer_state_dict = (
             state["optimizer_state_dict"] if "optimizer_state_dict" in state else None
         )
 
         model = LanguageModel(
-            state["dictionary"],
-            state["is_forward_lm"],
-            state["hidden_size"],
-            state["nlayers"],
-            state["embedding_size"],
-            state["nout"],
-            state["dropout"],
+            dictionary=state["dictionary"],
+            is_forward_lm=state["is_forward_lm"],
+            hidden_size=state["hidden_size"],
+            nlayers=state["nlayers"],
+            embedding_size=state["embedding_size"],
+            nout=state["nout"],
+            document_delimiter=document_delimiter,
+            dropout=state["dropout"],
         )
         model.load_state_dict(state["state_dict"])
         model.eval()
@@ -245,6 +253,7 @@ def save_checkpoint(
             "nlayers": self.nlayers,
             "embedding_size": self.embedding_size,
             "nout": self.nout,
+            "document_delimiter": self.document_delimiter,
             "dropout": self.dropout,
             "optimizer_state_dict": optimizer.state_dict(),
             "epoch": epoch,
@@ -263,6 +272,7 @@ def save(self, file: Path):
             "nlayers": self.nlayers,
             "embedding_size": self.embedding_size,
             "nout": self.nout,
+            "document_delimiter": self.document_delimiter,
             "dropout": self.dropout,
         }
 
diff --git a/flair/trainers/language_model_trainer.py b/flair/trainers/language_model_trainer.py
--- a/flair/trainers/language_model_trainer.py
+++ b/flair/trainers/language_model_trainer.py
@@ -31,7 +31,8 @@ def __init__(
         forward: bool = True,
         split_on_char: bool = True,
         random_case_flip: bool = True,
-        shuffle_lines: bool = True,
+        document_delimiter: str = '\n',
+        shuffle: bool = True,
     ):
 
         assert path.exists()
@@ -43,7 +44,8 @@ def __init__(
         self.forward = forward
         self.random_case_flip = random_case_flip
         self.expand_vocab = expand_vocab
-        self.shuffle_lines = shuffle_lines
+        self.document_delimiter = document_delimiter
+        self.shuffle = shuffle
 
         if path.is_dir():
             self.files = sorted([f for f in path.iterdir() if f.exists()])
@@ -74,9 +76,11 @@ def charsplit(
         """Tokenizes a text file on character basis."""
         assert path.exists()
 
-        lines = open(path, "r", encoding="utf-8").readlines()
+        lines = [doc + self.document_delimiter
+                 for doc in open(path, "r", encoding="utf-8").read().split(self.document_delimiter) if doc]
+
         log.info(f"read text file with {len(lines)} lines")
-        if self.shuffle_lines:
+        if self.shuffle:
             random.shuffle(lines)
             log.info(f"shuffled")
 
@@ -130,6 +134,7 @@ def charsplit(
                         break
                     ids[token] = self.dictionary.get_idx_for_item(char)
                     token -= 1
+
         return ids
 
     @staticmethod
@@ -174,13 +179,13 @@ def __init__(
         forward: bool = True,
         character_level: bool = True,
         random_case_flip: bool = True,
-        shuffle_lines: bool = True,
+        document_delimiter: str = '\n',
     ):
         self.dictionary: Dictionary = dictionary
         self.forward = forward
         self.split_on_char = character_level
         self.random_case_flip = random_case_flip
-        self.shuffle_lines = shuffle_lines
+        self.document_delimiter: str = document_delimiter
 
         if type(path) == str:
             path = Path(path)
@@ -192,7 +197,8 @@ def __init__(
             self.forward,
             self.split_on_char,
             self.random_case_flip,
-            shuffle_lines=self.shuffle_lines,
+            document_delimiter=self.document_delimiter,
+            shuffle=True,
         )
 
         # TextDataset returns a list. valid and test are only one file, so return the first element
@@ -203,7 +209,8 @@ def __init__(
             self.forward,
             self.split_on_char,
             self.random_case_flip,
-            shuffle_lines=False,
+            document_delimiter=document_delimiter,
+            shuffle=False,
         )[0]
         self.test = TextDataset(
             path / "test.txt",
@@ -212,7 +219,8 @@ def __init__(
             self.forward,
             self.split_on_char,
             self.random_case_flip,
-            shuffle_lines=False,
+            document_delimiter=document_delimiter,
+            shuffle=False,
         )[0]
 
 


Test patch: