File size: 4,828 Bytes
5b97e57 | 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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | import json
import re
from dataclasses import dataclass
@dataclass(frozen=True)
class OutputCleanConfig:
mode: str = "prompt"
strip_think: bool = True
strip_code_fences: bool = True
strip_role_prefixes: bool = True
strip_json_wrappers: bool = True
strip_leading_preamble: bool = True
strip_planning: bool = True
keep_first_paragraph_only: bool = False
_ROLE_PREFIX_RE = re.compile(r"^\s*(assistant|final|output|response|result|prompt)\s*:\s*", re.IGNORECASE)
_CODE_FENCE_RE = re.compile(r"^\s*```[\w-]*\s*$", re.IGNORECASE)
_THINK_BLOCK_RE = re.compile(r"<think[^>]*>.*?</think>", flags=re.IGNORECASE | re.DOTALL)
_THINK_OPEN_RE = re.compile(r"<think[^>]*>", flags=re.IGNORECASE)
_THINK_CLOSE_RE = re.compile(r"</think\s*>", flags=re.IGNORECASE)
_MARKER_RE = re.compile(
r"(?im)^\s*(final|final answer|answer|output|result|prompt)\s*[:\-]\s*",
)
_IM_TOKEN_RE = re.compile(
r"(?i)<\|?im_(start|end)\|?>|<im_(start|end)>|<\|endoftext\|>",
)
_PLANNING_RE = re.compile(
r"(?is)\b("
r"i\s+(should|need|must|will|want|am\s+going\s+to|have\s+to)\b|"
r"let's\b|"
r"first\b|next\b|then\b|"
r"wait\b|"
r"so\s+i\s+need\s+to\b|"
r"i\s+should\s+focus\s+on\b"
r")"
)
def clean_model_output(text: str, config: OutputCleanConfig | None = None) -> str:
if not text:
return ""
cfg = config or OutputCleanConfig()
cleaned = (text or "").strip()
# Remove common chat template tokens that sometimes leak into output.
cleaned = _IM_TOKEN_RE.sub("", cleaned).strip()
if cfg.strip_think:
cleaned = _THINK_BLOCK_RE.sub("", cleaned)
cleaned = _THINK_CLOSE_RE.sub("", cleaned)
if _THINK_OPEN_RE.search(cleaned):
cleaned = _THINK_OPEN_RE.sub("", cleaned)
parts = re.split(r"\n\s*\n", cleaned, maxsplit=1)
if len(parts) == 2:
cleaned = parts[1]
cleaned = cleaned.strip()
cleaned = _IM_TOKEN_RE.sub("", cleaned).strip()
if cfg.strip_code_fences and "```" in cleaned:
lines = [ln for ln in cleaned.splitlines() if not _CODE_FENCE_RE.match(ln)]
cleaned = "\n".join(lines).strip()
if cfg.strip_json_wrappers:
maybe = _extract_from_json(cleaned, mode=cfg.mode)
if maybe is not None:
cleaned = maybe.strip()
if cfg.strip_leading_preamble:
cleaned = _drop_preamble(cleaned).strip()
if cfg.strip_planning and cfg.mode == "prompt":
without_planning = _strip_planning_paragraphs(cleaned)
if without_planning:
cleaned = without_planning
if cfg.strip_role_prefixes:
lines = cleaned.splitlines()
if lines:
lines[0] = _ROLE_PREFIX_RE.sub("", lines[0])
cleaned = "\n".join(lines).strip()
cleaned = _MARKER_RE.sub("", cleaned).strip()
if cfg.keep_first_paragraph_only:
parts = re.split(r"\n\s*\n", cleaned, maxsplit=1)
cleaned = parts[0].strip()
return cleaned
def _extract_from_json(text: str, mode: str) -> str | None:
candidate = text.strip()
if not candidate:
return None
if not (candidate.startswith("{") and candidate.endswith("}")):
return None
try:
payload = json.loads(candidate)
except Exception:
return None
if not isinstance(payload, dict):
return None
preferred_keys = []
if mode == "prompt":
preferred_keys = ["prompt", "final", "output", "text", "content"]
else:
preferred_keys = ["text", "content", "output", "final"]
for key in preferred_keys:
value = payload.get(key)
if isinstance(value, str) and value.strip():
return value
return None
def _drop_preamble(text: str) -> str:
lines = text.splitlines()
if not lines:
return text
keep_from = 0
for i, ln in enumerate(lines[:20]):
if _MARKER_RE.match(ln):
keep_from = i
if re.search(r"(?i)\bhere(?:'s| is)\b", ln) and i < 6:
keep_from = max(keep_from, i + 1)
return "\n".join(lines[keep_from:]).strip()
def _strip_planning_paragraphs(text: str) -> str:
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", (text or "").strip()) if p.strip()]
if not paragraphs:
return ""
kept: list[str] = []
dropping = True
for p in paragraphs:
is_planning = bool(_PLANNING_RE.search(p))
if dropping and is_planning:
continue
dropping = False
kept.append(p)
# If everything looked like planning, don't destroy the output.
if not kept:
return text.strip()
return "\n\n".join(kept).strip()
|