parse olmo output

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 import xml.etree.ElementTree as ET
 import os
 import torch
+import json
 from transformers import AutoProcessor, AutoModelForImageTextToText, pipeline, Qwen2VLForConditionalGeneration
 import spaces
 
@@ -259,6 +260,14 @@ def run_hf_ocr(image_path, model_name="RolmOCR"):
             generated_content = ocr_results[0]["generated_text"]
 
             if isinstance(generated_content, str):
+                # Check if it's JSON format from olmOCR
+                if model_name == "olmOCR":
+                    try:
+                        json_data = json.loads(generated_content)
+                        if "natural_text" in json_data:
+                            return json_data["natural_text"]
+                    except (json.JSONDecodeError, KeyError, TypeError):
+                        pass
                 return generated_content
 
             if isinstance(generated_content, list) and generated_content:
@@ -272,6 +281,14 @@ def run_hf_ocr(image_path, model_name="RolmOCR"):
                     ),
                     None,
                 ):
+                    # Check if it's JSON format from olmOCR
+                    if model_name == "olmOCR":
+                        try:
+                            json_data = json.loads(assistant_message)
+                            if "natural_text" in json_data:
+                                return json_data["natural_text"]
+                        except (json.JSONDecodeError, KeyError, TypeError):
+                            pass
                     return assistant_message
 
                 # Fallback if the specific assistant message structure isn't found but there's content
@@ -284,11 +301,27 @@ def run_hf_ocr(image_path, model_name="RolmOCR"):
                         and isinstance(generated_content[1], dict)
                         and "content" in generated_content[1]
                     ):
-                        return generated_content[1][
-                            "content"
-                        ]  # Assuming second part is assistant
+                        content = generated_content[1]["content"]
+                        # Check if it's JSON format from olmOCR
+                        if model_name == "olmOCR":
+                            try:
+                                json_data = json.loads(content)
+                                if "natural_text" in json_data:
+                                    return json_data["natural_text"]
+                            except (json.JSONDecodeError, KeyError, TypeError):
+                                pass
+                        return content  # Assuming second part is assistant
                     else:
-                        return generated_content[0]["content"]
+                        content = generated_content[0]["content"]
+                        # Check if it's JSON format from olmOCR
+                        if model_name == "olmOCR":
+                            try:
+                                json_data = json.loads(content)
+                                if "natural_text" in json_data:
+                                    return json_data["natural_text"]
+                            except (json.JSONDecodeError, KeyError, TypeError):
+                                pass
+                        return content
 
             print(f"Unexpected OCR output structure from HF model: {ocr_results}")
             return "Error: Could not parse OCR model output. Check console."