wjm55 commited on
Commit
525b601
·
1 Parent(s): 52222e9

fixed ner output again

Browse files
Files changed (1) hide show
  1. app.py +11 -8
app.py CHANGED
@@ -24,6 +24,13 @@ DEFAULT_NER_LABELS = "person, organization, location, date, event"
24
  # "Qwen/Qwen2-VL-7B-Instruct": AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").cuda().eval()
25
 
26
  # }
 
 
 
 
 
 
 
27
  def array_to_image_path(image_array):
28
  # Convert numpy array to PIL Image
29
  img = Image.fromarray(np.uint8(image_array))
@@ -140,16 +147,12 @@ def run_example(image, model_id="Qwen/Qwen2-VL-7B-Instruct", run_ner=False, ner_
140
  if last_end < len(ocr_text):
141
  highlighted_text.append((ocr_text[last_end:], None))
142
 
143
- # Store the original text and entities as attributes of the highlighted_text list
144
- highlighted_text.original_text = ocr_text
145
- highlighted_text.entities = ner_results
146
-
147
- return highlighted_text
148
 
149
  # If NER is disabled, return the text without highlighting
150
- result = [(ocr_text, None)]
151
- result.original_text = ocr_text
152
- result.entities = []
153
  return result
154
 
155
  css = """
 
24
  # "Qwen/Qwen2-VL-7B-Instruct": AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", trust_remote_code=True, torch_dtype="auto", _attn_implementation="flash_attention_2").cuda().eval()
25
 
26
  # }
27
+
28
+ class TextWithMetadata(list):
29
+ def __init__(self, *args, **kwargs):
30
+ super().__init__(*args)
31
+ self.original_text = kwargs.get('original_text', '')
32
+ self.entities = kwargs.get('entities', [])
33
+
34
  def array_to_image_path(image_array):
35
  # Convert numpy array to PIL Image
36
  img = Image.fromarray(np.uint8(image_array))
 
147
  if last_end < len(ocr_text):
148
  highlighted_text.append((ocr_text[last_end:], None))
149
 
150
+ # Create TextWithMetadata instance with the highlighted text and metadata
151
+ result = TextWithMetadata(highlighted_text, original_text=ocr_text, entities=ner_results)
152
+ return result
 
 
153
 
154
  # If NER is disabled, return the text without highlighting
155
+ result = TextWithMetadata([(ocr_text, None)], original_text=ocr_text, entities=[])
 
 
156
  return result
157
 
158
  css = """