RAG, Measured: the code

These are the files that do most of the work. The comments explain why each part works the way it does. If anything is unclear, or you would do it differently, tell me. I will also send you the full code for free, with nothing to sign up for. Ask for the full code

export type Chunk = { id: string; text: string; start: number; end: number; parent?: string; page?: number; pages?: number[] };

type Span = { start: number; end: number };

const SEPARATORS = ["\n\n", "\n", " ", ""];

// Pieces are contiguous slices, with each separator kept on the end of the piece before it.
function pieces(text: string, start: number, end: number, size: number, seps: string[]): Span[] {
  const sepIndex = seps.findIndex((s) => s === "" || text.slice(start, end).includes(s));
  const sep = seps[sepIndex];
  const rest = seps.slice(sepIndex + 1);
  const out: Span[] = [];
  let from = start;
  const emit = (to: number) => {
    if (to <= from) return;
    if (to - from > size && rest.length > 0) out.push(...pieces(text, from, to, size, rest));
    else out.push({ start: from, end: to });
    from = to;
  };
  if (sep === "") {
    for (let i = start + 1; i <= end; i++) emit(i);
    return out;
  }
  let at = text.indexOf(sep, from);
  while (at !== -1 && at < end) {
    emit(Math.min(at + sep.length, end));
    at = text.indexOf(sep, from);
  }
  emit(end);
  return out;
}

function trimmed(text: string, span: Span): Span {
  let { start, end } = span;
  while (start < end && /\s/.test(text[start])) start++;
  while (end > start && /\s/.test(text[end - 1])) end--;
  return { start, end };
}

export function recursiveChunks(text: string, size: number, overlap: number, prefix = "c"): Chunk[] {
  const parts = pieces(text, 0, text.length, size, SEPARATORS);
  const windows: Span[] = [];
  let current: Span[] = [];
  const length = () => (current.length ? current[current.length - 1].end - current[0].start : 0);
  for (const p of parts) {
    const len = p.end - p.start;
    if (current.length && length() + len > size) {
      windows.push({ start: current[0].start, end: current[current.length - 1].end });
      while (current.length && (length() > overlap || length() + len > size)) current.shift();
    }
    current.push(p);
  }
  if (current.length) windows.push({ start: current[0].start, end: current[current.length - 1].end });
  const out: Chunk[] = [];
  for (const w of windows) {
    const t = trimmed(text, w);
    if (t.end > t.start) out.push({ id: `${prefix}${out.length}`, text: text.slice(t.start, t.end), start: t.start, end: t.end });
  }
  return out;
}

export function paragraphChunks(paras: { text: string; page?: number }[], minChars = 200, prefix = "p"): Chunk[] {
  const out: Chunk[] = [];
  let offset = 0;
  let buf: { start: number; end: number; page?: number } | null = null;
  const source = paras.map((p) => p.text).join("\n\n");
  const flush = () => {
    if (!buf) return;
    out.push({ id: `${prefix}${out.length}`, text: source.slice(buf.start, buf.end), start: buf.start, end: buf.end, page: buf.page });
    buf = null;
  };
  for (const p of paras) {
    const span = { start: offset, end: offset + p.text.length };
    offset = span.end + 2;
    if (buf) buf.end = span.end;
    else buf = { ...span, page: p.page };
    if (buf.end - buf.start >= minChars) flush();
  }
  if (buf) {
    const last = out[out.length - 1];
    if (last) {
      last.end = (buf as Span).end;
      last.text = source.slice(last.start, last.end);
    } else flush();
  }
  return out;
}

export function sentenceChunks(parents: Chunk[]): Chunk[] {
  const out: Chunk[] = [];
  for (const parent of parents) {
    const re = /(?<=[.!?])\s+(?=[A-Z0-9])/g;
    let from = 0;
    let i = 0;
    const push = (to: number) => {
      const t = parent.text.slice(from, to).trim();
      if (t) {
        const s = parent.start + parent.text.indexOf(t, from);
        out.push({ id: `${parent.id}.${i++}`, text: t, start: s, end: s + t.length, parent: parent.id, page: parent.page });
      }
    };
    for (const m of parent.text.matchAll(re)) {
      push(m.index);
      from = m.index + m[0].length;
    }
    push(parent.text.length);
  }
  return out;
}

import { words } from "../core/text.ts";

export function shingles(s: string, n = 5): Set<string> {
  const w = words(s);
  if (w.length < n) return new Set(w.length ? [w.join(" ")] : []);
  const out = new Set<string>();
  for (let i = 0; i + n <= w.length; i++) out.add(w.slice(i, i + n).join(" "));
  return out;
}

function share(of: Set<string>, within: Set<string>): number {
  if (of.size === 0) return 0;
  let n = 0;
  for (const s of of) if (within.has(s)) n++;
  return n / of.size;
}

export function qasperHit(chunk: string, evidence: string[], threshold = 0.5): boolean {
  const c = shingles(chunk);
  return evidence.some((e) => {
    const s = shingles(e);
    return share(s, c) >= threshold || share(c, s) >= threshold;
  });
}

export function financeHit(chunkPage: number, evidencePages: number[]): boolean {
  return evidencePages.includes(chunkPage);
}

// Import this first in every Node entry point. Transformers.js otherwise runs Node on the native ONNX
// Runtime, whose quantised kernels give vectors up to 0.02 away from the browser's WebAssembly ones.
import * as ort from "onnxruntime-web";

ort.env.wasm.numThreads = 1;
// Local files; left unset, Transformers.js points the runtime at a CDN.
ort.env.wasm.wasmPaths = new URL("../../node_modules/onnxruntime-web/dist/", import.meta.url).href;
(globalThis as Record<symbol, unknown>)[Symbol.for("onnxruntime")] = ort;

Ask for the full codeBack to the case study