WIP: New Foundation Model

benchmark/recognition.py

@@ -103,6 +103,7 @@ def normalize_text(text: str) -> str:
     help="Comma-separated list of languages to benchmark.",
     default=None,
 )
+@click.option("--xla_eager", is_flag=True, help="Use XLA eager mode for Surya.")
 def main(
     results_dir: str,
     max_rows: int,
@@ -112,7 +113,13 @@ def main(
     tess_cpus: int,
     textract_cpus: int,
     languages: str | None,
+    xla_eager: bool = False,
 ):
+    if xla_eager:
+        import torch_xla
+
+        torch_xla.experimental.eager_mode(True)
+
     foundation_predictor = FoundationPredictor()
     rec_predictor = RecognitionPredictor(foundation_predictor)
 

surya/common/donut/encoder.py

@@ -17,7 +17,7 @@
 from transformers.utils import ModelOutput
 from transformers import DonutSwinConfig
 
-from surya.common.util import mark_step
+from surya.common.xla import mark_step
 
 _EXPECTED_OUTPUT_SHAPE = [1, 49, 1024]
 

surya/common/predictor.py

@@ -9,15 +9,16 @@
 class BasePredictor:
     model_loader_cls = ModelLoader
     batch_size: Optional[int] = None
-    default_batch_sizes = {
-        "cpu": 1,
-        "mps": 1,
-        "cuda": 1
-    }
+    default_batch_sizes = {"cpu": 1, "mps": 1, "cuda": 1}
     disable_tqdm: bool = settings.DISABLE_TQDM
     torch_dtype = settings.MODEL_DTYPE
 
-    def __init__(self, checkpoint: Optional[str] = None, device: torch.device | str | None = settings.TORCH_DEVICE_MODEL, dtype: Optional[torch.dtype | str] = None):
+    def __init__(
+        self,
+        checkpoint: Optional[str] = None,
+        device: torch.device | str | None = settings.TORCH_DEVICE_MODEL,
+        dtype: Optional[torch.dtype | str] = None,
+    ):
         if dtype is None:
             dtype = self.torch_dtype
 
@@ -43,15 +44,17 @@ def get_batch_size(self):
         return batch_size
 
     @staticmethod
-    def pad_to_batch_size(tensor: torch.Tensor, batch_size: int):
+    def pad_to_batch_size(
+        tensor: torch.Tensor, batch_size: int, pad_value: int = 0
+    ) -> torch.Tensor:
         current_batch_size = tensor.shape[0]
         if current_batch_size >= batch_size:
             return tensor
 
         pad_size = batch_size - current_batch_size
         padding = (0, 0) * (tensor.dim() - 1) + (0, pad_size)
 
-        return F.pad(tensor, padding, mode='constant', value=0)
+        return F.pad(tensor, padding, mode="constant", value=pad_value)
 
     def __call__(self, *args, **kwargs):
-        raise NotImplementedError()
+        raise NotImplementedError()