Swap over to t5 editor

benchmark.py

@@ -119,7 +119,7 @@ if __name__ == "__main__":
     score_headers = benchmark_files
     for method in methods:
         summary_table.append([method, write_data[method]["avg_score"], write_data[method]["time_per_page"], write_data[method]["time_per_doc"]])
-        score_table.append([method, *[write_data[method]["scores"][h] for h in score_headers]])
+        score_table.append([method, *[write_data[method]["files"][h]["score"] for h in score_headers]])
 
     print(tabulate(summary_table, headers=["Method", "Average Score", "Time per page", "Time per document"]))
     print("")

marker/cleaners/equations.py

@@ -108,7 +108,7 @@ def get_nougat_text_batched(images, reformat_region_lens, nougat_model, batch_si
         for j, output in enumerate(model_output["predictions"]):
             disclaimer = ""
             token_count = get_total_nougat_tokens(output, nougat_model)
-            if token_count >= max_length:
+            if token_count >= max_length - 1:
                 disclaimer = "[TRUNCATED]"
 
             image_idx = idx * batch_size + j

marker/postprocessors/editor.py

@@ -3,11 +3,11 @@ from itertools import chain
 from typing import Optional
 import re
 
-from transformers import BloomForTokenClassification, AutoTokenizer
+from transformers import AutoTokenizer
 from marker.settings import settings
 import torch
 import torch.nn.functional as F
-from marker.postprocessors.t5 import T5ForTokenClassification
+from marker.postprocessors.t5 import T5ForTokenClassification, byt5_tokenize
 
 tokenizer = AutoTokenizer.from_pretrained(settings.EDITOR_MODEL_NAME)
 
@@ -37,24 +37,18 @@ def edit_full_text(text: str, model: Optional[T5ForTokenClassification], batch_s
     if not model:
         return text, {}
 
-    tokenized = tokenizer(
-        text,
-        truncation=True,
-        max_length=settings.EDITOR_MAX_LENGTH,
-        return_overflowing_tokens=True,
-        padding="max_length",
-    )
+    tokenized = byt5_tokenize(text, settings.EDITOR_MAX_LENGTH)
     input_ids = tokenized["input_ids"]
+    char_token_lengths = tokenized["char_token_lengths"]
 
     # Tokenize, and make sure reverse tokenization works
     model_tokens = [tokenizer.convert_ids_to_tokens(t, skip_special_tokens=True) for t in input_ids]
-    full_text = "".join(model_tokens)
+    model_tokens_str = [tokenizer.convert_tokens_to_string(t) for t in model_tokens]
+    full_text = "".join(model_tokens_str)
     assert full_text == text
 
     # List of characters in the text
-    model_tokens = [tokenizer.convert_ids_to_tokens(t) for t in input_ids]
-    flat_model_tokens = list(chain.from_iterable(model_tokens))
-    flat_str_tokens = list(text)
+    flat_input_ids = list(chain.from_iterable(input_ids))
 
     # Run model
     token_masks = []
@@ -72,47 +66,50 @@ def edit_full_text(text: str, model: Optional[T5ForTokenClassification], batch_s
         # We want to be conservative to not edit the text too much
         probs = F.softmax(logits, dim=-1)
         max_prob = torch.max(probs, dim=-1)
-        cutoff_prob = max_prob.values < 0.9
+        cutoff_prob = max_prob.values < settings.EDITOR_CUTOFF_THRESH
         labels = logits.argmax(-1).squeeze()
         labels[cutoff_prob] = model.config.label2id["equal"]
-
         labels = labels.tolist()
         if len(labels) == settings.EDITOR_MAX_LENGTH:
             labels = [labels]
         labels = list(chain.from_iterable(labels))
         token_masks.extend(labels)
 
-    # Strip special tokens
-    assert len(token_masks) == len(flat_model_tokens)
-    token_masks = [mask for mask, token in zip(token_masks, flat_model_tokens) if token not in ["<pad>", "<s>", "</s>"]]
+    # Strip special tokens 0,1.  Keep unknown token, although it should never be used
+    assert len(token_masks) == len(flat_input_ids)
+    token_masks = [mask for mask, token in zip(token_masks, flat_input_ids) if token >= 2]
 
-    assert len(token_masks) == len(flat_str_tokens)
+    assert len(token_masks) == len(list(text.encode("utf-8")))
 
     edit_stats = defaultdict(int)
-    out_tokens = []
-    for i, (str_token, mask) in enumerate(zip(flat_str_tokens, token_masks)):
-        label = model.config.id2label[mask]
-
-        match label:
-            case "equal":
-                out_tokens.append(str_token)
-                edit_stats[label] += 1
-            case "delete":
-                # If we delete whitespace, roll with it, otherwise ignore
-                if str_token.strip():
-                    out_tokens.append(str_token)
-                else:
-                    edit_stats[label] += 1
-            case "newline-1":
-                out_tokens.append("\n")
-                out_tokens.append(str_token)
-                edit_stats[label] += 1
-            case "space-1":
-                out_tokens.append(" ")
-                out_tokens.append(str_token)
-                edit_stats[label] += 1
-
-    return "".join(out_tokens), edit_stats
+    out_text = []
+    start = 0
+    for i, char in enumerate(text):
+        char_token_length = char_token_lengths[i]
+        masks = token_masks[start: start + char_token_length]
+        labels = [model.config.id2label[mask] for mask in masks]
+        if all(l == "delete" for l in labels):
+            # If we delete whitespace, roll with it, otherwise ignore
+            if char.strip():
+                out_text.append(char)
+            else:
+                edit_stats["delete"] += 1
+        elif labels[0] == "newline-1":
+            out_text.append("\n")
+            out_text.append(char)
+            edit_stats["newline-1"] += 1
+        elif labels[0] == "space-1":
+            out_text.append(" ")
+            out_text.append(char)
+            edit_stats["space-1"] += 1
+        else:
+            out_text.append(char)
+            edit_stats["equal"] += 1
+
+        start += char_token_length
+
+    out_text = "".join(out_text)
+    return out_text, edit_stats
 
 
 

marker/postprocessors/t5.py

@@ -1,91 +1,141 @@
-from transformers import T5ForSequenceClassification, T5Config
+from transformers import T5Config, T5PreTrainedModel
 import torch
+from torch import nn
+from copy import deepcopy
 from typing import Optional, Tuple, Union, List
+from itertools import chain
 
-from transformers.modeling_outputs import Seq2SeqSequenceClassifierOutput
+from transformers.modeling_outputs import TokenClassifierOutput
+from transformers.models.t5.modeling_t5 import T5Stack
+from transformers.utils.model_parallel_utils import get_device_map, assert_device_map
 
 
-class T5ForTokenClassification(T5ForSequenceClassification):
+def byt5_tokenize(text: str, max_length: int, pad_token_id: int = 0):
+    byte_codes = []
+    for char in text:
+        # Add 3 to account for special tokens
+        byte_codes.append([byte + 3 for byte in char.encode('utf-8')])
+
+    tokens = list(chain.from_iterable(byte_codes))
+    # Map each token to the character it represents
+    char_token_lengths = [len(b) for b in byte_codes]
+
+    batched_tokens = []
+    attention_mask = []
+    for i in range(0, len(tokens), max_length):
+        batched_tokens.append(tokens[i:i + max_length])
+        attention_mask.append([1] * len(batched_tokens[-1]))
+
+    # Pad last item
+    if len(batched_tokens[-1]) < max_length:
+        batched_tokens[-1] += [pad_token_id] * (max_length - len(batched_tokens[-1]))
+        attention_mask[-1] += [0] * (max_length - len(attention_mask[-1]))
+
+    return {"input_ids": batched_tokens, "attention_mask": attention_mask, "char_token_lengths": char_token_lengths}
+
+
+
+
+# From https://github.com/osainz59/t5-encoder
+class T5ForTokenClassification(T5PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"encoder.embed_tokens.weight"]
+
+    def __init__(self, config: T5Config):
+        super().__init__(config)
+        self.model_dim = config.d_model
+
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.is_encoder_decoder = False
+        encoder_config.use_cache = False
+        self.encoder = T5Stack(encoder_config, self.shared)
+
+        classifier_dropout = (
+            config.classifier_dropout if hasattr(config, 'classifier_dropout') else config.dropout_rate
+        )
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.d_model, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+
+    def parallelize(self, device_map=None):
+        self.device_map = (
+            get_device_map(len(self.encoder.block), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.encoder.block))
+        self.encoder.parallelize(self.device_map)
+        self.classifier.to(self.encoder.first_device)
+        self.model_parallel = True
+
+    def deparallelize(self):
+        self.encoder.deparallelize()
+        self.encoder = self.encoder.to("cpu")
+        self.classifier = self.classifier.to("cpu")
+        self.model_parallel = False
+        self.device_map = None
+        torch.cuda.empty_cache()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+
+    def get_encoder(self):
+        return self.encoder
+
+    def _prune_heads(self, heads_to_prune):
+        for layer, heads in heads_to_prune.items():
+            self.encoder.block[layer].layer[0].SelfAttention.prune_heads(heads)
 
     def forward(
         self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        decoder_input_ids: Optional[torch.LongTensor] = None,
-        decoder_attention_mask: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        decoder_head_mask: Optional[torch.Tensor] = None,
-        cross_attn_head_mask: Optional[torch.Tensor] = None,
-        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
-        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, Seq2SeqSequenceClassifierOutput]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
-            config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-        Returns:
-        """
+    ) -> Union[Tuple[torch.FloatTensor], TokenClassifierOutput]:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        if labels is not None:
-            use_cache = False
-
-        if input_ids is None and inputs_embeds is not None:
-            raise NotImplementedError(
-                f"Passing input embeddings is currently not supported for {self.__class__.__name__}"
-            )
-
-        # Copied from models.bart.modeling_bart.BartModel.forward different to other models, T5 automatically creates
-        # decoder_input_ids from input_ids if no decoder_input_ids are provided
-        if decoder_input_ids is None and decoder_inputs_embeds is None:
-            if input_ids is None:
-                raise ValueError(
-                    "If no `decoder_input_ids` or `decoder_inputs_embeds` are "
-                    "passed, `input_ids` cannot be `None`. Please pass either "
-                    "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`."
-                )
-            decoder_input_ids = self._shift_right(input_ids)
-
-        outputs = self.transformer(
-            input_ids,
+
+        outputs = self.encoder(
+            input_ids=input_ids,
             attention_mask=attention_mask,
-            decoder_input_ids=decoder_input_ids,
-            decoder_attention_mask=decoder_attention_mask,
-            head_mask=head_mask,
-            decoder_head_mask=decoder_head_mask,
-            cross_attn_head_mask=cross_attn_head_mask,
-            encoder_outputs=encoder_outputs,
             inputs_embeds=inputs_embeds,
-            decoder_inputs_embeds=decoder_inputs_embeds,
-            use_cache=use_cache,
+            head_mask=head_mask,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-        # Make predictions for all tokens
+
         sequence_output = outputs[0]
-        logits = self.classification_head(sequence_output)
 
-        assert labels.numel() * self.config.num_labels == logits.numel()
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
         loss = None
 
         if not return_dict:
-            output = (logits,) + outputs[1:]
+            output = (logits,) + outputs[2:]
             return ((loss,) + output) if loss is not None else output
 
-        return Seq2SeqSequenceClassifierOutput(
+        return TokenClassifierOutput(
             loss=loss,
             logits=logits,
-            past_key_values=outputs.past_key_values,
-            decoder_hidden_states=outputs.decoder_hidden_states,
-            decoder_attentions=outputs.decoder_attentions,
-            cross_attentions=outputs.cross_attentions,
-            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
-            encoder_hidden_states=outputs.encoder_hidden_states,
-            encoder_attentions=outputs.encoder_attentions,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions
         )

marker/settings.py

@@ -54,7 +54,7 @@ class Settings(BaseSettings):
     NOUGAT_MODEL_MAX: int = 512 # Max inference length for nougat
     NOUGAT_TOKEN_BUFFER: int = 256 # Number of tokens to buffer above max for nougat
     NOUGAT_HALLUCINATION_WORDS: List[str] = ["[MISSING_PAGE_POST]", "## References\n", "**Figure Captions**\n", "Footnote",
-                                  "\par\par\par", "## Chapter", "Fig.", "particle", "[REPEATS]", "[TRUNCATED]", "### "]
+                                  "\par\par\par", "## Chapter", "Fig.", "particle", "[REPEATS]", "[TRUNCATED]", "### ", "effective field strength", "\Phi_{\rm eff}"]
     NOUGAT_DPI: int = 96 # DPI to render images at, matches default settings for nougat
     NOUGAT_MODEL_NAME: str = "0.1.0-small" # Name of the model to use
     NOUGAT_BATCH_SIZE: int = 6 if TORCH_DEVICE == "cuda" else 1 # Batch size for nougat, don't batch on cpu
@@ -74,8 +74,9 @@ class Settings(BaseSettings):
     # Final editing model
     EDITOR_BATCH_SIZE: int = 4
     EDITOR_MAX_LENGTH: int = 2048
-    EDITOR_MODEL_NAME: str = "vikp/pdf_postprocessor"
-    ENABLE_EDITOR_MODEL: bool = False # The editor model can create false positives
+    EDITOR_MODEL_NAME: str = "vikp/pdf_postprocessor_t5"
+    ENABLE_EDITOR_MODEL: bool = True # The editor model can create false positives
+    EDITOR_CUTOFF_THRESH: float = 0.75 # Ignore predictions below this probability
 
     # Ray
     RAY_CACHE_PATH: Optional[str] = None # Where to save ray cache