Benchmarking inference providers

import dotenv from "dotenv";
import * as fs from "fs";

if (fs.existsSync(".env.local")) {
  dotenv.config({ path: ".env.local", override: true });
}

import { OpenAI } from "openai";
import { wrapOpenAI } from "braintrust";

import { templates } from "autoevals";
import * as yaml from "js-yaml";
import mustache from "mustache";

const client = wrapOpenAI(
  new OpenAI({
    apiKey: process.env.BRAINTRUST_API_KEY,
    baseURL: "https://api.braintrust.dev/v1/proxy",
    defaultHeaders: { "x-bt-use-cache": "never" },
  })
);

function parseToolResponse(response: string) {
  const functionRegex = /<function=(\w+)>(.*?)(?:<\/function>|$)/;
  const match = response.match(functionRegex);

  if (match) {
    const [, functionName, argsString] = match;
    try {
      const args = JSON.parse(argsString);
      return {
        functionName,
        args,
      };
    } catch (error) {
      console.error("Error parsing function arguments:", error);
      return null;
    }
  }

  return null;
}

const template = yaml.load(templates["factuality"]);

const selectTool = {
  name: "select_choice",
  description: "Call this function to select a choice.",
  parameters: {
    properties: {
      reasons: {
        description:
          "Write out in a step by step manner your reasoning to be sure that your conclusion is correct. Avoid simply stating the correct answer at the outset.",
        title: "Reasoning",
        type: "string",
      },
      choice: {
        description: "The choice",
        title: "Choice",
        type: "string",
        enum: Object.keys(template.choice_scores),
      },
    },
    required: ["reasons", "choice"],
    title: "CoTResponse",
    type: "object",
  },
};

async function LLaMaFactuality({
  model,
  input,
  output,
  expected,
}: {
  model: string;
  input: string;
  output: string;
  expected: string;
}) {
  const toolPrompt = `You have access to the following functions:

Use the function '${selectTool.name}' to '${selectTool.description}':
${JSON.stringify(selectTool)}

If you choose to call a function ONLY reply in the following format with no prefix or suffix:

<function=example_function_name>{"example_name": "example_value"}</function>

Reminder:
- If looking for real time information use relevant functions before falling back to brave_search
- Function calls MUST follow the specified format, start with <function= and end with </function>
- Required parameters MUST be specified
- Only call one function at a time
- Put the entire function call reply on one line

Here are a few examples:

`;

  const response = await client.chat.completions.create({
    model,
    messages: [
      {
        role: "system",
        content: toolPrompt,
      },
      {
        role: "user",
        content: mustache.render(template.prompt, {
          input,
          output,
          expected,
        }),
      },
    ],
    temperature: 0,
    max_tokens: 2048,
  });

  try {
    const parsed = parseToolResponse(response.choices[0].message.content);
    return {
      name: "Factuality",
      score: template.choice_scores[parsed?.args.choice],
      metadata: {
        rationale: parsed?.args.reasons,
        choice: parsed?.args.choice,
      },
    };
  } catch (e) {
    return {
      name: "Factuality",
      score: -1,
      metadata: {
        error: `${e}`,
      },
    };
  }
}

console.log(
  await LLaMaFactuality({
    model: "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
    input: "What is the weather in Tokyo?",
    output: "The weather in Tokyo is scorching.",
    expected: "The weather in Tokyo is extremely hot.",
  })
);

(node:12633) [DEP0040] DeprecationWarning: The \`punycode\` module is deprecated. Please use a userland alternative instead.
(Use \`node --trace-deprecation ...\` to show where the warning was created)

{
  name: 'Factuality',
  score: 0.6,
  metadata: {
    rationale: "The submitted answer 'The weather in Tokyo is scorching' is a superset of the expert answer 'The weather in Tokyo is extremely hot' because it includes the same information and adds more detail. The word 'scorching' is a synonym for 'extremely hot', so the submitted answer is fully consistent with the expert answer.",
    choice: 'B'
  }
}

interface CoqaCase {
  input: {
    input: string;
    output: string;
    expected: string;
  };
  expected: number;
}

const data: CoqaCase[] = JSON.parse(
  fs.readFileSync("../LLaMa-3_1-Tools/coqa-factuality.json", "utf-8")
);

console.log("LLaMa-3.1-8B Factuality");
console.log(
  await LLaMaFactuality({
    model: "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
    ...data[1].input,
  })
);

LLaMa-3.1-8B Factuality
{
  name: 'Factuality',
  score: 0,
  metadata: {
    rationale: "The submitted answer 'in a barn' does not contain the word 'white' which is present in the expert answer. Therefore, it is not a subset or superset of the expert answer. The submitted answer also does not contain all the same details as the expert answer. There is a disagreement between the submitted answer and the expert answer.",
    choice: 'D'
  }
}

const providers = [
  {
    provider: "Provider 1",
    models: [
      "accounts/fireworks/models/llama-v3p1-8b-instruct",
      "accounts/fireworks/models/llama-v3p1-70b-instruct",
      "accounts/fireworks/models/llama-v3p1-405b-instruct",
    ],
  },
  {
    provider: "Provider 2",
    models: [
      "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
      "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
      "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
    ],
  },
  {
    provider: "Provider 3",
    models: ["llama3-1-8b", "llama3-1-70b", "llama3-1-405b"],
  },
  {
    provider: "Self-hosted vLLM",
    models: ["meta-llama/Meta-Llama-3.1-405B-Instruct-FP8"],
  },
];

import { Eval } from "braintrust";
import { Score, NumericDiff } from "autoevals";

function NonNull({ output }: { output: number | null }) {
  return output !== null && output !== undefined ? 1 : 0;
}

async function CorrectScore({
  output,
  expected,
}: {
  output: number | null;
  expected: number | null;
}): Promise<Score> {
  if (output === null || expected === null) {
    return {
      name: "CorrectScore",
      score: 0,
      metadata: {
        error: output === null ? "output is null" : "expected is null",
      },
    };
  }
  return {
    ...(await NumericDiff({ output, expected })),
    name: "CorrectScore",
  };
}

async function runProviderBenchmark(provider: (typeof providers)[number]) {
  const evals = [];
  for (const model of provider.models) {
    const size = model.toLowerCase().includes("8b")
      ? "8b"
      : model.toLowerCase().includes("70b")
        ? "70b"
        : "405b";

    evals.push(
      Eval("LLaMa-3.1-Multi-Provider-Benchmark", {
        data: data,
        task: async (input) =>
          (await LLaMaFactuality({ model, ...input }))?.score,
        scores: [CorrectScore, NonNull],
        metadata: {
          size,
          provider: provider.provider,
          model,
        },
        experimentName: `${provider.provider} (${size})`,
        maxConcurrency: 3,
        trialCount: 3,
      })
    );
  }
  await Promise.all(evals);
}

await Promise.all(providers.map(runProviderBenchmark));