[WIP] onnx support

setup.py

@@ -31,5 +31,23 @@ def get_requirements(req_file):
     long_description=open("README.md", "r", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
     python_requires=">=3.6",
-    install_requires=get_requirements("requirements.txt"),
+    extras_require={
+        "onnx": ["onnxruntime"],
+        "optimum": ["optimum[onnxruntime]"],
+        "optimum-gpu": ["optimum[onnxruntime-gpu]"]
+    },
+    install_requires=[
+        "torch>=1.5",
+        "torchvision>=0.6.0",
+        "pandas",
+        "pytest",
+        "pysbd",
+        "layoutparser[effdet]>=0.2",
+        "transformers>4.5", # Enforce the version for now
+        "datasets",
+        "pdfplumber",
+        "pdf2image",
+        "tqdm",
+        "scikit-learn"
+    ]
 )

src/vila/__init__.py

@@ -7,4 +7,4 @@
     HierarchicalPDFPredictor,
 )
 
-__version__ = "0.4.2"
+__version__ = "0.5.0+cw07"

src/vila/automodel.py

@@ -1,3 +1,5 @@
+import types
+
 from .models import HierarchicalModelConfig, HierarchicalModelForTokenClassification
 
 from transformers import (
@@ -6,7 +8,8 @@
     MODEL_NAMES_MAPPING,
     TOKENIZER_MAPPING,
 )
-from transformers.models.auto.modeling_auto import auto_class_factory
+#from transformers.models.auto.modeling_auto import auto_class_factory
+from transformers.models.auto.modeling_auto import _BaseAutoModelClass, auto_class_update
 from transformers import BertTokenizer, BertTokenizerFast, AutoTokenizer
 
 CONFIG_MAPPING.update([("hierarchical_model", HierarchicalModelConfig)])
@@ -21,8 +24,12 @@
     [(HierarchicalModelConfig, HierarchicalModelForTokenClassification)]
 )
 
-AutoModelForTokenClassification = auto_class_factory(
-    "AutoModelForTokenClassification",
-    MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
-    head_doc="token classification",
-)
+cls = types.new_class("AutoModelForTokenClassification", (_BaseAutoModelClass,))
+cls._model_mapping = MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING
+cls.__name__ = "AutoModelForTokenClassification"
+AutoModelForTokenClassification = auto_class_update(cls, head_doc="token classification")
+#AutoModelForTokenClassification = auto_class_factory(
+#    "AutoModelForTokenClassification",
+#    MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
+#    head_doc="token classification",
+#)

src/vila/predictors.py

@@ -4,6 +4,7 @@
 import inspect
 import logging
 import copy
+import os
 
 import numpy as np
 import torch
@@ -75,7 +76,7 @@ def normalize_bbox(
         scale_factor = target_width / page_width if page_width > page_height else target_height / page_height
 
         logger.debug(f"Scaling page as page width {page_width} is larger than target width {target_width} or height {page_height} is larger than target height {target_height}")
-        
+
         x1 = float(x1) * scale_factor
         x2 = float(x2) * scale_factor
 
@@ -100,14 +101,14 @@ def unnormalize_bbox(
 
     # Right now only execute this for only "large" PDFs
     # TODO: Change it for all PDFs
-    
+
     if page_width > target_width or page_height > target_height:
 
         # Aspect ratio preserving scaling
         scale_factor = target_width / page_width if page_width > page_height else target_height / page_height
 
         logger.debug(f"Scaling page as page width {page_width} is larger than target width {target_width} or height {page_height} is larger than target height {target_height}")
-        
+
         x1 = float(x1) / scale_factor
         x2 = float(x2) / scale_factor
 
@@ -129,17 +130,30 @@ def __init__(self, model, preprocessor, device):
             self.device = device
             model.to(self.device)
 
-        self.model.eval()
+        # Optimum-wrapped ONNX models don't have an eval mode
+        if hasattr(self.model, "eval"):
+            self.model.eval()
         self._used_cols = columns_used_in_model_inputs(self.model)
 
     @classmethod
     def from_pretrained(
         cls, model_path, preprocessor=None, device=None, **preprocessor_config
     ):
-
-        model = AutoModelForTokenClassification.from_pretrained(model_path)
         tokenizer = AutoTokenizer.from_pretrained(model_path)
 
+        if os.path.exists(os.path.join(model_path, "model.onnx")):
+            try:
+                from optimum.onnxruntime import ORTModelForTokenClassification
+                model = ORTModelForTokenClassification.from_pretrained(model_path, file_name="model.onnx")
+            except:
+                raise Exception("""
+                    The provided model is an ONNX graph, and requires additional packages to be installed.
+                    Please install `vila[optimum]` / `vila[optimum-gpu]`, or switch to an uncompiled 
+                    pytorch model to proceed.
+                """)
+        else:
+            model = AutoModelForTokenClassification.from_pretrained(model_path)
+
         if preprocessor is None:
             preprocessor_config = VILAPreprocessorConfig.from_pretrained(
                 model_path, **preprocessor_config