Support subword and word language models

Closes #199 (closed)

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

dan/datasets/extract/extract.py

@@ -33,6 +33,7 @@ from dan.datasets.extract.utils import (
     Tokenizer,
     download_image,
     get_bbox,
+    get_vocabulary,
     insert_token,
     normalize_linebreaks,
     normalize_spaces,
@@ -354,6 +355,8 @@ 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")
+        # Add unknown token to charset
+        self.charset.add(self.unknown_token)
 
         # Build LM tokens
         for token in sorted(list(self.charset)):
@@ -363,62 +366,40 @@ class ArkindexExtractor:
             self.language_tokens.append(
                 self.mapping.encode[token]
             ) if token in self.mapping.encode else self.language_tokens.append(token)
-
         self.language_tokens.append(self.mapping.ctc.encoded)
-        assert all(
-            [len(token) == 1 for token in self.language_lexicon]
-        ), "Tokens should be single characters."
 
         # Build LM corpus
-        train_corpus = [text.replace("\n", " ") for text in self.data["train"].values()]
+        train_corpus = [
+            text.replace(self.mapping.linebreak.display, self.mapping.space.display)
+            for text in self.data["train"].values()
+        ]
+
         tokenizer = Tokenizer(
-            train_corpus,
+            training_corpus=train_corpus,
+            charset=self.language_tokens,
+            unknown_token=self.unknown_token,
             outdir=self.output / "language_model",
             mapping=self.mapping,
             tokens=self.tokens,
             subword_vocab_size=self.subword_vocab_size,
         )
-        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 != ""
-            ]
-        )
+        for level, tokenize in (
+            ("characters", tokenizer.char_tokenize),
+            ("words", tokenizer.word_tokenize),
+            ("subwords", tokenizer.subword_tokenize),
+        ):
+            self.language_corpus[level] = list(map(tokenize, train_corpus))
 
         # 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)
-        ]
+        for level in ["words", "subwords"]:
+            self.language_lexicon[level] = [
+                f"{token} {tokenizer.char_tokenize(token)}"
+                for token in get_vocabulary(self.language_corpus[level])
+            ]
 
     def export(self):
         (self.output / "labels.json").write_text(