async run()

in packages/core/src/models/bert.js [14:65]
48 lines of code
4 McCabe index (conditional complexity)

  async run() {
    const {
      result: [test, expected, cleanup],
      time: setupTime,
    } = await time(async () => {
      const model_id = "Snowflake/snowflake-arctic-embed-xs";
      const tokenizer = await AutoTokenizer.from_pretrained(model_id);
      const model = await BertModel.from_pretrained(model_id, {
        ...this.options,
      });
      const inputs = await tokenizer(["hello", "hello world"], {
        truncation: true,
        padding: true,
      });

      return [
        async () => {
          const { last_hidden_state } = await model(inputs);
          return pick(last_hidden_state, ["type", "dims"]);
        },
        {
          type: "float32",
          dims: [2, 4, 384],
        },
        () => model.dispose(),
      ];
    });

    const times = [];
    const numRuns = DEFAULT_NUM_WARMUP_RUNS + this.num_runs;
    for (let i = 0; i < numRuns; ++i) {
      const { result, time: executionTime } = await time(test);
      const { pass, message } = toBeCloseToNested(result, expected);
      if (!pass) {
        console.log(result);
        console.log(expected);
        throw new Error(message());
      }
      if (i >= DEFAULT_NUM_WARMUP_RUNS) times.push(executionTime);
    }
    const stats = {
      [this.name]: computeStatistics(times),
    };

    const { time: disposeTime } = await time(cleanup);

    return {
      setupTime,
      stats,
      disposeTime,
    };
  }