Retrieval by words

import { fileURLToPath } from "node:url";
import { LAWS, type Law } from "./corpus.js";

// Retrieval by words: deterministic, in process, over the state's own file.
// The model sees the laws this returns and nothing else from the corpus, so
// a test can say exactly which laws a question reaches.

/** How many laws a question is answered from: enough that a paraphrase still reaches its law, few enough that the model is not reading a chapter. */
export const RETRIEVED_LAWS = 8;

// Words that say nothing about which law is meant. "park" and "far" stay: in
// a question about road law they are the whole question.
const STOP = new Set(
  "a an and are as at be by can could do does for from how i if in is it its me my of on or should so that the their them there this to was what when where which who why will with would you your".split(
    " ",
  ),
);

export function words(text: string): string[] {
  return (text.toLowerCase().match(/[a-z0-9]+/g) ?? [])
    .filter((word) => !STOP.has(word))
    // "hydrants" must find "hydrant"; "bus" must not become "bu".
    .map((word) =>
      word.length > 3 && word.endsWith("s") && !word.endsWith("ss") ? word.slice(0, -1) : word,
    );
}

// What BM25 needs of a laws array that does not depend on the question: each
// law's word counts and length, the average length, and how many laws each
// word is in. Built once per array and kept, because the pipeline retrieves
// over the same array several times a question.
interface Index {
  length: number;
  /** How many laws are ranked by words: every one but the state's own pages. */
  size: number;
  docs: { length: number; counts: Map<string, number> }[];
  average: number;
  frequency: Map<string, number>;
}

const INDEXES = new WeakMap<Law[], Index>();

// A row of the state's own page is not ranked by words. Every page row carries
// its family's question, and five of them a family would fill the shortlist
// for any question that shares words with it. A page row reaches the model
// by its family's key and by meaning, never by words.
function ranked(law: Law): boolean {
  return law.kind !== "process";
}

function indexOf(laws: Law[]): Index {
  const kept = INDEXES.get(laws);
  if (kept !== undefined && kept.length === laws.length) return kept;
  const frequency = new Map<string, number>();
  const docs = laws.map((law) => {
    const tokens = ranked(law) ? words(`${law.law} ${law.question} ${law.concept.replace(/-/g, " ")}`) : [];
    const counts = new Map<string, number>();
    for (const word of tokens) counts.set(word, (counts.get(word) ?? 0) + 1);
    for (const word of counts.keys()) frequency.set(word, (frequency.get(word) ?? 0) + 1);
    return { length: tokens.length, counts };
  });
  const counted = docs.filter((_, index) => ranked(laws[index]));
  const built: Index = {
    length: laws.length,
    size: counted.length,
    docs,
    average: counted.reduce((sum, doc) => sum + doc.length, 0) / Math.max(1, counted.length),
    frequency,
  };
  INDEXES.set(laws, built);
  return built;
}

/** The laws most likely to answer `question` by the words they share with it, best first. A law sharing no word is never returned. */
export function bm25(laws: Law[], question: string, limit = RETRIEVED_LAWS): Law[] {
  if (laws.length === 0) return [];
  const { docs, average, frequency, size } = indexOf(laws);
  const asked = [...new Set(words(question).flatMap((word) => spelledAs(word, frequency)))];
  if (asked.length === 0) return [];
  const k1 = 1.2;
  const b = 0.75;
  return docs
    .map((doc, index) => {
      let score = 0;
      for (const word of asked) {
        const n = frequency.get(word);
        if (n === undefined) continue;
        const tf = doc.counts.get(word) ?? 0;
        if (tf === 0) continue;
        const idf = Math.log(1 + (size - n + 0.5) / (n + 0.5));
        score += (idf * tf * (k1 + 1)) / (tf + k1 * (1 - b + (b * doc.length) / average));
      }
      return { law: laws[index], score };
    })
    .filter((scored) => scored.score > 0)
    .sort((a, b) => b.score - a.score)
    .slice(0, limit)
    .map((scored) => scored.law);
}

// The corpus's words a learner's word stands for. A word the laws use is
// itself. One they do not is matched to the laws' words it contains
// ("firehydrant" → fire, hydrant) or is one letter away from ("behing" →
// behind), because a learner types on a phone.
function spelledAs(word: string, vocabulary: Map<string, number>): string[] {
  if (vocabulary.has(word)) return [word];
  if (word.length < 4) return [];
  const found: string[] = [];
  for (const known of vocabulary.keys()) {
    if (known.length >= 4 && word.includes(known)) found.push(known);
    else if (word.length >= 5 && Math.abs(known.length - word.length) <= 1 && oneEditApart(word, known)) {
      found.push(known);
    }
  }
  return found;
}

/** True when one insertion, deletion or substitution turns `a` into `b`. */
function oneEditApart(a: string, b: string): boolean {
  if (a === b) return false;
  let i = 0;
  while (i < a.length && i < b.length && a[i] === b[i]) i += 1;
  if (a.length === b.length) return a.slice(i + 1) === b.slice(i + 1);
  return a.length > b.length ? a.slice(i + 1) === b.slice(i) : a.slice(i) === b.slice(i + 1);
}

function main(): void {
  const question = process.argv.slice(2).join(" ");
  if (question.length === 0) {
    console.error('usage: tsx ask/retrieve.ts "<question>"');
    process.exit(1);
  }
  console.log(`words: ${words(question).join(" ")}`);
  for (const law of bm25(LAWS, question)) console.log(`${law.id}  ${law.law}`);
}

if (process.argv[1] === fileURLToPath(import.meta.url)) main();

npx tsx ask/retrieve.ts "how far from a fire hydrant can i park"