taprosoft commited on
Commit
6c0ba50
·
1 Parent(s): a2b4d23

fix: monkey patch unimernet

Browse files
Files changed (2) hide show
  1. app.py +2 -1
  2. utils.py +15 -0
app.py CHANGED
@@ -2,8 +2,9 @@ import os
2
  import zipfile
3
  from collections import defaultdict
4
 
5
- from utils import prepare_env_mineru
6
 
 
7
  prepare_env_mineru() # noqa
8
 
9
 
 
2
  import zipfile
3
  from collections import defaultdict
4
 
5
+ from utils import patch_unimernet_model, prepare_env_mineru
6
 
7
+ patch_unimernet_model() # noqa
8
  prepare_env_mineru() # noqa
9
 
10
 
utils.py CHANGED
@@ -31,6 +31,21 @@ def trim_pages(pdf_path, output_path, start_page=0, trim_pages=5):
31
  return str(output_file_path)
32
 
33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
34
  def fix_problematic_imports():
35
  import sys
36
  import types
 
31
  return str(output_file_path)
32
 
33
 
34
+ def patch_unimernet_model():
35
+ from unimernet.models.unimernet.encoder_decoder import CustomMBartForCausalLM
36
+
37
+ # Save the original __init__ method
38
+ original_init = CustomMBartForCausalLM.__init__
39
+
40
+ # Define a new __init__ method
41
+ def new_init(self, config):
42
+ config._attn_implementation = "eager"
43
+ original_init(self, config)
44
+
45
+ # Monkey patch the __init__ method
46
+ CustomMBartForCausalLM.__init__ = new_init
47
+
48
+
49
  def fix_problematic_imports():
50
  import sys
51
  import types