ember_draft::DraftTable Class Reference

#include <draft_table.h>

Collaboration diagram for ember_draft::DraftTable:

Classes
struct	Order

Public Member Functions
bool	load (const std::string &path, int vocab_size, std::string *why)
int	orders () const
size_t	entries () const
std::vector< int >	draft (const std::vector< int > &h, int k, double theta) const

Static Public Member Functions
static uint64_t	mix64 (uint64_t x)

Private Member Functions
bool	lookup (const std::vector< int > &ctx, int *tok, double *c) const

Static Private Member Functions
static bool	fail (std::string *why, const std::string &w)

Private Attributes
std::vector< char >	buf_
std::vector< Order >	tabs_

Detailed Description

Definition at line 27 of file draft_table.h.

Member Function Documentation

draft()

std::vector< int > ember_draft::DraftTable::draft ( const std::vector< int > &  h, int  k, double  theta  ) const

inline

Definition at line 82 of file draft_table.h.

                                                                           {
    std::vector<int> d;
    if (tabs_.empty() || k <= 0) return d;
    const int no = orders();
    std::vector<int> ctx(h.end() - std::min<size_t>(h.size(), static_cast<size_t>(no)), h.end());
    double p = 1.0;
    while (static_cast<int>(d.size()) < k) {
      int tok = -1;
      double c = 0.0;
      if (!lookup(ctx, &tok, &c)) break;
      p *= c;
      if (p < theta) break;
      d.push_back(tok);
      ctx.push_back(tok);
      if (static_cast<int>(ctx.size()) > no) ctx.erase(ctx.begin());
    }
    return d;
  }

References c, j, k, lookup(), orders(), p, tabs_, and tok.

Referenced by tinyllama::TinyLlamaSession::generate().

entries()

size_t ember_draft::DraftTable::entries ( ) const

inline

Definition at line 75 of file draft_table.h.

                         {
    size_t s = 0;
    for (const auto& t : tabs_) s += t.n;
    return s;
  }

References j, and tabs_.

Referenced by tinyllama::corpus_draft_table().

fail()

static bool ember_draft::DraftTable::fail ( std::string *  why, const std::string &  w  )

inlinestaticprivate

Definition at line 149 of file draft_table.h.

                                                       {
    if (why) *why = w;
    return false;
  }

References j, and why.

Referenced by load().

load()

bool ember_draft::DraftTable::load ( const std::string &  path, int  vocab_size, std::string *  why  )

inline

Definition at line 31 of file draft_table.h.

                                                                   {
    std::ifstream f(path, std::ios::binary);
    if (!f) return fail(why, "cannot open " + path);
    buf_.assign(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>());
    if (buf_.size() < 16 || std::memcmp(buf_.data(), "EMBERDT1", 8) != 0) return fail(why, "not a draft table");
    uint32_t vocab = 0, orders = 0;
    std::memcpy(&vocab, buf_.data() + 8, 4);
    std::memcpy(&orders, buf_.data() + 12, 4);
    if (orders < 1 || orders > 8) return fail(why, "orders must be 1..8");
    if (static_cast<int>(vocab) != vocab_size)
      return fail(why, "built for vocabulary size " + std::to_string(vocab) + ", the model has " +
                           std::to_string(vocab_size));
    size_t p = 16 + 8 * static_cast<size_t>(orders);
    if (buf_.size() < p) return fail(why, "truncated header");
    std::vector<uint64_t> ns(orders);
    std::memcpy(ns.data(), buf_.data() + 16, 8 * orders);
    tabs_.assign(orders, Order{});
    auto pad8 = [](size_t b) { return (b + 7) / 8 * 8; };
    for (uint32_t o = 0; o < orders; ++o) {
      const size_t n = ns[o];
      const size_t need = pad8(8 * n) + pad8(2 * n) + pad8(n);
      if (buf_.size() - p < need) return fail(why, "truncated tables");
      tabs_[o].n = n;
      tabs_[o].keys.resize(n);
      tabs_[o].next.resize(n);
      tabs_[o].conf.resize(n);
      std::memcpy(tabs_[o].keys.data(), buf_.data() + p, 8 * n);
      p += pad8(8 * n);
      std::memcpy(tabs_[o].next.data(), buf_.data() + p, 2 * n);
      p += pad8(2 * n);
      std::memcpy(tabs_[o].conf.data(), buf_.data() + p, n);
      p += pad8(n);
      for (size_t i = 0; i < n; ++i)
        if (tabs_[o].next[i] >= vocab || (i > 0 && tabs_[o].keys[i] <= tabs_[o].keys[i - 1]))
          return fail(why, "entries out of range or not ascending");
    }
    if (p != buf_.size()) return fail(why, "trailing bytes");
    buf_.clear();
    buf_.shrink_to_fit();
    return true;
  }

References buf_, fail(), j, n, orders(), p, tabs_, and why.

Referenced by tinyllama::corpus_draft_table().

lookup()

bool ember_draft::DraftTable::lookup ( const std::vector< int > &  ctx, int *  tok, double *  c  ) const

inlineprivate

Definition at line 119 of file draft_table.h.

                                                                    {
    const int len = static_cast<int>(ctx.size());
    for (int o = std::min(orders(), len); o >= 1; --o) {
      const Order& t = tabs_[o - 1];
      if (t.n == 0) continue;
      uint64_t key = 0;
      bool ok = true;
      for (int j = 0; j < o; ++j) {
        const int v = ctx[len - 1 - j];
        if (v < 0 || v > 0xFFFF) { ok = false; break; }
        if (j < 4) key |= static_cast<uint64_t>(v) << (16 * j);
      }
      if (!ok) continue;
      if (o >= 4) {
        key = mix64(key);
        for (int j = 4; j < o; ++j)
          key = mix64(key ^ ((static_cast<uint64_t>(ctx[len - 1 - j]) + 1) * 0x9E3779B97F4A7C15ULL));
        key >>= 16;
      }
      const auto it = std::lower_bound(t.keys.begin(), t.keys.end(), key);
      if (it != t.keys.end() && *it == key) {
        const size_t i = static_cast<size_t>(it - t.keys.begin());
        *tok = t.next[i];
        *c = static_cast<double>(t.conf[i]) / 255.0;
        return true;
      }
    }
    return false;
  }

References c, j, mix64(), orders(), tabs_, tok, and v.

Referenced by draft().

mix64()

static uint64_t ember_draft::DraftTable::mix64 ( uint64_t  x )

inlinestatic

Definition at line 101 of file draft_table.h.

                                    {
    x ^= x >> 30;
    x *= 0xBF58476D1CE4E5B9ULL;
    x ^= x >> 27;
    x *= 0x94D049BB133111EBULL;
    x ^= x >> 31;
    return x;
  }

References j.

Referenced by lookup().

orders()

int ember_draft::DraftTable::orders ( ) const

inline

Definition at line 73 of file draft_table.h.

{ return static_cast<int>(tabs_.size()); }

References tabs_.

Referenced by tinyllama::corpus_draft_table(), draft(), load(), and lookup().

Member Data Documentation

buf_

std::vector<char> ember_draft::DraftTable::buf_

private

Definition at line 154 of file draft_table.h.

Referenced by load().

tabs_

std::vector<Order> ember_draft::DraftTable::tabs_

private

Definition at line 155 of file draft_table.h.

Referenced by draft(), entries(), load(), lookup(), and orders().

---

The documentation for this class was generated from the following file:

• draft_table.h