File size: 6,108 Bytes
1b8f835 | 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 | """PUP-T for one prompt: render, read the text back, rewrite the caption and render again (no training).
Same decisions as the measured research loop (p5/klein_rl/text_loop.py, wave 4):
- round 0 renders the PUP caption with the user seed;
- every candidate is read by the verifier and scored deterministically against the quoted texts of the prompt;
- the best candidate so far ranks by (1-CER, exact match), and ties keep the earliest;
- a prompt stops when its best candidate is solved (exact match and 1-CER = 1), when the rewrite is unusable, or
when the last round did not improve on the candidate already rewritten;
- otherwise the best caption is rewritten from the expected vs observed text, and CANDIDATES_PER_ROUND candidates
are rendered with seed, seed+1, ..., at most MAX_REFINES times.
Prompts without quoted text get the round-0 render only.
"""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from .quotes import enforce_quote_guardrail, extract_quoted_literals
from .typography import parse_vlm_ocr_readout, score_typography_ocr
MAX_REFINES = 2
CANDIDATES_PER_ROUND = 2
CAPTION_MIN_CHARS = 20
CAPTION_MAX_CHARS = 4000
EPS = 1e-9
REVISE_TASK = """An image generator was given this caption:
{caption}
The image must show exactly these texts: {expected}
A careful reading of the generated image found: {observed}
Rewrite the caption so the generator renders every required text exactly, letter for letter: keep the scene and
everything else that is requested; put each required text in double quotes with its exact spelling, casing and
punctuation; place it prominently on a large plain high-contrast surface in a simple bold sans-serif font; state that each
required text appears exactly once and nowhere else, with no other text in the image; keep every other detail of the caption and do not shorten it.
Answer with one JSON object only: {{"revised_caption": "..."}}"""
@dataclass
class Candidate:
round: int
cand: int
caption: str
seed: int
image: object
observed: list
ocr_reward_01: float
exact_match_rate: float
@dataclass
class Result:
image: object
best: Candidate | None
caption: str
expected: list
candidates: list = field(default_factory=list)
revisions: list = field(default_factory=list)
stop: str = ''
@property
def renders(self) -> int:
return max(1, len(self.candidates))
def parse_revision(text: str) -> str:
"""revised_caption from the verifier's JSON; '' when absent, malformed or out of bounds."""
decoder = json.JSONDecoder()
text = text or ''
pos = text.find('{')
while pos >= 0:
try:
obj, _end = decoder.raw_decode(text, pos)
except (ValueError, RecursionError):
obj = None
if isinstance(obj, dict) and 'revised_caption' in obj:
value = obj['revised_caption']
value = value.strip() if isinstance(value, str) else ''
return value if CAPTION_MIN_CHARS <= len(value) <= CAPTION_MAX_CHARS else ''
pos = text.find('{', pos + 1)
return ''
def _rank(c: Candidate):
return (c.ocr_reward_01, c.exact_match_rate >= 1.0 - EPS)
def _solved(c: Candidate) -> bool:
return c.exact_match_rate >= 1.0 - EPS and c.ocr_reward_01 >= 1.0 - EPS
def _best(candidates):
best = None
for c in candidates: # rounds and candidates in order, so ties keep the earliest
if best is None or _rank(c) > _rank(best):
best = c
return best
def caption_for(prompt: str, upsample) -> str:
"""PUP caption with every quoted literal of the prompt kept verbatim (the measured guardrail)."""
repaired, _check = enforce_quote_guardrail(prompt, upsample(prompt))
return repaired
def generate(prompt: str, seed: int, *, upsample, render, read=None, revise=None, caption: str | None = None,
max_refines: int = MAX_REFINES) -> Result:
"""upsample(prompt) -> caption; render(caption, seed) -> image; read(image) -> raw OCR reply;
revise(caption, expected, observed) -> raw JSON reply. read/revise are needed only for prompts with quoted text."""
expected = list(extract_quoted_literals(prompt))
if int(max_refines) < 0:
raise ValueError('max_refines must be >= 0')
if caption is None and upsample is None:
raise ValueError('upsample is required unless a caption is given')
if expected and (read is None or revise is None):
raise ValueError('prompts with quoted text need read and revise')
caption = caption if caption is not None else caption_for(prompt, upsample)
image0 = render(caption, seed)
if not expected:
return Result(image=image0, best=None, caption=caption, expected=[], stop='no_quoted_text')
def check(rnd, cand, cap, s, image):
observed = parse_vlm_ocr_readout(read(image))
score = score_typography_ocr(expected, observed)
return Candidate(rnd, cand, cap, s, image, observed, float(score['ocr_reward_01']),
float(score['exact_match_rate']))
result = Result(image=image0, best=None, caption=caption, expected=expected)
result.candidates.append(check(0, 0, caption, seed, image0))
for rnd in range(max_refines):
best = _best(result.candidates)
if _solved(best):
result.stop = 'solved'
break
if best.round != rnd:
result.stop = 'no_improvement'
break
revised = parse_revision(revise(best.caption, expected, best.observed))
result.revisions.append({'round': rnd + 1, 'from': [best.round, best.cand], 'revised_caption': revised})
if not revised:
result.stop = 'unusable_rewrite'
break
for j in range(CANDIDATES_PER_ROUND):
result.candidates.append(check(rnd + 1, j, revised, seed + j, render(revised, seed + j)))
else:
result.stop = 'solved' if _solved(_best(result.candidates)) else 'max_refines'
result.best = _best(result.candidates)
result.image = result.best.image
return result
|