๐Ÿ“ฆ EqualifyEverything / equalify-iris-bench

๐Ÿ“„ prepare.mjs ยท 395 lines
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395// Stage 1: turn a CSV of PDF URLs into a corpus Iris can actually be asked to run.
//
// Iris does not fetch URLs โ€” `POST /v1/sessions` is multipart only โ€” so the URLs
// are client-side input and every file has to be downloaded here. Two things then
// have to happen before any of it is uploaded:
//
//   * Page cap. A PDF over the deployment's `max_pages` is rejected with a 400.
//     An unfiltered real-world corpus skews long, so dropping those would bias
//     the whole accuracy baseline toward short documents โ€” they are split into
//     cap-sized chunks instead, and labelled, because a chunk's review score is
//     not comparable to a whole document's (the reviewer sees a document with no
//     beginning).
//   * Everything that is not a PDF. At corpus scale a meaningful share of any URL
//     list is 404s, login walls, HTML error pages served as 200, and encrypted
//     files. Finding that out here costs a download; finding it out during the run
//     costs a session, a queue slot and a place in the failure statistics.
//
// Idempotent and resumable: every URL's outcome is appended to prepared.jsonl and
// skipped on a later pass, and downloads are content-addressed in cache/.

import { join, resolve } from "node:path";
import { existsSync, readFileSync, writeFileSync } from "node:fs";
import { fetchLimits, pxFromPts, RASTER_DPI } from "./limits.mjs";
import { appendJsonl, args, ensureDir, exec, hasCommand, log, num, pool, readJsonl, sha256 } from "./util.mjs";

const UA = "iris-bench/0.1 (+https://github.com/EqualifyEverything/iris-bench)";

// A ceiling on the download itself, not on what Iris accepts. Its purpose is to
// stop one pathological 400 MB scan from filling the disk; anything over it is
// recorded as skipped rather than quietly ignored.
const DEFAULT_MAX_DOWNLOAD_MB = 200;

// How many cap-sized chunks one oversize PDF may contribute. A 600-page document
// would otherwise become 24 sessions and dominate the corpus on its own. The
// dropped tail is recorded on the parent and logged in the summary โ€” a bound that
// does not announce itself reads as "we covered everything".
const DEFAULT_MAX_CHUNKS = 4;

// --- CSV -------------------------------------------------------------------

// RFC 4180 enough for real exports: quoted fields, embedded commas and newlines,
// doubled quotes.
function parseCsv(text) {
  const rows = [];
  let row = [];
  let field = "";
  let quoted = false;
  for (let i = 0; i < text.length; i++) {
    const c = text[i];
    if (quoted) {
      if (c === '"') {
        if (text[i + 1] === '"') {
          field += '"';
          i++;
        } else quoted = false;
      } else field += c;
    } else if (c === '"') quoted = true;
    else if (c === ",") {
      row.push(field);
      field = "";
    } else if (c === "\n") {
      row.push(field);
      field = "";
      rows.push(row);
      row = [];
    } else if (c !== "\r") field += c;
  }
  if (field.length || row.length) {
    row.push(field);
    rows.push(row);
  }
  return rows;
}

const isUrl = (s) => /^https?:\/\//i.test((s ?? "").trim());

// Which column holds the URLs, and whether row 0 is a header. Logged rather than
// assumed: a harness that silently picked the wrong column would produce a
// perfectly plausible run over the wrong 2000 things.
function pickUrls(rows) {
  if (!rows.length) return { urls: [], column: null };
  const header = rows[0];
  let column;
  let body;
  if (header.some(isUrl)) {
    column = header.findIndex(isUrl);
    body = rows;
    log(`csv: no header row detected; reading URLs from column ${column}`);
  } else {
    const named = header.findIndex((h) => /url|link|href|pdf|document/i.test(h));
    column = named === -1 ? 0 : named;
    body = rows.slice(1);
    log(`csv: header row detected; reading URLs from column ${column} ("${header[column]}")`);
  }
  const urls = [];
  const seen = new Set();
  let dupes = 0;
  for (const r of body) {
    const u = (r[column] ?? "").trim();
    if (!isUrl(u)) continue;
    if (seen.has(u)) {
      dupes++;
      continue;
    }
    seen.add(u);
    urls.push(u);
  }
  if (dupes) log(`csv: ${dupes} duplicate URL(s) collapsed`);
  return { urls, column };
}

