Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """ | |
| Section-aware paper text for the poster SFT input (v2 of paper_side.load_paper_text). | |
| Why this exists. load_paper_text reads the LayoutLMv3 text boxes in page order | |
| and stops at 12,000 characters. Measured on the built dataset that cut lands | |
| inside Related Work / Preliminaries for almost every paper: only 7.5% of inputs | |
| contained any Experiments/Results section and 0.6% a Conclusion, while the | |
| posters the model is trained to emit are mostly Method / Results / Conclusion. | |
| Only 13.6% of poster text blocks shared a 7-gram with the input. | |
| What it does instead -- the way a person reads a paper before making a poster: | |
| 1. read every text box (same LayoutLMv3 classes, same reading order) up to | |
| References / Acknowledgements / Appendix, never a fixed character cut; | |
| 2. split the stream into sections at the `title` boxes, work out each | |
| heading's level (numbering) and *role* (introduction, background, method, | |
| experiments, conclusion, ...); | |
| 3. spend a per-role character budget: Introduction, Method, Experiments and | |
| Conclusion get most of it, Related Work gets none, and inside a section | |
| every subsection keeps its opening paragraphs rather than the first | |
| subsection eating the whole budget; | |
| 4. tell the model which section each figure/table lives in, so a figure in | |
| the Method section reads as the overview figure and one in Experiments as | |
| a result. | |
| The output is still plain text with `## ` headings, now `## [role] heading`, | |
| so the prompt format changes only inside the <paper> and <assets> blocks. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import re | |
| from pathlib import Path | |
| from typing import Dict, List, Optional, Tuple | |
| from paper_side import TEXT_CLASSES, reading_order | |
| # PyMuPDF returns the PDF's own characters: typographic ligatures (fi fl) and the | |
| # line-break hyphens of justified text ("calibra- tion") come through verbatim. | |
| # Measured on 800 built inputs: 6.5 hyphenation artefacts and ~3 ligatures per | |
| # 1000 words -- text the model would otherwise learn to copy into posters. | |
| _LIGATURES = {'\ufb00': 'ff', '\ufb01': 'fi', '\ufb02': 'fl', '\ufb03': 'ffi', '\ufb04': 'ffl', '\ufb05': 'st', '\ufb06': 'st'} | |
| _PUA = re.compile('[\ue000-\uf8ff]') | |
| _HYPHEN_BREAK = re.compile(r'\b([A-Za-z]+)- ([a-z]+)\b') | |
| def clean_text(t: str) -> str: | |
| for k, v in _LIGATURES.items(): # not NFKC: that would also flatten x^2 and R^d | |
| t = t.replace(k, v) | |
| t = _PUA.sub('', t) # private-use glyphs (broken symbol fonts) | |
| return re.sub(r'\s+', ' ', t).strip() | |
| def dehyphenate(texts: List[str]) -> List[str]: | |
| """Resolve line-break hyphens with the paper's own vocabulary. | |
| "calibra- tion" -> "calibration" (the joined word occurs elsewhere, or neither | |
| form does: a syllable break); "self- supervised" -> "self-supervised" (the | |
| hyphenated form occurs elsewhere in the paper: a real compound). | |
| """ | |
| joined_vocab, hyphen_vocab = set(), set() | |
| for t in texts: | |
| for w in re.findall(r'[A-Za-z]+(?:-[A-Za-z]+)+', t): | |
| hyphen_vocab.add(w.lower()) | |
| for w in re.findall(r'\b[A-Za-z]{4,}\b', t): | |
| joined_vocab.add(w.lower()) | |
| def fix(m): | |
| a, b = m.group(1), m.group(2) | |
| if (a + b).lower() in joined_vocab: | |
| return a + b | |
| if (a + '-' + b).lower() in hyphen_vocab: | |
| return a + '-' + b | |
| return a + b | |
| return [_HYPHEN_BREAK.sub(fix, t) for t in texts] | |
| # --------------------------------------------------------------------------- # | |
| # budgets (characters). ~3.4 chars/token for this kind of English prose. | |
| # --------------------------------------------------------------------------- # | |
| # | |
| # Measured with the Qwen3.8 tokenizer on the test split: this prose runs ~2.5 | |
| # chars/token (math, citations), so 22k chars ≈ 9k text tokens; with ~2.2k | |
| # visual tokens and a ~3.8k-token target the whole sample sits ~11.5k tokens | |
| # median, p90 ≈ 13.5k, inside a 16384 window with a small margin. | |
| ROLE_BUDGET = { | |
| 'abstract': 0, # already carried in the <paper> header | |
| 'introduction': 4000, | |
| 'related': 0, # posters never carry it | |
| 'background': 1500, | |
| 'method': 8000, | |
| 'experiments': 7000, | |
| 'conclusion': 2000, | |
| 'other': 1200, | |
| } | |
| TOTAL_BUDGET = 22000 | |
| SUBSECTION_FLOOR = 500 # every subsection keeps at least its opening | |
| SECTION_FLOOR = 800 | |
| STOP_ROLES = {'references', 'acknowledgements', 'appendix'} | |
| _RULES: List[Tuple[str, str]] = [ | |
| ('references', r'\b(references?|bibliography)\b'), | |
| ('acknowledgements', r'\backnowledg'), | |
| ('appendix', r'\b(appendix|appendices|supplementa)'), | |
| ('conclusion', r'\b(conclusions?|concluding|discussions?|limitations?|future work|' | |
| r'summary|broader impacts?|societal impact|ethic)'), | |
| ('experiments', r'\b(experiments?|experimental|results?|evaluations?|empirical|' | |
| r'benchmarks?|ablations?|case stud|applications?|findings)'), | |
| ('related', r'\b(related|prior|previous|existing) work|literature'), | |
| ('introduction', r'\bintroduction\b|\bmotivation\b'), | |
| ('abstract', r'\babstract\b'), | |
| ('background', r'\b(preliminar|background|notation|problem (setup|formulation|' | |
| r'statement|definition|setting)|setup|setting|definitions?|formulation)'), | |
| ] | |
| _NUM = re.compile(r'^\s*(\d+(?:\.\d+)*)[.\s]') | |
| _ROMAN = re.compile(r'^\s*(?:[IVX]+)[.\s]+[A-Z]') | |
| _LETTER = re.compile(r'^\s*[A-Z][.)]\s+[A-Z]') | |
| def _role_of(text: str) -> Optional[str]: | |
| t = text.lower() | |
| for role, pat in _RULES: | |
| if re.search(pat, t): | |
| return role | |
| return None | |
| def _level_of(text: str, role: Optional[str], have_l1: bool) -> int: | |
| m = _NUM.match(text) | |
| if m: | |
| return min(m.group(1).count('.') + 1, 3) | |
| if _ROMAN.match(text): | |
| return 1 | |
| if _LETTER.match(text): | |
| return 2 | |
| if role in ('abstract', 'introduction', 'related', 'conclusion') or role in STOP_ROLES: | |
| return 1 | |
| return 2 if have_l1 else 1 | |
| def _looks_like_heading(text: str) -> bool: | |
| if len(text) > 140 or len(text) < 3: | |
| return False | |
| if re.match(r'^(theorem|lemma|corollary|proposition|definition|proof|remark|' | |
| r'assumption|example|algorithm|figure|table)\b', text.lower()): | |
| return False | |
| return True | |
| # --------------------------------------------------------------------------- # | |
| # extraction | |
| # --------------------------------------------------------------------------- # | |
| def read_items(pid: str, root: str, pdf_path: str, max_pages: int = 12) -> Tuple[List[Dict], Dict]: | |
| """Text boxes in reading order: [{'kind','text','page','bbox','key'}], plus page info. | |
| `key` is a sortable reading-order position (page, column, y) shared with | |
| assets so a figure can be attributed to the section it sits in. | |
| """ | |
| d = Path(root) / pid | |
| layout_p = d / 'layout.json' | |
| if not layout_p.exists() or not pdf_path or not Path(pdf_path).exists(): | |
| return [], {} | |
| layout = json.loads(layout_p.read_text(encoding='utf-8')) | |
| meta = json.loads((d / 'assets.json').read_text(encoding='utf-8')) | |
| scale = 72.0 / meta.get('dpi', 200) | |
| import fitz | |
| doc = fitz.open(pdf_path) | |
| items: List[Dict] = [] | |
| pages: Dict[int, Dict] = {} | |
| try: | |
| for page_no in range(1, min(max_pages, len(doc)) + 1): | |
| boxes = [b for b in layout if b.get('page') == page_no | |
| and b.get('class_name') in TEXT_CLASSES] | |
| page = doc[page_no - 1] | |
| pw = page.rect.width / scale | |
| ph = page.rect.height / scale | |
| pages[page_no] = {'w': pw, 'h': ph} | |
| if not boxes: | |
| continue | |
| ordered = reading_order(boxes, pw) | |
| body_w = max(b['bbox'][2] - b['bbox'][0] for b in boxes) | |
| for b in ordered: | |
| x1, y1, x2, y2 = (v * scale for v in b['bbox']) | |
| t = page.get_text('text', clip=fitz.Rect(x1 - 1, y1 - 1, x2 + 1, y2 + 1)) | |
| t = clean_text(t) | |
| if len(t) < 3: | |
| continue | |
| items.append({'kind': 'title' if TEXT_CLASSES[b['class_name']] == 'title' else 'text', | |
| 'text': t, 'page': page_no, 'bbox': b['bbox'], | |
| 'key': order_key(b['bbox'], page_no, pw, body_w)}) | |
| finally: | |
| doc.close() | |
| fixed = dehyphenate([it['text'] for it in items]) | |
| for it, t in zip(items, fixed): | |
| it['text'] = t | |
| return items, pages | |
| def order_key(bbox, page_no: int, page_w: float, body_w: float) -> Tuple[int, int, float]: | |
| x1, y1, x2, _ = bbox | |
| spanning = (x2 - x1) > 0.75 * body_w | |
| col = 0 if spanning or (x1 + x2) / 2 < page_w / 2 else 1 | |
| return (page_no, col, float(y1)) | |
| def split_sections(items: List[Dict]) -> List[Dict]: | |
| """-> [{'heading','level','role','page','key','groups':[{'heading','level','paras':[str]}]}] | |
| Level-1 sections carry a role; deeper headings become groups inside the | |
| enclosing level-1 section. Text before the first real heading (paper | |
| title, authors, affiliations, and the abstract when it has no heading box) | |
| is kept in a 'front' section that the budget step drops. | |
| """ | |
| sections: List[Dict] = [] | |
| cur: Optional[Dict] = None | |
| seen_intro = seen_exp = False | |
| front = {'heading': '', 'level': 1, 'role': 'front', 'page': 0, 'key': (0, 0, 0.0), | |
| 'groups': [{'heading': '', 'level': 1, 'paras': []}]} | |
| def new_section(heading, level, role, it): | |
| return {'heading': heading, 'level': level, 'role': role, 'page': it['page'], | |
| 'key': it['key'], 'groups': [{'heading': '', 'level': level, 'paras': []}]} | |
| for it in items: | |
| if it['kind'] == 'title' and _looks_like_heading(it['text']): | |
| role = _role_of(it['text']) | |
| level = _level_of(it['text'], role, cur is not None) | |
| if level == 1: | |
| if role is None: | |
| # positional default: between intro and experiments = method, | |
| # after experiments = still experiments (Analysis, Discussion...) | |
| role = 'experiments' if seen_exp else ('method' if seen_intro else 'other') | |
| if role in STOP_ROLES: | |
| break | |
| if not seen_intro and role not in ('abstract', 'introduction') and not _NUM.match(it['text']): | |
| # paper title, author names, affiliations: front matter, not a section | |
| role = 'front' | |
| seen_intro |= role in ('introduction', 'abstract') | |
| seen_exp |= role == 'experiments' | |
| cur = new_section(it['text'], 1, role, it) | |
| sections.append(cur) | |
| else: | |
| if cur is None: # sub-heading before any section: treat as front text | |
| front['groups'][0]['paras'].append(it['text']) | |
| continue | |
| cur['groups'].append({'heading': it['text'], 'level': level, 'paras': []}) | |
| continue | |
| target = cur if cur is not None else front | |
| target['groups'][-1]['paras'].append(it['text']) | |
| return [front] + sections | |
| # --------------------------------------------------------------------------- # | |
| # budgeting | |
| # --------------------------------------------------------------------------- # | |
| def _cut(text: str, n: int) -> str: | |
| """Cut at a sentence boundary at or before n characters.""" | |
| if len(text) <= n: | |
| return text | |
| head = text[:n] | |
| m = max(head.rfind('. '), head.rfind('? '), head.rfind('! ')) | |
| if m > n * 0.5: | |
| return head[:m + 1] | |
| m = head.rfind(' ') | |
| return (head[:m] if m > 0 else head) + ' …' | |
| def _take_paras(paras: List[str], budget: int) -> Tuple[str, bool]: | |
| out, used, cut = [], 0, False | |
| for p in paras: | |
| if used + len(p) + 2 <= budget: | |
| out.append(p) | |
| used += len(p) + 2 | |
| continue | |
| room = budget - used | |
| if room > 200: | |
| out.append(_cut(p, room)) | |
| cut = True | |
| break | |
| return '\n\n'.join(out), cut | |
| def _spend(groups: List[Dict], budget: int) -> str: | |
| """Distribute a section's budget over its subsection groups, floor first.""" | |
| lens = [sum(len(p) + 2 for p in g['paras']) for g in groups] | |
| total = sum(lens) | |
| if total == 0: | |
| return '' | |
| if total <= budget: | |
| alloc = lens | |
| else: | |
| n = len(groups) | |
| floor = min(SUBSECTION_FLOOR, budget // max(n, 1)) | |
| alloc = [min(l, floor) for l in lens] | |
| rest = budget - sum(alloc) | |
| extra_need = [l - a for l, a in zip(lens, alloc)] | |
| need_total = sum(extra_need) | |
| if rest > 0 and need_total > 0: | |
| alloc = [a + int(rest * e / need_total) for a, e in zip(alloc, extra_need)] | |
| parts = [] | |
| for g, b in zip(groups, alloc): | |
| mark = '#' * max(g['level'] + 1, 3) # section = ##, subsections = ###, #### | |
| if not g['paras']: | |
| if g['heading']: | |
| parts.append(mark + ' ' + g['heading']) | |
| continue | |
| body, cut = _take_paras(g['paras'], b) | |
| if not body: | |
| continue | |
| head = (mark + ' ' + g['heading'] + '\n\n') if g['heading'] else '' | |
| parts.append(head + body + (' […]' if cut else '')) | |
| return '\n\n'.join(parts) | |
| def budget_body(sections: List[Dict], total: int = TOTAL_BUDGET, | |
| role_budget: Dict[str, int] = ROLE_BUDGET, | |
| abstract_missing: bool = False) -> Tuple[str, Dict]: | |
| """Render the sections into the body text under the per-role budgets.""" | |
| if total < TOTAL_BUDGET: | |
| # long poster / many figures: shrink every role's share proportionally | |
| # so the section mix stays the same and only the depth per section drops | |
| f = total / TOTAL_BUDGET | |
| role_budget = {k: int(v * f) for k, v in role_budget.items()} | |
| keep = [s for s in sections if s['role'] not in ('front', 'related', 'abstract')] | |
| if abstract_missing: | |
| for s in sections: | |
| if s['role'] in ('abstract', 'front'): | |
| keep.insert(0, dict(s, role='abstract')) | |
| break | |
| lens = {id(s): sum(len(p) + 2 for g in s['groups'] for p in g['paras']) for s in keep} | |
| by_role: Dict[str, List[Dict]] = {} | |
| for s in keep: | |
| by_role.setdefault(s['role'], []).append(s) | |
| alloc: Dict[int, int] = {} | |
| for role, secs in by_role.items(): | |
| b = role_budget.get(role, role_budget['other']) | |
| if role == 'abstract': | |
| b = 1800 | |
| need = [lens[id(s)] for s in secs] | |
| if sum(need) <= b: | |
| for s, n in zip(secs, need): | |
| alloc[id(s)] = n | |
| else: | |
| floor = min(SECTION_FLOOR, b // max(len(secs), 1)) | |
| a = [min(n, floor) for n in need] | |
| rest = b - sum(a) | |
| extra = [n - x for n, x in zip(need, a)] | |
| et = sum(extra) | |
| for i, s in enumerate(secs): | |
| alloc[id(s)] = a[i] + (int(rest * extra[i] / et) if et > 0 and rest > 0 else 0) | |
| # leftover (no conclusion, short intro...) goes to experiments, then method | |
| used = sum(alloc.values()) | |
| for role in ('experiments', 'method', 'introduction'): | |
| rest = total - used | |
| if rest <= 0: | |
| break | |
| for s in by_role.get(role, []): | |
| room = lens[id(s)] - alloc[id(s)] | |
| if room > 0: | |
| add = min(room, rest) | |
| alloc[id(s)] += add | |
| used += add | |
| rest -= add | |
| # hard cap | |
| if used > total: | |
| scale = total / used | |
| alloc = {k: int(v * scale) for k, v in alloc.items()} | |
| parts, stats = [], {} | |
| for s in keep: | |
| text = _spend(s['groups'], alloc[id(s)]) | |
| if not text.strip(): | |
| continue | |
| parts.append(f"## [{s['role']}] {s['heading']}".rstrip() + '\n\n' + text) | |
| stats[s['role']] = stats.get(s['role'], 0) + len(text) | |
| body = '\n\n'.join(parts) | |
| stats['total'] = len(body) | |
| stats['roles'] = sorted(by_role) | |
| return body, stats | |
| def load_paper_body(pid: str, root: str, pdf_path: str, abstract_missing: bool = False, | |
| max_pages: int = 12, total: int = TOTAL_BUDGET) -> Tuple[str, Dict, List[Dict]]: | |
| items, _ = read_items(pid, root, pdf_path, max_pages) | |
| if not items: | |
| return '', {'total': 0, 'roles': []}, [] | |
| sections = split_sections(items) | |
| body, stats = budget_body(sections, total=min(total, TOTAL_BUDGET), abstract_missing=abstract_missing) | |
| stats['budget'] = min(total, TOTAL_BUDGET) | |
| return body, stats, sections | |
| # --------------------------------------------------------------------------- # | |
| # assets: section attribution + caption fallback | |
| # --------------------------------------------------------------------------- # | |
| def section_of_asset(sections: List[Dict], page: int, bbox, page_w: float, body_w: float) -> Optional[Dict]: | |
| key = order_key(bbox, page, page_w, body_w) | |
| best = None | |
| for s in sections: | |
| if s['role'] == 'front': | |
| continue | |
| if s['key'] <= key: | |
| best = s | |
| return best | |
| _CAP_LINE = re.compile(r'^\s*(Figure|Fig\.|Table)\s*(\d+|[IVX]+)\s*[:.]', re.I) | |
| def section_by_reference(sections: List[Dict], caption: str, kind: str) -> Optional[Dict]: | |
| """The section whose text first cites this figure/table ("Figure 3", "Tab. 1"). | |
| Figures sit where the page layout allows (often at the top of the next | |
| column, under the next heading); they are *discussed* in the section that | |
| cites them, and that is the section a poster would place them in. | |
| """ | |
| m = _CAP_LINE.match(caption or '') | |
| if not m or kind not in ('figure', 'table'): | |
| return None | |
| num = m.group(2) | |
| if kind == 'table': | |
| pat = re.compile(r'\b(Table|Tab\.)\s*' + re.escape(num) + r'(?!\d)') | |
| else: | |
| pat = re.compile(r'\b(Figure|Fig\.|Figs?\.)\s*' + re.escape(num) + r'(?!\d)') | |
| for s in sections: | |
| if s['role'] in ('front', 'related'): | |
| continue | |
| for g in s['groups']: | |
| for p in g['paras']: | |
| if pat.search(p): | |
| return s | |
| return None | |
| def caption_fallback(doc, page_no: int, bbox, kind: str, scale: float, max_chars: int = 300) -> str: | |
| """When no LayoutLMv3 caption box matched: the nearest 'Figure N:'/'Table N:' | |
| text block on the page, below the asset (tables: above or below).""" | |
| if page_no > len(doc): | |
| return '' | |
| page = doc[page_no - 1] | |
| ph = page.rect.height / scale | |
| ax1, ay1, ax2, ay2 = bbox | |
| want = 'table' if kind == 'table' else 'figure' | |
| best, best_gap = '', 1e9 | |
| for blk in page.get_text('dict').get('blocks', []): | |
| lines = blk.get('lines') or [] | |
| if not lines: | |
| continue | |
| text = ' '.join(s['text'] for l in lines for s in l['spans']) | |
| text = re.sub(r'\s+', ' ', text).strip() | |
| m = _CAP_LINE.match(text) | |
| if not m: | |
| continue | |
| if (m.group(1).lower().startswith('tab')) != (want == 'table'): | |
| continue | |
| bx1, by1, bx2, by2 = (v / scale for v in blk['bbox']) | |
| lo, hi = max(ax1, bx1), min(ax2, bx2) | |
| if hi <= lo: | |
| continue | |
| gap = by1 - ay2 if by1 >= ay2 else ay1 - by2 | |
| if gap < -5 or gap > 0.2 * ph: | |
| continue | |
| if gap < best_gap: | |
| best, best_gap = text, gap | |
| if len(best) > max_chars: | |
| best = best[:max_chars].rsplit(' ', 1)[0] + '…' | |
| return best | |