Merge pull request #212 from VikParuchuri/dev

Fix MPS issue with pytorch 2.5

ocr_text.py

@@ -4,14 +4,11 @@
 import time
 from collections import defaultdict
 
-import torch
-
-from surya.input.langs import replace_lang_with_code, get_unique_langs
+from surya.input.langs import replace_lang_with_code
 from surya.input.load import load_from_folder, load_from_file, load_lang_file
 from surya.model.detection.model import load_model as load_detection_model, load_processor as load_detection_processor
 from surya.model.recognition.model import load_model as load_recognition_model
 from surya.model.recognition.processor import load_processor as load_recognition_processor
-from surya.model.recognition.tokenizer import _tokenize
 from surya.ocr import run_ocr
 from surya.postprocessing.text import draw_text_on_image
 from surya.settings import settings

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surya-ocr"
-version = "0.6.4"
+version = "0.6.5"
 description = "OCR, layout, reading order, and table recognition in 90+ languages"
 authors = ["Vik Paruchuri <[email protected]>"]
 readme = "README.md"
@@ -23,7 +23,7 @@ include = [
 [tool.poetry.dependencies]
 python = "^3.10"
 transformers = "^4.41.0"
-torch = "^2.3.0"
+torch = "^2.4.1"
 pydantic = "^2.5.3"
 pydantic-settings = "^2.1.0"
 python-dotenv = "^1.0.0"
@@ -33,7 +33,7 @@ opencv-python = "^4.9.0.80"
 tabulate = "^0.9.0"
 filetype = "^1.2.0"
 ftfy = "^6.1.3"
-pdftext = "^0.3.16"
+pdftext = "^0.3.17"
 
 [tool.poetry.group.dev.dependencies]
 jupyter = "^1.0.0"

surya/model/recognition/decoder.py

@@ -173,9 +173,9 @@ def forward(
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
         attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states.contiguous(),
-            key_states.contiguous(),
-            value_states.contiguous(),
+            query_states,
+            key_states,
+            value_states,
             attn_mask=None,
             dropout_p=self.attention_dropout if self.training else 0.0,
             scale=self.head_dim**-0.5,
@@ -261,9 +261,9 @@ def forward(
                 causal_mask = torch.where(position_mask, torch.finfo(causal_mask.dtype).min, causal_mask)
 
         attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states.contiguous(),
-            key_states.contiguous(),
-            value_states.contiguous(),
+            query_states,
+            key_states,
+            value_states,
             attn_mask=causal_mask,
             dropout_p=self.attention_dropout if self.training else 0.0,
             scale=self.head_dim**-0.5,

surya/model/recognition/encoder.py

@@ -363,9 +363,9 @@ def forward(
             attention_mask = attention_mask + relative_position_bias
 
         attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_layer.contiguous(),
-            key_layer.contiguous(),
-            value_layer.contiguous(),
+            query_layer,
+            key_layer,
+            value_layer,
             attn_mask=attention_mask,
             dropout_p=self.dropout_p if self.training else 0.0,
             scale=self.attention_head_size**-0.5,

surya/model/recognition/tokenizer.py

@@ -1,7 +1,4 @@
-from itertools import chain
-import random
-from typing import List, Optional, Tuple, Union
-from tokenizers import AddedToken
+from typing import List, Union
 from transformers import ByT5Tokenizer
 import numpy as np
 import torch

surya/model/table_rec/decoder.py

@@ -172,9 +172,9 @@ def forward(
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
         attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states.contiguous(),
-            key_states.contiguous(),
-            value_states.contiguous(),
+            query_states,
+            key_states,
+            value_states,
             attn_mask=None,
             dropout_p=self.attention_dropout if self.training else 0.0,
             scale=self.head_dim**-0.5,
@@ -260,9 +260,9 @@ def forward(
                 causal_mask = torch.where(position_mask, torch.finfo(causal_mask.dtype).min, causal_mask)
 
         attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states.contiguous(),
-            key_states.contiguous(),
-            value_states.contiguous(),
+            query_states,
+            key_states,
+            value_states,
             attn_mask=causal_mask,
             dropout_p=self.attention_dropout if self.training else 0.0,
             scale=self.head_dim**-0.5,