// --- download + inspect ----------------------------------------------------

async function download(url, maxBytes) {
  let res;
  try {
    res = await fetch(url, {
      redirect: "follow",
      signal: AbortSignal.timeout(120_000),
      headers: { "user-agent": UA, accept: "application/pdf,*/*" },
    });
  } catch (e) {
    return { klass: "download_failed", error: `${e.name}: ${e.message}` };
  }
  if (!res.ok) return { klass: "download_failed", error: `http_${res.status}`, http_status: res.status };

  const declared = Number(res.headers.get("content-length") ?? 0);
  if (declared && declared > maxBytes) {
    return { klass: "too_large_bytes", error: `content-length ${declared} > ${maxBytes}`, bytes: declared };
  }
  let buf;
  try {
    buf = Buffer.from(await res.arrayBuffer());
  } catch (e) {
    return { klass: "download_failed", error: `body: ${e.message}` };
  }
  if (buf.length > maxBytes) {
    return { klass: "too_large_bytes", error: `${buf.length} bytes > ${maxBytes}`, bytes: buf.length };
  }
  // Magic bytes, not content-type. A great many servers hand out PDFs as
  // application/octet-stream, and a great many hand out an HTML "sign in" page as
  // application/pdf โ€” only one of those two errors is detectable from the header.
  if (!buf.subarray(0, 1024).includes("%PDF-")) {
    return {
      klass: "not_pdf",
      error: `no %PDF- header (content-type: ${res.headers.get("content-type") ?? "none"})`,
      bytes: buf.length,
      final_url: res.url,
    };
  }
  return { buf, bytes: buf.length, final_url: res.url, content_type: res.headers.get("content-type") };
}

function field(text, key) {
  const m = text.match(new RegExp(`^${key}:\\s*(.*)$`, "m"));
  return m ? m[1].trim() : null;
}

async function inspect(path, maxPages) {
  const info = await exec("pdfinfo", [path]);
  if (info.code !== 0) {
    return { klass: "pdfinfo_failed", error: info.stderr.trim().slice(0, 300) };
  }
  const pages = Number(field(info.stdout, "Pages"));
  if (!Number.isInteger(pages) || pages < 1) {
    return { klass: "pdfinfo_failed", error: `unreadable page count: ${field(info.stdout, "Pages")}` };
  }
  const encrypted = /^yes/i.test(field(info.stdout, "Encrypted") ?? "no");

  // Per-page physical sizes, for the pages that would actually be rasterized.
  // Iris measures each RENDERED page against the vision model's per-image byte
  // cap and its hard dimension ceiling, so a physically huge page (an ARCH-D
  // drawing) fails a 400 that no property of the file itself predicts.
  const sizes = await exec("pdfinfo", ["-f", "1", "-l", String(Math.min(pages, maxPages)), path]);
  let maxEdgePts = 0;
  for (const m of sizes.stdout.matchAll(/^Page\s+\d+ size:\s+([\d.]+) x ([\d.]+) pts/gm)) {
    maxEdgePts = Math.max(maxEdgePts, Number(m[1]), Number(m[2]));
  }
  return {
    pages,
    encrypted,
    max_edge_pts: maxEdgePts || null,
    predicted_max_edge_px: maxEdgePts ? pxFromPts(maxEdgePts) : null,
    pdf_version: field(info.stdout, "PDF version"),
    tagged: /^yes/i.test(field(info.stdout, "Tagged") ?? "no"),
  };
}

// --- splitting -------------------------------------------------------------

// qpdf if it is here, poppler otherwise. Both are ordinary installs and either
// will do; what matters is not silently declining to split, since that would
// re-introduce the long-document bias splitting exists to remove.
async function splitter() {
  if (await hasCommand("qpdf")) return "qpdf";
  if ((await hasCommand("pdfseparate")) && (await hasCommand("pdfunite"))) return "poppler";
  return null;
}

