Merge pull request #3641 from flairNLP/GH-3635-lazy-tokenization

Lazy Tokenization in Flair

Author: alanakbik

flair/data.py

@@ -818,22 +818,21 @@ def __init__(
         """
         super().__init__()
 
-        self.tokens: list[Token] = []
+        self._tokens: Optional[list[Token]] = None
+        self._text: str = ""  # Change from Optional[str] to str with empty string default
 
-        # private field for all known spans
-        self._known_spans: dict[str, _PartOfSentence] = {}
+        # private field for all known spans with explicit typing
+        self._known_spans: dict[str, Union[Span, Relation]] = {}
 
         self.language_code: Optional[str] = language_code
 
         self._start_position = start_position
 
         # the tokenizer used for this sentence
         if isinstance(use_tokenizer, Tokenizer):
-            tokenizer = use_tokenizer
-
+            self._tokenizer = use_tokenizer
         elif isinstance(use_tokenizer, bool):
-            tokenizer = SegtokTokenizer() if use_tokenizer else SpaceTokenizer()
-
+            self._tokenizer = SegtokTokenizer() if use_tokenizer else SpaceTokenizer()
         else:
             raise AssertionError("Unexpected type of parameter 'use_tokenizer'. Parameter should be bool or Tokenizer")
 
@@ -848,24 +847,79 @@ def __init__(
         self._next_sentence: Optional[Sentence] = None
         self._position_in_dataset: Optional[tuple[Dataset, int]] = None
 
-        # if text is passed, instantiate sentence with tokens (words)
-        if isinstance(text, str):
-            text = Sentence._handle_problem_characters(text)
-            words = tokenizer.tokenize(text)
-        elif text and isinstance(text[0], Token):
-            for t in text:
-                self._add_token(t)
-            self.tokens[-1].whitespace_after = 0
-            return
+        # if list of strings or tokens is passed, create tokens directly
+        if not isinstance(text, str):
+            self._tokens = []
+
+            # First construct the text from tokens to ensure proper text reconstruction
+            if len(text) > 0:
+                # Type check the input list and cast
+                if all(isinstance(t, Token) for t in text):
+                    tokens = cast(list[Token], text)
+                    reconstructed_text = ""
+                    for i, token in enumerate(tokens):
+                        reconstructed_text += token.text
+                        if i < len(tokens) - 1:  # Add whitespace between tokens
+                            reconstructed_text += " " * token.whitespace_after
+                    self._text = reconstructed_text
+                elif all(isinstance(t, str) for t in text):
+                    strings = cast(list[str], text)
+                    self._text = " ".join(strings)
+                else:
+                    raise TypeError("All elements must be either Token or str")
+            else:
+                self._text = ""
+
+            # Now add the tokens
+            current_position = 0
+            for i, item in enumerate(text):
+                # create Token if string, otherwise use existing Token
+                if isinstance(item, str):
+                    # For strings, create new Token with default whitespace
+                    token = Token(text=item)
+                    token.whitespace_after = 0 if i == len(text) - 1 else 1
+                elif isinstance(item, Token):
+                    # For existing Tokens, preserve their whitespace_after
+                    token = item
+
+                # Set start position for the token
+                token.start_position = current_position
+                current_position += len(token.text) + token.whitespace_after
+
+                self._add_token(token)
+
+            if len(text) > 0:
+                # convention: the last token has no whitespace after
+                self.tokens[-1].whitespace_after = 0
         else:
-            words = cast(list[str], text)
-            text = " ".join(words)
+            self._text = Sentence._handle_problem_characters(text)
+
+        # log a warning if the dataset is empty
+        if self._text == "":
+            log.warning("Warning: An empty Sentence was created! Are there empty strings in your dataset?")
+
+    @property
+    def tokens(self) -> list[Token]:
+        """Gets the tokens of this sentence. Automatically triggers tokenization if not yet tokenized."""
+        if self._tokens is None:
+            self._tokenize()
+        if self._tokens is None:
+            raise ValueError("Tokens are None after tokenization - this indicates a bug in the tokenization process")
+        return self._tokens
+
+    def _tokenize(self) -> None:
+        """Internal method that performs tokenization."""
+
+        # tokenize the text
+        words = self._tokenizer.tokenize(self._text)
 
         # determine token positions and whitespace_after flag
         current_offset: int = 0
         previous_token: Optional[Token] = None
+        self._tokens = []
+
         for word in words:
-            word_start_position: int = text.index(word, current_offset)
+            word_start_position: int = self._text.index(word, current_offset)
             delta_offset: int = word_start_position - current_offset
 
             token: Token = Token(text=word, start_position=word_start_position)
@@ -878,17 +932,56 @@ def __init__(
             previous_token = token
 
         # the last token has no whitespace after
-        if len(self) > 0:
-            self.tokens[-1].whitespace_after = 0
+        if len(self._tokens) > 0:
+            self._tokens[-1].whitespace_after = 0
 
-        # log a warning if the dataset is empty
-        if text == "":
-            log.warning("Warning: An empty Sentence was created! Are there empty strings in your dataset?")
+    def __iter__(self):
+        """Allows iteration over tokens. Triggers tokenization if not yet tokenized."""
+        return iter(self.tokens)
+
+    def __len__(self) -> int:
+        """Returns the number of tokens in this sentence. Triggers tokenization if not yet tokenized."""
+        return len(self.tokens)
 
     @property
     def unlabeled_identifier(self):
         return f'Sentence[{len(self)}]: "{self.text}"'
 
+    @property
+    def text(self) -> str:
+        """Returns the original text of this sentence. Does not trigger tokenization."""
+        return self._text
+
+    def to_original_text(self) -> str:
+        """Returns the original text of this sentence."""
+        return self._text
+
+    def to_tagged_string(self, main_label: Optional[str] = None) -> str:
+        # For sentence-level labels, we don't need tokenization
+        if not self._tokens:
+            output = f'Sentence: "{self.text}"'
+            if self.labels:
+                output += self._printout_labels(main_label)
+            return output
+
+        # Only tokenize if we have token-level labels or spans to print
+        already_printed = [self]
+        output = super().__str__()
+
+        label_append = []
+        for label in self.get_labels(main_label):
+            if label.data_point in already_printed:
+                continue
+            label_append.append(
+                f'"{label.data_point.text}"{label.data_point._printout_labels(main_label=main_label, add_score=False)}'
+            )
+            already_printed.append(label.data_point)
+
+        if len(label_append) > 0:
+            output += f"{flair._arrow}[" + ", ".join(label_append) + "]"
+
+        return output
+
     def get_relations(self, label_type: Optional[str] = None) -> list[Relation]:
         relations: list[Relation] = []
         for label in self.get_labels(label_type):
@@ -951,11 +1044,13 @@ def to(self, device: str, pin_memory: bool = False):
             token.to(device, pin_memory)
 
     def clear_embeddings(self, embedding_names: Optional[list[str]] = None):
+        # clear sentence embeddings
         super().clear_embeddings(embedding_names)
 
-        # clear token embeddings
-        for token in self:
-            token.clear_embeddings(embedding_names)
+        # clear token embeddings if sentence is tokenized
+        if self._is_tokenized():
+            for token in self.tokens:
+                token.clear_embeddings(embedding_names)
 
     def left_context(self, context_length: int, respect_document_boundaries: bool = True) -> list[Token]:
         sentence = self
@@ -987,29 +1082,6 @@ def right_context(self, context_length: int, respect_document_boundaries: bool =
     def __str__(self) -> str:
         return self.to_tagged_string()
 
-    def to_tagged_string(self, main_label: Optional[str] = None) -> str:
-        already_printed = [self]
-
-        output = super().__str__()
-
-        label_append = []
-        for label in self.get_labels(main_label):
-            if label.data_point in already_printed:
-                continue
-            label_append.append(
-                f'"{label.data_point.text}"{label.data_point._printout_labels(main_label=main_label, add_score=False)}'
-            )
-            already_printed.append(label.data_point)
-
-        if len(label_append) > 0:
-            output += f"{flair._arrow}[" + ", ".join(label_append) + "]"
-
-        return output
-
-    @property
-    def text(self) -> str:
-        return self.to_original_text()
-
     def to_tokenized_string(self) -> str:
         if self.tokenized is None:
             self.tokenized = " ".join([t.text for t in self.tokens])
@@ -1056,15 +1128,6 @@ def infer_space_after(self):
             last_token = token
         return self
 
-    def to_original_text(self) -> str:
-        # if sentence has no tokens, return empty string
-        if len(self) == 0:
-            return ""
-        # otherwise, return concatenation of tokens with the correct offsets
-        return (self[0].start_position - self.start_position) * " " + "".join(
-            [t.text + t.whitespace_after * " " for t in self.tokens]
-        ).strip()
-
     def to_dict(self, tag_type: Optional[str] = None) -> dict[str, Any]:
         return {
             "text": self.to_original_text(),
@@ -1090,12 +1153,6 @@ def __getitem__(self, subscript):
         else:
             return self.tokens[subscript]
 
-    def __iter__(self):
-        return iter(self.tokens)
-
-    def __len__(self) -> int:
-        return len(self.tokens)
-
     def __repr__(self) -> str:
         return self.__str__()
 
@@ -1233,20 +1290,59 @@ def get_labels(self, label_type: Optional[str] = None):
         return []
 
     def remove_labels(self, typename: str):
-        # labels also need to be deleted at all tokens
-        for token in self:
-            token.remove_labels(typename)
-
-        # labels also need to be deleted at all known spans
-        for span in self._known_spans.values():
-            span.remove_labels(typename)
+        # only access tokens if already tokenized
+        if self._is_tokenized():
+            # labels also need to be deleted at all tokens
+            for token in self.tokens:
+                token.remove_labels(typename)
 
-        # remove spans without labels
-        self._known_spans = {k: v for k, v in self._known_spans.items() if len(v.labels) > 0}
+            # labels also need to be deleted at all known spans
+            for span in self._known_spans.values():
+                span.remove_labels(typename)
 
-        # delete labels at object itself
+        # delete labels at object itself first
         super().remove_labels(typename)
 
+    def _is_tokenized(self) -> bool:
+        return self._tokens is not None
+
+    def truncate(self, max_tokens: int) -> None:
+        """Truncates the sentence to a maximum number of tokens and updates all annotations accordingly."""
+        if len(self.tokens) <= max_tokens:
+            return
+
+        # Truncate tokens
+        self._tokens = self.tokens[:max_tokens]
+
+        # Remove spans that reference removed tokens
+        self._known_spans = {
+            identifier: span
+            for identifier, span in self._known_spans.items()
+            if isinstance(span, Span) and all(token.idx <= max_tokens for token in span.tokens)
+        }
+
+        # Remove relations that reference removed spans
+        self._known_spans = {
+            identifier: relation
+            for identifier, relation in self._known_spans.items()
+            if not isinstance(relation, Relation)
+            or (
+                all(token.idx <= max_tokens for token in relation.first.tokens)
+                and all(token.idx <= max_tokens for token in relation.second.tokens)
+            )
+        }
+
+        # Clean up any labels that reference removed spans/relations
+        for typename in list(self.annotation_layers.keys()):
+            self.annotation_layers[typename] = [
+                label
+                for label in self.annotation_layers[typename]
+                if (
+                    not isinstance(label.data_point, (Span, Relation))
+                    or label.data_point.unlabeled_identifier in self._known_spans
+                )
+            ]
+
 
 class DataPair(DataPoint, typing.Generic[DT, DT2]):
     def __init__(self, first: DT, second: DT2) -> None:
@@ -1375,7 +1471,7 @@ class Corpus(typing.Generic[T_co]):
     """The main object in Flair for holding a dataset used for training and testing.
 
     A corpus consists of three splits: A `train` split used for training, a `dev` split used for model selection
-    and/or early stopping and a `test` split used for testing. All three splits are optional, so it is possible
+    or early stopping and a `test` split used for testing. All three splits are optional, so it is possible
     to create a corpus only using one or two splits. If the option `sample_missing_splits` is set to True,
     missing splits will be randomly sampled from the training split.
     """

flair/datasets/base.py

@@ -207,7 +207,7 @@ def _parse_document_to_sentence(
                 sentence.add_label(self.tag_type, label)
 
             if self.max_tokens_per_doc > 0:
-                sentence.tokens = sentence.tokens[: min(len(sentence), self.max_tokens_per_doc)]
+                sentence.truncate(self.max_tokens_per_doc)
 
             return sentence
         return None

flair/datasets/document_classification.py

@@ -268,12 +268,15 @@ def _parse_line_to_sentence(self, line: str, label_prefix: str, tokenizer: Union
         if text and (labels or self.allow_examples_without_labels):
             sentence = Sentence(text, use_tokenizer=tokenizer)
 
+            if 0 < self.truncate_to_max_tokens < len(sentence):
+                # Create new sentence with truncated text
+                truncated_text = " ".join(token.text for token in sentence.tokens[: self.truncate_to_max_tokens])
+                sentence = Sentence(truncated_text, use_tokenizer=tokenizer)
+
+            # Add the labels
             for label in labels:
                 sentence.add_label(self.label_type, label)
 
-            if sentence is not None and 0 < self.truncate_to_max_tokens < len(sentence):
-                sentence.tokens = sentence.tokens[: self.truncate_to_max_tokens]
-
             return sentence
         return None