src/token-stream.js

// Incremental ByteLevel decoding keeps long generations linear in output length.
export class TokenStream {
  constructor(tokenizer) {
    this.tokenizer = tokenizer;
    this.decoder = new TextDecoder("utf-8", { ignoreBOM: true });
    this.added = new Map(
      tokenizer.tokenizer.added_tokens.map((t) => [t.id, t.content]),
    );
    this.bytes = tokenizer.decoder.byte_decoder;
    this.cache = new Map();
  }
  push(id) {
    if (this.added.has(id)) return this.decoder.decode() + this.added.get(id);
    let bytes = this.cache.get(id);
    if (!bytes) {
      bytes = Uint8Array.from(
        [...this.tokenizer.model.vocab[id]].map((c) => this.bytes[c]),
      );
      this.cache.set(id, bytes);
    }
    return this.decoder.decode(bytes, { stream: true });
  }
  flush() {
    return this.decoder.decode();
  }
}