async function splitRange(tool, src, from, to, dest, tmpDir) {
  if (tool === "qpdf") {
    const r = await exec("qpdf", [src, "--pages", src, `${from}-${to}`, "--", dest]);
    // qpdf exits 3 on warnings while still writing a valid file, which is common
    // in real-world PDFs and must not read as a failure.
    return r.code === 0 || (r.code === 3 && existsSync(dest)) ? null : r.stderr.trim().slice(0, 300);
  }
  const sep = await exec("pdfseparate", ["-f", String(from), "-l", String(to), src, join(tmpDir, "p-%d.pdf")]);
  if (sep.code !== 0) return sep.stderr.trim().slice(0, 300);
  const parts = [];
  for (let p = from; p <= to; p++) {
    const f = join(tmpDir, `p-${p}.pdf`);
    if (existsSync(f)) parts.push(f);
  }
  if (!parts.length) return "pdfseparate produced no pages";
  const uni = await exec("pdfunite", [...parts, dest]);
  return uni.code === 0 ? null : uni.stderr.trim().slice(0, 300);
}

// --- main ------------------------------------------------------------------

async function main() {
  const a = args();
  if (!a.csv) {
    console.error(
      "usage: node src/prepare.mjs --csv urls.csv [--out .] [--concurrency 8]\n" +
        `                           [--max-download-mb ${DEFAULT_MAX_DOWNLOAD_MB}] [--max-chunks ${DEFAULT_MAX_CHUNKS}] [--limit N]`,
    );
    process.exit(2);
  }
  const base = process.env.IRIS_BASE_URL ?? "https://iris.equalify.uic.edu/v1";
  // Absolute, so the `path` recorded for every corpus item resolves from wherever the
  // run stage is later invoked. corpus.jsonl is machine-local by nature (it points at
  // gigabytes of cached PDFs), so there is nothing to gain from relative paths and a
  // resumed run started from the wrong directory to lose.
  const out = ensureDir(resolve(a.out ?? "."));
  const cache = ensureDir(join(out, "cache"));
  const chunkDir = ensureDir(join(cache, "chunks"));
  const preparedPath = join(out, "prepared.jsonl");
  const corpusPath = join(out, "corpus.jsonl");
  const maxBytes = num(a["max-download-mb"], DEFAULT_MAX_DOWNLOAD_MB) * 1024 * 1024;
  const maxChunks = num(a["max-chunks"], DEFAULT_MAX_CHUNKS);
  const concurrency = num(a.concurrency, 8);

  if (!(await hasCommand("pdfinfo"))) {
    console.error("pdfinfo not found. Install poppler-utils (brew install poppler / apt install poppler-utils).");
    process.exit(2);
  }
  const limits = await fetchLimits(base);
  const tool = await splitter();
  if (!tool) {
    log("WARNING: neither qpdf nor pdfseparate+pdfunite found โ€” oversize PDFs will be recorded, not split.");
    log("         That drops long documents from the corpus and biases it toward short ones. Install qpdf.");
  }

  const { urls } = pickUrls(parseCsv(readFileSync(a.csv, "utf8")));
  const done = new Map(readJsonl(preparedPath).map((r) => [r.url, r]));
  let todo = urls.filter((u) => !done.has(u));
  if (a.limit) todo = todo.slice(0, num(a.limit, todo.length));
  log(`csv: ${urls.length} unique URL(s); ${done.size} already prepared; ${todo.length} to fetch`);

  // sha256 -> url, so the same file behind two URLs is downloaded twice (cheap,
  // and unavoidable without fetching) but run once.
  const bySha = new Map();
  for (const r of done.values()) if (r.sha256 && !r.duplicate_of) bySha.set(r.sha256, r.url);

  let n = 0;
  await pool(todo, concurrency, async (url) => {
    const record = { url, prepared_at: new Date().toISOString() };
    try {
      const dl = await download(url, maxBytes);
      if (dl.klass) {
        appendJsonl(preparedPath, { ...record, ...dl });
        return;
      }
      const sha = sha256(dl.buf);
      record.sha256 = sha;
      record.bytes = dl.bytes;
      record.final_url = dl.final_url;
      record.content_type = dl.content_type;

      const existing = bySha.get(sha);
      if (existing && existing !== url) {
        // Byte-identical to something already in the corpus: keep the record so the
        // URL is accounted for, but do not run it twice.
        appendJsonl(preparedPath, { ...record, klass: "duplicate", duplicate_of: existing });
        return;
      }
      bySha.set(sha, url);

      const path = join(cache, `${sha}.pdf`);
      if (!existsSync(path)) writeFileSync(path, dl.buf);

      const info = await inspect(path, limits.maxPages);
      Object.assign(record, info);
      if (info.klass) {
        appendJsonl(preparedPath, record);
        return;
      }
      if (info.encrypted) {
        appendJsonl(preparedPath, { ...record, klass: "encrypted" });
        return;
      }

      // A page that will rasterize past the model's hard dimension ceiling is a
      // predicted 400. Flagged, never excluded: the DPI behind the prediction is
      // Iris-internal and unpublished, so this can be wrong, and being wrong must
      // cost a failed run rather than a missing document.
      const risks = [];
      if (limits.maxDimensionPx && info.predicted_max_edge_px > limits.maxDimensionPx) {
        risks.push(`predicted ${info.predicted_max_edge_px}px long edge at ${RASTER_DPI}dpi > max_dimension_px`);
      }

      if (info.pages <= limits.maxPages) {
        appendJsonl(preparedPath, { ...record, klass: "ok", id: sha, path, risks });
        return;
      }

      // Oversize: split into cap-sized chunks.
      if (!tool) {
        appendJsonl(preparedPath, { ...record, klass: "oversize_pages", risks, chunks: 0, chunks_dropped: null });
        return;
      }
      const wanted = Math.ceil(info.pages / limits.maxPages);
      const made = [];
      let failure = null;
      for (let c = 0; c < Math.min(wanted, maxChunks); c++) {
        const from = c * limits.maxPages + 1;
        const to = Math.min(info.pages, from + limits.maxPages - 1);
        const id = `${sha}-p${from}-${to}`;
        const dest = join(chunkDir, `${id}.pdf`);
        const tmp = ensureDir(join(chunkDir, `.tmp-${id}`));
        const err = existsSync(dest) ? null : await splitRange(tool, path, from, to, dest, tmp);
        if (err) {
          failure = err;
          break;
        }
        made.push({ id, from, to, path: dest, pages: to - from + 1 });
      }
      appendJsonl(preparedPath, {
        ...record,
        klass: "oversize_pages",
        risks,
        split_with: tool,
        chunks: made.length,
        chunks_dropped: Math.max(0, wanted - made.length),
        split_error: failure,
      });
      for (const c of made) {
        appendJsonl(preparedPath, {
          url,
          prepared_at: record.prepared_at,
          klass: "ok",
          id: c.id,
          path: c.path,
          sha256: sha,
          parent_sha: sha,
          parent_pages: info.pages,
          page_from: c.from,
          page_to: c.to,
          pages: c.pages,
          bytes: null,
          risks,
        });
      }
    } catch (e) {
      appendJsonl(preparedPath, { ...record, klass: "prepare_error", error: `${e.name}: ${e.message}` });
    } finally {
      if (++n % 25 === 0) log(`prepared ${n}/${todo.length}`);
    }
  });

  // corpus.jsonl is the runnable subset, rewritten from scratch each pass so it is
  // a deterministic function of prepared.jsonl rather than an append-only history.
  const all = readJsonl(preparedPath);
  const runnable = all.filter((r) => r.klass === "ok");
  runnable.sort((x, y) => (x.id < y.id ? -1 : x.id > y.id ? 1 : 0)); // stable staging order
  writeFileSync(corpusPath, runnable.map((r) => JSON.stringify(r)).join("\n") + (runnable.length ? "\n" : ""));

  const counts = {};
  for (const r of all) counts[r.klass ?? "unknown"] = (counts[r.klass ?? "unknown"] ?? 0) + 1;
  const chunked = runnable.filter((r) => r.parent_sha).length;
  const droppedChunks = all.reduce((s, r) => s + (r.chunks_dropped ?? 0), 0);
  const atRisk = runnable.filter((r) => r.risks?.length).length;
  const pages = runnable.reduce((s, r) => s + (r.pages ?? 0), 0);

  log("--- corpus ---");
  for (const [k, v] of Object.entries(counts).sort((x, y) => y[1] - x[1])) log(`  ${k}: ${v}`);
  log(`  runnable: ${runnable.length} session(s), ${pages} page(s) โ€” ${chunked} of them chunks of oversize PDFs`);
  if (droppedChunks) log(`  NOT covered: ${droppedChunks} chunk(s) past --max-chunks=${maxChunks}`);
  if (atRisk) log(`  ${atRisk} runnable item(s) carry a predicted-rejection risk flag (see .risks)`);
  log(`wrote ${corpusPath}`);
}

await main();