nochinator/thought-vectors-chat

1// The chat loop, ported line-for-line from OnnxChat in2// scripts/quantize_web.py (the parity-proven python reference). No imports:3// the ONNX runtime, sessions, and tokenizer are injected so the same module4// runs in the browser and under Node for the parity test.5 6export const PAD_ID = 0;7export const BOS_ID = 1;8export const EOS_ID = 2;9 10const D = 384;11const K = 8; // thoughts per turn12 13export class OnnxChat {14  /**15   * @param ort onnxruntime-web module (for Tensor construction)16   * @param sessions {encoder, thinker, decoder} InferenceSessions17   * @param tok {encode(text)->number[], decode(ids)->string}18   */19  constructor(ort, sessions, tok) {20    this.ort = ort;21    this.s = sessions;22    this.tok = tok;23    this.history = [];24  }25 26  reset() {27    this.history = [];28  }29 30  ids64(ids) {31    return new this.ort.Tensor(32      "int64", BigInt64Array.from(ids, BigInt), [1, ids.length]);33  }34 35  /** Greedy reply; onToken(text-so-far) fires as tokens decode. */36  async reply(text, onToken) {37    this.history.push(text.trim());38    const turns = this.history.slice(-6);39    const n = turns.length;40    const firstRole = (this.history.length - n) % 2;41    const th = new Float32Array(n * K * D);42    const roles = [], dist = [];43    for (let j = 0; j < n; j++) {44      const ids = [BOS_ID, ...this.tok.encode(turns[j]).slice(0, 254), EOS_ID];45      const enc = await this.s.encoder.run({ ids: this.ids64(ids) });46      th.set(enc.thoughts.data, j * K * D);47      roles.push((firstRole + j) % 2);48      dist.push(Math.min(n - j, 6));49    }50    const out = await this.s.thinker.run({51      ctx_th: new this.ort.Tensor("float32", th, [1, n, K, D]),52      ctx_roles: this.ids64(roles),53      dist: this.ids64(dist),54    });55    const score = out.score.data;56    let best = 0;57    for (let h = 1; h < score.length; h++) if (score[h] < score[best]) best = h;58    const thoughts = new this.ort.Tensor(59      "float32", out.hyps.data.slice(best * K * D, (best + 1) * K * D),60      [1, K, D]);61 62    const ids = [BOS_ID];63    for (let step = 0; step < 255; step++) {64      const fed = ids.length < 2 ? [...ids, 0] : ids;65      const dec = await this.s.decoder.run({66        thoughts,67        ids: this.ids64(fed),68        pos: new this.ort.Tensor(69          "int64", BigInt64Array.from([BigInt(ids.length - 1)]), [1]),70      });71      const lg = dec.logits.data;72      if (ids.length >= 3) { // no_repeat_ngram=373        const p0 = ids[ids.length - 2], p1 = ids[ids.length - 1];74        for (let k = 0; k < ids.length - 2; k++) {75          if (ids[k] === p0 && ids[k + 1] === p1) lg[ids[k + 2]] = -Infinity;76        }77      }78      let nxt = 0;79      for (let v = 1; v < lg.length; v++) if (lg[v] > lg[nxt]) nxt = v;80      ids.push(nxt);81      if (onToken) onToken(this.tok.decode(ids));82      if (nxt === EOS_ID) break;83    }84    const reply = this.tok.decode(ids);85    this.history.push(reply);86    return reply;87  }88}89 90// The paper's §6.5 matched token-LM baseline. Ported line-for-line from91// OnnxLmChat in scripts/quantize_web.py. Deliberately un-improved: flat92// token history (no turn windowing) and no repeat-ngram ban, because its93// repetition loops and apology-default register are the paper's documented94// finding about this paradigm, not bugs for this loop to paper over.95export class OnnxLmChat {96  /**97   * @param ort onnxruntime-web module98   * @param session {lm} InferenceSession99   * @param tok {encode(text)->number[], decode(ids)->string}100   * @param maxLen model's max_seq_len (384 for the released baseline)101   * @param maxNew max new tokens per reply (64 for the released baseline)102   */103  constructor(ort, session, tok, maxLen, maxNew = 64) {104    this.ort = ort;105    this.s = session;106    this.tok = tok;107    this.maxLen = maxLen;108    this.maxNew = maxNew;109    this.history = [];110  }111 112  reset() {113    this.history = [];114  }115 116  ids64(ids) {117    return new this.ort.Tensor(118      "int64", BigInt64Array.from(ids, BigInt), [1, ids.length]);119  }120 121  async reply(text, onToken) {122    this.history.push(text.trim());123    let ids = [];124    for (const t of this.history) {125      ids = ids.concat([BOS_ID], this.tok.encode(t), [EOS_ID]);126    }127    const room = this.maxLen - this.maxNew - 1;128    const out = ids.slice(-room);129    out.push(BOS_ID);130    const gen = [];131    for (let step = 0; step < this.maxNew; step++) {132      const dec = await this.s.lm.run({133        ids: this.ids64(out),134        pos: new this.ort.Tensor(135          "int64", BigInt64Array.from([BigInt(out.length - 1)]), [1]),136      });137      const lg = dec.logits.data;138      let nxt = 0;139      for (let v = 1; v < lg.length; v++) if (lg[v] > lg[nxt]) nxt = v;140      if (nxt === EOS_ID) break;141      gen.push(nxt);142      out.push(nxt);143      if (onToken) onToken(this.tok.decode(gen));144    }145    const reply = this.tok.decode(gen);146    this.history.push(reply);147    return reply;148  }149}150