import { yieldIfNeeded as u } from "../utilities/yielder.js";
import b, { SPECIALS as T } from "./BaseTokeniser.js";
const l = ["<eos>", "<unk>"];
class x extends b {
  vocabSize = 0;
  eosToken = 0;
  bosToken = 0;
  unkToken = 0;
  vocab = [];
  cache = /* @__PURE__ */ new Map();
  _trained = !1;
  constructor(t) {
    if (super(), Array.isArray(t)) {
      if (this.vocab = t, this.vocab.length > 0)
        this.vocabSize = this.vocab.length, T.forEach((i) => {
          const e = this.vocab.indexOf(i);
          e !== -1 && this.addSpecialToken(i, e);
        }), this.eosToken = this.getSpecialTokenIndex("<eos>"), this.bosToken = this.getSpecialTokenIndex("<bos>") ?? this.eosToken, this.unkToken = this.getSpecialTokenIndex("") ?? -1, this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("<unk>")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("<pad>")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf("_")), this.unkToken === -1 && (this.unkToken = this.vocab.indexOf(" ")), this.unkToken === -1 && (this.unkToken = this.eosToken), this.vocab = this.vocab.map((i) => i === "<pad>" ? "" : i), this.vocab.forEach((i, e) => {
          this.cache.set(i, e);
        });
      else
        throw new Error("Vocab cannot be empty");
      this._trained = !0;
    } else
      this.vocabSize = t, this.vocab = new Array(this.vocabSize).fill(""), this.addSpecialTokens(), this.eosToken = this.getSpecialTokenIndex("<eos>"), this.bosToken = this.getSpecialTokenIndex("<bos>") ?? this.eosToken, this.unkToken = this.getSpecialTokenIndex(""), this.vocab.forEach((i, e) => {
        this.cache.set(i, e);
      }), this.cache.set("", this.unkToken);
  }
  addToken(t, i) {
    if (this.cache.has(t))
      return this.cache.get(t);
    let e;
    if (i !== void 0 ? e = i : (e = this.vocab.indexOf("", this.unkToken + 1), e === -1 && (e = this.vocabSize)), e >= this.vocabSize)
      throw new Error("Vocab size exceeded");
    return this.vocab[e] = t, this.cache.set(t, e), e;
  }
  get trained() {
    return this.vocab.length === this.vocabSize && this._trained;
  }
  destroy() {
    this.cache.clear(), this.vocab = [];
  }
  async train(t, i, e) {
    this.datasetID = e;
    const a = /* @__PURE__ */ new Set();
    let h = performance.now();
    for (const n of t)
      n.forEach((o) => {
        for (const r of o.content)
          a.add(r);
      }), h = await u(h, i, 0);
    const s = Array.from(a), k = this.vocab.indexOf("", this.unkToken + 1), d = this.vocabSize - l.length;
    if (k === -1)
      return this.generateID(), this.vocabSize;
    if (this._trained = !0, s.length > d) {
      const n = /* @__PURE__ */ new Map();
      t.forEach((o) => {
        o.forEach((r) => {
          for (const f of r.content)
            n.set(f, (n.get(f) || 0) + 1);
        });
      }), s.sort((o, r) => (n.get(o) || 0) - (n.get(r) || 0)), s.splice(0, s.length - d);
    }
    let c = k;
    if (c !== -1) {
      const n = new Set(this.vocab);
      for (const o of s)
        if (!n.has(o) && (this.vocab[c] = o, n.add(o), c = this.vocab.indexOf("", c + 1), c === -1))
          break;
    }
    return this.cache.clear(), this.vocab.forEach((n, o) => {
      this.cache.set(n, o);
    }), this.generateID(), this.emit("trainStatus", "trained"), this.vocabSize;
  }
  tokenise(t, i) {
    if (!this.trained)
      throw new Error("Tokeniser not trained");
    return t.map((a) => i ? a.split("").map((h) => this.cache.get(h) ?? this.unkToken) : a.split("").map((h) => {
      const s = this.cache.get(h);
      return s !== void 0 ? this.vocab[s] : "";
    }));
  }
  detokenise(t) {
    return t.map((e) => Array.from(e).map((a) => this.vocab[a] || "").join(""));
  }
  encode(t) {
    return this.tokenise([t], !0)[0];
  }
  decode(t) {
    return this.detokenise([t])[0];
  }
  getVocab() {
    return this.vocab;
  }
  getMerges() {
    return [];
  }
  async createTrainingData(t, i = 5) {
    const e = await this.tokenise(t, !0), a = [], h = [];
    for (let s = 0; s < e.length - i; s++)
      a.push(...e[s].slice(0, i)), h.push(e[s + 1][0]);
    return [a, h];
  }
}
export {
  x as default
};
