import { t as y } from "../index-DSGwv2Yx.js";
import "../dataset-DlqAN81i.js";
import { g as I } from "../readers-17HLdxVM.js";
import "../index-Cp39cXWe.js";
function w(p, o, t, l) {
  const s = [t.bosToken], a = [!1], u = {
    user: t.getSpecialTokenIndex("<|user_start|>"),
    assistant: t.getSpecialTokenIndex("<|assistant_start|>"),
    system: t.getSpecialTokenIndex("<|system_start|>"),
    text: void 0
  }, c = {
    user: t.getSpecialTokenIndex("<|user_end|>"),
    assistant: t.getSpecialTokenIndex("<|assistant_end|>"),
    system: t.getSpecialTokenIndex("<|system_end|>"),
    text: void 0
  };
  for (const e of p) {
    const r = u[e.role], h = c[e.role];
    if (!r || !h)
      throw new Error(`Missing special tokens for role: ${e.role}`);
    s.push(r), a.push(!1);
    const m = e.role === "assistant", x = t.encode(e.content);
    for (const T of x) {
      s.push(T);
      const S = t.isSpecialToken(T);
      a.push(m && !S);
    }
    s.push(h), a.push(m);
  }
  s.push(t.eosToken), a.push(!1);
  const n = l + 1;
  if (s.length < n) {
    const e = n - s.length, r = t.getSpecialTokenIndex("<pad>");
    for (let h = 0; h < e; h++)
      s.push(r), a.push(!1);
  } else s.length > n && (s.length = n, a.length = n);
  const f = new Int32Array(s.slice(0, l)), i = s.slice(1, l + 1), k = a.slice(1, l + 1), d = new Int32Array(i.length);
  let g = !1;
  for (let e = 0; e < i.length; e++) {
    const r = k[e] ? i[e] : o;
    d[e] = r, r !== o && (g = !0);
  }
  return g ? { xs: f, ys: d } : null;
}
class A {
  tokenizer;
  blockSize;
  constructor(o, t = 128) {
    this.tokenizer = o, this.blockSize = t;
  }
  /**
   * Create SFT dataset from structured conversations.
   * - Always starts from the beginning of each conversation.
   * - Pads with eosToken and masks padding.
   * - Masks non-assistant tokens in labels with ignoreIndex (default -100).
   */
  async createSFTDataset(o, t = 32, l = -100) {
    if (!o.length)
      throw new Error("No conversations provided.");
    const s = this.tokenizer, a = this.blockSize;
    for (const c of o)
      c.shuffle();
    return I(function* () {
      for (; ; ) {
        const c = Math.floor(Math.random() * o.length), n = o[c], f = n.nextConversation();
        if (!f) {
          n.shuffle();
          continue;
        }
        const i = w(f, l, s, a);
        i && (yield i);
      }
    }).batch(t).map((c) => {
      const n = c;
      return y(() => ({
        xs: n.xs.cast("int32"),
        ys: n.ys.cast("int32")
      }));
    }).prefetch(2);
  }
}
export {
  A as SFTDatasetBuilder,
  w as buildSFTExample
};
