Skip to content

Commit

Permalink
squash! llama : std::move llm_bigram_bpe from work_queue
Browse files Browse the repository at this point in the history
Introduced a MovablePriorityQueue class to allow moving elements
out of the priority queue for llm_bigram_bpe.
  • Loading branch information
danbev committed Aug 18, 2024
1 parent 6c6db7b commit 823948c
Showing 1 changed file with 17 additions and 3 deletions.
20 changes: 17 additions & 3 deletions src/llama-vocab.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -321,6 +321,21 @@ struct llm_tokenizer_spm {

// TODO: there are a lot of common parts between spm and bpe tokenizers, should be refactored and reused

template<typename T, typename Container = std::vector<T>, typename Compare = std::less<typename Container::value_type>>
class MovablePriorityQueue : public std::priority_queue<T, Container, Compare> {
public:
using std::priority_queue<T, Container, Compare>::priority_queue;

T pop_move() {
T item = std::move(this->c.front());
std::pop_heap(this->c.begin(), this->c.end(), this->comp);
this->c.pop_back();
return item;
}

void pop() = delete;
};

struct llm_bigram_bpe {
struct comparator {
bool operator()(const llm_bigram_bpe & l, const llm_bigram_bpe & r) const {
Expand All @@ -329,7 +344,7 @@ struct llm_bigram_bpe {
};

using queue_storage = std::vector<llm_bigram_bpe>;
using queue = std::priority_queue<llm_bigram_bpe, queue_storage, comparator>;
using queue = MovablePriorityQueue<llm_bigram_bpe, queue_storage, comparator>;
llm_symbol::index left;
llm_symbol::index right;
std::string text;
Expand Down Expand Up @@ -520,8 +535,7 @@ struct llm_tokenizer_bpe {

// build token(s)
while (!work_queue.empty()) {
auto bigram = std::move(const_cast<llm_bigram_bpe&>(work_queue.top()));
work_queue.pop();
auto bigram = work_queue.pop_move();

auto & left_symbol = symbols[bigram.left];
auto & right_symbol = symbols[bigram.right];
Expand Down

0 comments on commit 823948c

Please sign in to comment.