scoring.ts

import type { EvalAnswer, EvalQuestion } from "./eval-question.ts";

export interface ModelAnswer {
  answer: EvalAnswer | null;
  /** `p=0.73` for a situation, `switching_tool 0.81` for an intent, or the error code. */
  shown: string;
}

export interface CaseOutcome {
  caseId: string;
  pair: string | null;
  expected: EvalAnswer;
  ruleAnswer: EvalAnswer;
  /** Situations only: "yes" whenever any built-in signal fired. */
  signalPresentAnswer: EvalAnswer | null;
  modelAnswers: ModelAnswer[];
  modelMajority: EvalAnswer | null;
  /** Repeats that disagreed with the majority answer. */
  flips: number;
}

export interface QuestionSummary {
  questionId: string;
  isControl: boolean;
  caseCount: number;
  modelCorrect: number;
  ruleCorrect: number;
  modelOnlyCorrect: number;
  ruleOnlyCorrect: number;
  signTestPValue: number;
  pairCount: number;
  modelPairsCorrect: number;
  rulePairsCorrect: number;
  signalPresentPairsCorrect: number | null;
  /** Situations only: cases where the model's answer equals "any signal fired". */
  mirrorAgreements: number | null;
  casesWithFlips: number;
  modelErrors: number;
}

export function findMajorityAnswer(answers: ReadonlyArray<EvalAnswer | null>): {
  majority: EvalAnswer | null;
  flips: number;
} {
  const tallies = new Map<EvalAnswer | null, number>();
  for (const answer of answers) tallies.set(answer, (tallies.get(answer) ?? 0) + 1);
  let majority: EvalAnswer | null = null;
  let majorityCount = 0;
  for (const [answer, count] of tallies) {
    if (count > majorityCount) {
      majority = answer;
      majorityCount = count;
    }
  }
  return { majority, flips: answers.length - majorityCount };
}

export function summarizeQuestion(
  question: EvalQuestion,
  outcomes: readonly CaseOutcome[],
): QuestionSummary {
  const isModelCorrect = (outcome: CaseOutcome) => outcome.modelMajority === outcome.expected;
  const isRuleCorrect = (outcome: CaseOutcome) => outcome.ruleAnswer === outcome.expected;
  const modelOnlyCorrect = outcomes.filter(
    (outcome) => isModelCorrect(outcome) && !isRuleCorrect(outcome),
  ).length;
  const ruleOnlyCorrect = outcomes.filter(
    (outcome) => isRuleCorrect(outcome) && !isModelCorrect(outcome),
  ).length;
  const isSituation = question.kind === "situation";

  return {
    questionId: question.id,
    isControl: question.isControl,
    caseCount: outcomes.length,
    modelCorrect: outcomes.filter(isModelCorrect).length,
    ruleCorrect: outcomes.filter(isRuleCorrect).length,
    modelOnlyCorrect,
    ruleOnlyCorrect,
    signTestPValue: calculateTwoSidedSignTestPValue(modelOnlyCorrect, ruleOnlyCorrect),
    pairCount: groupByPair(outcomes).length,
    modelPairsCorrect: countPairsAnsweredCorrectly(outcomes, (outcome) => outcome.modelMajority),
    rulePairsCorrect: countPairsAnsweredCorrectly(outcomes, (outcome) => outcome.ruleAnswer),
    signalPresentPairsCorrect: isSituation
      ? countPairsAnsweredCorrectly(outcomes, (outcome) => outcome.signalPresentAnswer)
      : null,
    mirrorAgreements: isSituation
      ? outcomes.filter((outcome) => outcome.modelMajority === outcome.signalPresentAnswer).length
      : null,
    casesWithFlips: outcomes.filter((outcome) => outcome.flips > 0).length,
    modelErrors: outcomes.reduce(
      (total, outcome) =>
        total + outcome.modelAnswers.filter((modelAnswer) => modelAnswer.answer === null).length,
      0,
    ),
  };
}

/**
 * Exact two-sided sign test on the discordant cases (McNemar's exact test):
 * if the model and the rule were equally good, which side wins each
 * disagreement would be a coin flip.
 */
export function calculateTwoSidedSignTestPValue(
  modelOnlyCorrect: number,
  ruleOnlyCorrect: number,
): number {
  const discordant = modelOnlyCorrect + ruleOnlyCorrect;
  if (discordant === 0) return 1;
  const smallerSide = Math.min(modelOnlyCorrect, ruleOnlyCorrect);
  let tailProbability = 0;
  for (let wins = 0; wins <= smallerSide; wins++)
    tailProbability += calculateBinomialCoefficient(discordant, wins);
  return Math.min(1, (2 * tailProbability) / 2 ** discordant);
}

function calculateBinomialCoefficient(total: number, chosen: number): number {
  let coefficient = 1;
  for (let step = 1; step <= chosen; step++)
    coefficient = (coefficient * (total - chosen + step)) / step;
  return coefficient;
}

function groupByPair(outcomes: readonly CaseOutcome[]): CaseOutcome[][] {
  const pairs = new Map<string, CaseOutcome[]>();
  for (const outcome of outcomes) {
    if (outcome.pair === null) continue;
    pairs.set(outcome.pair, [...(pairs.get(outcome.pair) ?? []), outcome]);
  }
  return [...pairs.values()];
}

function countPairsAnsweredCorrectly(
  outcomes: readonly CaseOutcome[],
  pickAnswer: (outcome: CaseOutcome) => EvalAnswer | null,
): number {
  return groupByPair(outcomes).filter((members) =>
    members.every((outcome) => pickAnswer(outcome) === outcome.expected),
  ).length;
}