Spaces:
Sleeping
Sleeping
| """Builds a small mixed-modality corpus for trying / testing the pipeline. | |
| Facts are planted so that each one is only answerable through a specific modality: | |
| * report.pdf p1 - plain text layer | |
| * report.pdf p2 - Asia-Pacific revenue figure exists ONLY inside an embedded chart image (needs OCR) | |
| * report.pdf p3 - scanned page, no text layer (needs page rasterisation + OCR) | |
| * handbook.docx - paragraphs, a table, and an embedded photo | |
| * images/ - real photographs (scikit-image sample data) + a text screenshot | |
| * notes.md - unrelated distractor text | |
| usage: python examples/make_demo_corpus.py OUT_DIR | |
| """ | |
| from __future__ import annotations | |
| import io | |
| import sys | |
| import textwrap | |
| from pathlib import Path | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt # noqa: E402 | |
| from matplotlib import font_manager # noqa: E402 | |
| from PIL import Image, ImageDraw, ImageFont # noqa: E402 | |
| FONT = font_manager.findfont("DejaVu Sans") | |
| def _text_image(lines: list[str], size=(1240, 1754), font_size=34, margin=90) -> Image.Image: | |
| im = Image.new("RGB", size, "white") | |
| d = ImageDraw.Draw(im) | |
| f = ImageFont.truetype(FONT, font_size) | |
| y = margin | |
| for line in lines: | |
| for w in textwrap.wrap(line, 52) or [""]: | |
| d.text((margin, y), w, fill="black", font=f) | |
| y += int(font_size * 1.5) | |
| y += font_size // 2 | |
| return im | |
| def _png(im: Image.Image) -> bytes: | |
| b = io.BytesIO() | |
| im.save(b, format="PNG") | |
| return b.getvalue() | |
| def chart_png() -> bytes: | |
| regions = ["North America", "Europe", "Asia-Pacific", "Latin America"] | |
| vals = [38.5, 27.1, 41.2, 9.8] | |
| fig, ax = plt.subplots(figsize=(7, 4.2), dpi=130) | |
| bars = ax.bar(regions, vals, color=["#4C72B0", "#55A868", "#C44E52", "#8172B2"]) | |
| for b, v in zip(bars, vals): | |
| ax.text(b.get_x() + b.get_width() / 2, v + 0.8, f"{v}", ha="center", fontsize=13) | |
| ax.set_title("Revenue by region FY2025 (USD million)", fontsize=14) | |
| ax.set_ylim(0, 50) | |
| ax.spines[["top", "right"]].set_visible(False) | |
| fig.tight_layout() | |
| b = io.BytesIO() | |
| fig.savefig(b, format="png") | |
| plt.close(fig) | |
| return b.getvalue() | |
| def make_pdf(path: Path): | |
| import pymupdf as fitz | |
| doc = fitz.open() | |
| p = doc.new_page() | |
| p.insert_textbox(fitz.Rect(60, 60, 540, 780), ( | |
| "Nimbus Robotics Annual Report 2025\n\n" | |
| "Nimbus Robotics designs autonomous warehouse robots. In fiscal year 2025 total revenue reached " | |
| "116.6 million US dollars, up 22 percent from the previous year. The company employs 640 people " | |
| "across offices in Pune, Rotterdam and Austin.\n\n" | |
| "The flagship product is the Nimbus R2, a mobile picking robot with a payload of 25 kilograms. " | |
| "The chief executive officer is Meera Iyer, who founded the company in 2016."), fontsize=11) | |
| p = doc.new_page() | |
| p.insert_textbox(fitz.Rect(60, 60, 540, 200), ( | |
| "Regional performance\n\nGrowth was strongest in one region this year; the chart below breaks down " | |
| "revenue by geography."), fontsize=11) | |
| p.insert_image(fitz.Rect(60, 210, 540, 500), stream=chart_png()) | |
| p = doc.new_page() # scanned page: image only, no text layer | |
| scan = _text_image([ | |
| "Appendix B - Service and Warranty", | |
| "Every Nimbus R2 unit ships with a 36-month limited warranty covering drive motors and the lidar module.", | |
| "Battery packs are covered for 18 months or 4000 charge cycles, whichever comes first.", | |
| ]) | |
| p.insert_image(p.rect, stream=_png(scan)) | |
| doc.save(path) | |
| def make_docx(path: Path, photo: Image.Image): | |
| import docx | |
| from docx.shared import Inches | |
| d = docx.Document() | |
| d.add_heading("Warehouse Operations Handbook", 1) | |
| d.add_paragraph("All staff must complete forklift certification before operating lifting equipment. " | |
| "Night shifts start at 22:00 and end at 06:00.") | |
| t = d.add_table(rows=1, cols=2) | |
| t.rows[0].cells[0].text, t.rows[0].cells[1].text = "Equipment", "Maximum load" | |
| for k, v in [("Electric forklift", "2500 kg"), ("Pallet jack", "1800 kg"), ("Nimbus R2 robot", "25 kg")]: | |
| r = t.add_row().cells | |
| r[0].text, r[1].text = k, v | |
| d.add_paragraph("The staff break room on level 2 has a coffee machine; the photo below shows it.") | |
| b = io.BytesIO() | |
| photo.save(b, format="PNG") | |
| b.seek(0) | |
| d.add_picture(b, width=Inches(3)) | |
| d.save(path) | |
| def build(out: str | Path) -> Path: | |
| from skimage import data as D | |
| out = Path(out) | |
| (out / "images").mkdir(parents=True, exist_ok=True) | |
| make_pdf(out / "report.pdf") | |
| make_docx(out / "handbook.docx", Image.fromarray(D.coffee())) | |
| Image.fromarray(D.chelsea()).save(out / "images" / "office_cat.jpg", quality=92) | |
| Image.fromarray(D.rocket()).save(out / "images" / "launch.jpg", quality=92) | |
| Image.fromarray(D.astronaut()).save(out / "images" / "visitor.png") | |
| _text_image(["NOTICE", "The server room door code is rotated every Monday at 09:00.", | |
| "Contact facilities on extension 4417."], size=(1100, 520), font_size=36, margin=50 | |
| ).save(out / "images" / "notice_board.png") | |
| (out / "notes.md").write_text( | |
| "# Reading notes\n\nPhotosynthesis converts light energy into chemical energy in chloroplasts. " | |
| "The Calvin cycle fixes carbon dioxide into sugars.\n", encoding="utf-8") | |
| return out | |
| if __name__ == "__main__": | |
| print(build(sys.argv[1] if len(sys.argv) > 1 else "demo_corpus")) | |