fixed the issue with entities and text output

app.py

@@ -140,10 +140,17 @@ def run_example(image, model_id="Qwen/Qwen2-VL-7B-Instruct", run_ner=False, ner_
         if last_end < len(ocr_text):
             highlighted_text.append((ocr_text[last_end:], None))
         
+        # Store the original text and entities as attributes of the highlighted_text list
+        highlighted_text.original_text = ocr_text
+        highlighted_text.entities = ner_results
+        
         return highlighted_text
     
     # If NER is disabled, return the text without highlighting
-    return [(ocr_text, None)]
+    result = [(ocr_text, None)]
+    result.original_text = ocr_text
+    result.entities = []
+    return result
 
 css = """
   /* Overall app styling */
@@ -265,37 +272,23 @@ with gr.Blocks(css=css) as demo:
             if isinstance(image, np.ndarray):
                 image = Image.fromarray(image)
             
-            # Create a temporary directory
             with tempfile.TemporaryDirectory() as temp_dir:
                 # Save image
                 img_path = os.path.join(temp_dir, f"{fname}.png")
                 image.save(img_path)
                 
-                # Extract text and entities from text_data
-                full_text = ""
-                entities = []
-                current_pos = 0
-                
-                # Process the highlighted text data
-                for segment, label in text_data:
-                    full_text += segment
-                    if label:  # If this segment has a label (is an entity)
-                        entities.append({
-                            "text": segment,
-                            "label": label,
-                            "start": current_pos,
-                            "end": current_pos + len(segment)
-                        })
-                    current_pos += len(segment)
+                # Get the original OCR text and entities
+                original_text = getattr(text_data, 'original_text', '')
+                entities = getattr(text_data, 'entities', [])
                 
                 # Save text
                 txt_path = os.path.join(temp_dir, f"{fname}.txt")
                 with open(txt_path, 'w', encoding='utf-8') as f:
-                    f.write(full_text)
+                    f.write(original_text)
                 
                 # Create JSON with text and entities
                 json_data = {
-                    "text": full_text,
+                    "text": original_text,
                     "entities": entities,
                     "image_file": f"{fname}.png"
                 }