Support subword and word language models

Closes #199 (closed)

Ref https://redmine.teklia.com/issues/4941

dan/datasets/extract/extract.py

@@ -30,6 +30,7 @@ from dan.datasets.extract.exceptions import (
     UnknownTokenInText,
 )
 from dan.datasets.extract.utils import (
+    Tokenizer,
     download_image,
     get_bbox,
     insert_token,
@@ -95,9 +96,9 @@ class ArkindexExtractor:
 
         self.data: Dict = defaultdict(dict)
         self.charset = set()
-        self.language_corpus = []
+        self.language_corpus = defaultdict(list)
         self.language_tokens = []
-        self.language_lexicon = []
+        self.language_lexicon = defaultdict(list)
 
         # Image download tasks to process
         self.tasks: List[Dict[str, str]] = []
@@ -271,12 +272,6 @@ class ArkindexExtractor:
             )
         return text.strip()
 
-    def format_text_language_model(self, text: str):
-        """
-        Format text for the language model. Return the text tokenized at character-level.
-        """
-        return " ".join(map(self.mapping.encode_token, list(text.strip())))
-
     def process_element(
         self,
         element: Element,
@@ -315,10 +310,6 @@ class ArkindexExtractor:
         self.data[split][str(image_path)] = text
         self.charset = self.charset.union(set(text))
 
-        # Language model should be built using only text from the training set
-        if split == "train":
-            self.language_corpus.append(self.format_text_language_model(text))
-
     def process_parent(
         self,
         pbar,
@@ -357,6 +348,9 @@ class ArkindexExtractor:
         """
         Convert charset to a LM-compatible charset. Ensure that special LM tokens do not appear in the charset.
         """
+        logger.info("Preparing language resources")
+
+        # Build LM tokens
         for token in sorted(list(self.charset)):
             assert (
                 token not in self.mapping.encode.values()
@@ -365,14 +359,60 @@ class ArkindexExtractor:
                 self.mapping.encode[token]
             ) if token in self.mapping.encode else self.language_tokens.append(token)
 
-        # Add the special blank token
         self.language_tokens.append(self.mapping.ctc.encoded)
-
-        # Build lexicon
         assert all(
             [len(token) == 1 for token in self.language_lexicon]
         ), "Tokens should be single characters."
-        self.language_lexicon = [f"{token} {token}" for token in self.language_tokens]
+
+        # Build LM corpus
+        train_corpus = [text for text in self.data["train"].values()]
+        tokenizer = Tokenizer(
+            train_corpus,
+            outdir=self.output / "language_model",
+            mapping=self.mapping,
+            tokens=self.tokens,
+        )
+        self.language_corpus["characters"] = [
+            tokenizer.char_tokenize(doc) for doc in train_corpus
+        ]
+        self.language_corpus["words"] = [
+            tokenizer.word_tokenize(doc) for doc in train_corpus
+        ]
+        self.language_corpus["subwords"] = [
+            tokenizer.subword_tokenize(doc) for doc in train_corpus
+        ]
+
+        # Build vocabulary
+        word_vocabulary = set(
+            [
+                word
+                for doc in self.language_corpus["words"]
+                for word in doc.split()
+                if word != ""
+            ]
+        )
+        subword_vocabulary = set(
+            [
+                subword
+                for doc in self.language_corpus["subwords"]
+                for subword in doc.split()
+                if subword != ""
+            ]
+        )
+
+        # Build LM lexicon
+        self.language_lexicon["characters"] = [
+            f"{token} {token}" for token in self.language_tokens
+        ]
+        self.language_lexicon["words"] = [
+            f"{word} {tokenizer.char_tokenize(word)}"
+            for word in sorted(word_vocabulary)
+            if word != ""
+        ]
+        self.language_lexicon["subwords"] = [
+            f"{subword} {tokenizer.char_tokenize(subword)}"
+            for subword in sorted(subword_vocabulary)
+        ]
 
     def export(self):
         (self.output / "labels.json").write_text(
@@ -382,15 +422,16 @@ class ArkindexExtractor:
                 indent=4,
             )
         )
-        (self.output / "language_model" / "corpus.txt").write_text(
-            "\n".join(self.language_corpus)
-        )
+        for level in ["characters", "words", "subwords"]:
+            (self.output / "language_model" / f"corpus_{level}.txt").write_text(
+                "\n".join(self.language_corpus[level])
+            )
+            (self.output / "language_model" / f"lexicon_{level}.txt").write_text(
+                "\n".join(self.language_lexicon[level])
+            )
         (self.output / "language_model" / "tokens.txt").write_text(
             "\n".join(self.language_tokens)
         )
-        (self.output / "language_model" / "lexicon.txt").write_text(
-            "\n".join(self.language_lexicon)
-        )
         (self.output / "charset.pkl").write_bytes(
             pickle.dumps(sorted(list(self.charset)))
         )