Spaces:
Running
on
Zero
Running
on
Zero
wjm55
commited on
Commit
·
525b601
1
Parent(s):
52222e9
fixed ner output again
Browse files
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 |
-
#
|
144 |
-
highlighted_text
|
145 |
-
|
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 = """
|