Spaces:
Running
Running
File size: 4,761 Bytes
b611f38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | """
Post-processing utilities for OCR outputs
Strips reasoning/thinking tokens and normalizes Text / Markdown / JSON formats.
"""
import re
import json
from typing import Dict, Any, Tuple, Optional
def clean_reasoning_and_special_tokens(text: str) -> str:
"""
Strips internal model reasoning/thinking tokens like <think>...</think>,
<thought>...</thought>, <|thought|>...<|/thought|>, etc., leaving only the final answer.
"""
if not isinstance(text, str):
return str(text)
# Remove standard thinking tags (multiline, dotall)
text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE)
text = re.sub(r"<thought>.*?</thought>", "", text, flags=re.DOTALL | re.IGNORECASE)
text = re.sub(r"<\|thought\|>.*?<\|/thought\|>", "", text, flags=re.DOTALL | re.IGNORECASE)
text = re.sub(r"\[THINKING\].*?\[/THINKING\]", "", text, flags=re.DOTALL | re.IGNORECASE)
text = re.sub(r"<\|reasoning\|>.*?<\|/reasoning\|>", "", text, flags=re.DOTALL | re.IGNORECASE)
# Remove dangling open tags if any
text = re.sub(r"<think>.*", "", text, flags=re.DOTALL | re.IGNORECASE)
# Strip common end of generation or special system tokens
special_tokens = [
"<|im_end|>", "<|endoftext|>", "<|eot_id|>", "<|end_of_text|>",
"<|output|>", "<|input|>", "</s>", "<s>"
]
for token in special_tokens:
text = text.replace(token, "")
return text.strip()
def extract_json_from_text(text: str) -> Optional[Dict[str, Any]]:
"""
Attempts to extract and parse JSON object/array from a model output string.
"""
if not text:
return None
cleaned = clean_reasoning_and_special_tokens(text)
# Direct json load check
try:
return json.loads(cleaned)
except Exception:
pass
# Look for ```json ... ``` codeblocks
json_block = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned)
if json_block:
try:
return json.loads(json_block.group(1).strip())
except Exception:
pass
# Look for { ... } or [ ... ]
brace_match = re.search(r"(\{[\s\S]*\}|\[[\s\S]*\])", cleaned)
if brace_match:
try:
return json.loads(brace_match.group(1).strip())
except Exception:
pass
return None
def format_ocr_output(raw_output: Any, default_type: str = "markdown") -> Dict[str, Any]:
"""
Normalizes any OCR model raw response into standardized dictionary:
{
'text': str,
'markdown': str,
'json': Optional[Any],
'output_type': 'markdown' | 'text' | 'json'
}
"""
if isinstance(raw_output, dict):
if "text" in raw_output or "markdown" in raw_output or "json" in raw_output:
text_val = clean_reasoning_and_special_tokens(raw_output.get("text", ""))
md_val = clean_reasoning_and_special_tokens(raw_output.get("markdown", text_val))
json_val = raw_output.get("json", None)
if json_val is None and md_val:
json_val = extract_json_from_text(md_val)
return {
"text": text_val,
"markdown": md_val,
"json": json_val,
"output_type": raw_output.get("output_type", default_type)
}
else:
# The raw output is a json object itself
json_str = json.dumps(raw_output, indent=2, ensure_ascii=False)
return {
"text": json_str,
"markdown": f"```json\n{json_str}\n```",
"json": raw_output,
"output_type": "json"
}
if isinstance(raw_output, (list, tuple)):
# List of items or detections
json_str = json.dumps(raw_output, indent=2, ensure_ascii=False)
lines = []
for item in raw_output:
if isinstance(item, str):
lines.append(item)
elif isinstance(item, dict) and "text" in item:
lines.append(str(item["text"]))
else:
lines.append(str(item))
text_content = "\n".join(lines)
return {
"text": text_content,
"markdown": text_content,
"json": raw_output,
"output_type": default_type
}
# Otherwise raw_output is a string
cleaned_text = clean_reasoning_and_special_tokens(str(raw_output))
parsed_json = extract_json_from_text(cleaned_text)
output_type = default_type
if parsed_json is not None and default_type == "json":
output_type = "json"
return {
"text": cleaned_text,
"markdown": cleaned_text,
"json": parsed_json,
"output_type": output_type
}
|