import { t as p } from "../index-DSGwv2Yx.js";
import { calculateLoss as d, calculateAccuracy as m } from "./loss.js";
import { buildSFTExample as x } from "./SFTDatasetBuilder.js";
import { t as h } from "../tensor-D8e0Gd7c.js";
class k {
  constructor(i, t, o, c) {
    if (this.model = i, this.masked = !!c, Array.isArray(t)) {
      if (!o)
        throw new Error("Tokeniser is required when dataset is an array of conversations");
      const a = t.map((s) => x(s, -100, o, i.config.blockSize)).filter((s) => s !== null);
      if (a.length === 0)
        return;
      this.xs = h(a.map((s) => s.xs)), this.ys = h(a.map((s) => s.ys));
    } else
      this.iterator = t.iterator();
  }
  iterator;
  xs;
  ys;
  masked = !1;
  dispose() {
    this.xs && this.xs.dispose(), this.ys && this.ys.dispose();
  }
  async calculateBatchLoss(i, t, o, c) {
    const [a, s] = p(() => {
      const r = this.model.forward({ training: !1 }, i), y = d(r, t, c, o), f = m(r, t);
      return r.dispose(), [y, f];
    }), n = await a.array(), u = await s.array(), e = n, l = u;
    return s.dispose(), a.dispose(), Array.isArray(e) ? e.map((r) => ({ loss: r, accuracy: l })) : { loss: e, accuracy: l };
  }
  async evaluate(i = 100) {
    let t = 0, o = 0, c = 0;
    if (this.iterator) {
      const a = await this.iterator;
      for (let s = 0; s < i; s++) {
        const n = await a.next();
        if (n.done) break;
        const u = n.value, { xs: e, ys: l } = u, r = await this.calculateBatchLoss(e, l, !1, this.masked);
        e.dispose(), l.dispose(), t += r.loss, o += r.accuracy, c++;
      }
      return { loss: t / c, accuracy: o / c };
    } else if (this.xs && this.ys)
      return this.calculateBatchLoss(this.xs, this.ys, !0, !0);
    throw new Error("No data available for evaluation");
  }
}
export {
  k as default
};
